github.com/uber/kraken@v0.1.4/utils/lockermap/map_test.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package lockermap
    15  
    16  import (
    17  	"fmt"
    18  	"strconv"
    19  	"sync"
    20  	"testing"
    21  
    22  	"github.com/stretchr/testify/require"
    23  )
    24  
    25  func TestMapTryStoreReturnsFalseOnDuplicates(t *testing.T) {
    26  	require := require.New(t)
    27  	var m Map
    28  
    29  	require.True(m.TryStore("k", new(sync.Mutex)))
    30  	require.False(m.TryStore("k", new(sync.Mutex)))
    31  }
    32  
    33  type testValue struct {
    34  	sync.Mutex
    35  	n int
    36  }
    37  
    38  func TestMapLoadHoldsLock(t *testing.T) {
    39  	require := require.New(t)
    40  	var m Map
    41  
    42  	require.True(m.TryStore("k", new(testValue)))
    43  
    44  	// Only a single goroutine should be able to increment n.
    45  	var wg sync.WaitGroup
    46  	for i := 0; i < 100; i++ {
    47  		wg.Add(1)
    48  		go require.True(m.Load("k", func(l sync.Locker) {
    49  			defer wg.Done()
    50  			v := l.(*testValue)
    51  			if v.n == 0 {
    52  				v.n++
    53  			} else {
    54  				require.Equal(1, v.n)
    55  			}
    56  		}))
    57  	}
    58  	wg.Wait()
    59  }
    60  
    61  type testValueCoordinatedLock struct {
    62  	mu      sync.Mutex
    63  	loading bool
    64  	locking chan bool
    65  	deleted chan bool
    66  }
    67  
    68  func (v *testValueCoordinatedLock) Lock() {
    69  	if v.loading {
    70  		v.locking <- true
    71  		<-v.deleted
    72  	}
    73  	v.mu.Lock()
    74  }
    75  
    76  func (v *testValueCoordinatedLock) Unlock() { v.mu.Unlock() }
    77  
    78  func newTestValueCoordinatedLock() *testValueCoordinatedLock {
    79  	return &testValueCoordinatedLock{
    80  		locking: make(chan bool),
    81  		deleted: make(chan bool),
    82  	}
    83  }
    84  
    85  func TestMapLoadReturnsFalseWhenKeyDeletedBeforeValueLocked(t *testing.T) {
    86  	require := require.New(t)
    87  	var m Map
    88  
    89  	v := newTestValueCoordinatedLock()
    90  
    91  	require.True(m.TryStore("k", v))
    92  
    93  	var wg sync.WaitGroup
    94  	wg.Add(1)
    95  
    96  	// The goroutine should be able to load k, but k is deleted before it can
    97  	// acquire the value lock.
    98  	v.loading = true
    99  	go func() {
   100  		defer wg.Done()
   101  		require.False(m.Load("k", func(l sync.Locker) {}))
   102  	}()
   103  
   104  	<-v.locking
   105  	v.loading = false
   106  	m.Delete("k")
   107  	v.deleted <- true
   108  
   109  	wg.Wait()
   110  }
   111  
   112  func TestMapRange(t *testing.T) {
   113  	require := require.New(t)
   114  	var m Map
   115  
   116  	var wg sync.WaitGroup
   117  
   118  	for i := 0; i < 10; i++ {
   119  		v := newTestValueCoordinatedLock()
   120  		require.True(m.TryStore(strconv.Itoa(i), v))
   121  		wg.Add(1)
   122  	}
   123  
   124  	go func() {
   125  		m.Range(func(k interface{}, v sync.Locker) bool {
   126  			fmt.Println("Iterating - ", k)
   127  			wg.Done()
   128  			return true
   129  		})
   130  	}()
   131  
   132  	wg.Wait()
   133  }
   134  
   135  func TestMapRangeSkipsWhenKeyDeletedBeforeValueLocked(t *testing.T) {
   136  	require := require.New(t)
   137  	var m Map
   138  
   139  	v := newTestValueCoordinatedLock()
   140  
   141  	require.True(m.TryStore("k", v))
   142  
   143  	var wg sync.WaitGroup
   144  	wg.Add(1)
   145  
   146  	// The goroutine should be able to load k, but k is deleted before it can
   147  	// acquire the value lock, and thus range should never execute on anything.
   148  	v.loading = true
   149  	go func() {
   150  		defer wg.Done()
   151  		m.Range(func(k interface{}, v sync.Locker) bool {
   152  			require.Fail("Should not be able to execute Range iteration")
   153  			return true
   154  		})
   155  	}()
   156  
   157  	<-v.locking
   158  	v.loading = false
   159  	m.Delete("k")
   160  	v.deleted <- true
   161  
   162  	wg.Wait()
   163  }