go.temporal.io/server@v1.23.0/common/collection/concurrent_tx_map_test.go (about)

     1  // The MIT License
     2  //
     3  // Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
     4  //
     5  // Copyright (c) 2020 Uber Technologies, Inc.
     6  //
     7  // Permission is hereby granted, free of charge, to any person obtaining a copy
     8  // of this software and associated documentation files (the "Software"), to deal
     9  // in the Software without restriction, including without limitation the rights
    10  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    11  // copies of the Software, and to permit persons to whom the Software is
    12  // furnished to do so, subject to the following conditions:
    13  //
    14  // The above copyright notice and this permission notice shall be included in
    15  // all copies or substantial portions of the Software.
    16  //
    17  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    18  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    19  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    20  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    21  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    22  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    23  // THE SOFTWARE.
    24  
    25  package collection
    26  
    27  import (
    28  	"errors"
    29  	"math/rand"
    30  	"sync"
    31  	"sync/atomic"
    32  	"testing"
    33  
    34  	"github.com/pborman/uuid"
    35  	"github.com/stretchr/testify/require"
    36  	"github.com/stretchr/testify/suite"
    37  )
    38  
    39  type (
    40  	ConcurrentTxMapSuite struct {
    41  		*require.Assertions // override suite.Suite.Assertions with require.Assertions; this means that s.NotNil(nil) will stop the test, not merely log an error
    42  		suite.Suite
    43  	}
    44  	boolType bool
    45  	intType  int
    46  )
    47  
    48  func TestConcurrentTxMapSuite(t *testing.T) {
    49  	suite.Run(t, new(ConcurrentTxMapSuite))
    50  }
    51  
    52  func (s *ConcurrentTxMapSuite) SetupTest() {
    53  	s.Assertions = require.New(s.T()) // Have to define our overridden assertions in the test setup. If we did it earlier, s.T() will return nil
    54  }
    55  
    56  func (s *ConcurrentTxMapSuite) TestLen() {
    57  	testMap := NewShardedConcurrentTxMap(1, UUIDHashCode)
    58  
    59  	key1 := "0001"
    60  	testMap.Put(key1, boolType(true))
    61  	s.Equal(1, testMap.Len(), "Wrong concurrent map size")
    62  
    63  	testMap.Put(key1, boolType(false))
    64  	s.Equal(1, testMap.Len(), "Wrong concurrent map size")
    65  
    66  	key2 := "0002"
    67  	testMap.Put(key2, boolType(false))
    68  	s.Equal(2, testMap.Len(), "Wrong concurrent map size")
    69  
    70  	testMap.PutIfNotExist(key2, boolType(false))
    71  	s.Equal(2, testMap.Len(), "Wrong concurrent map size")
    72  
    73  	testMap.Remove(key2)
    74  	s.Equal(1, testMap.Len(), "Wrong concurrent map size")
    75  
    76  	testMap.Remove(key2)
    77  	s.Equal(1, testMap.Len(), "Wrong concurrent map size")
    78  }
    79  
    80  func (s *ConcurrentTxMapSuite) TestGetAndDo() {
    81  	testMap := NewShardedConcurrentTxMap(1, UUIDHashCode)
    82  	key := uuid.New()
    83  	var value intType
    84  	fnApplied := false
    85  
    86  	interf, ok, err := testMap.GetAndDo(key, func(key interface{}, value interface{}) error {
    87  		fnApplied = true
    88  		return nil
    89  	})
    90  	s.Nil(interf, "GetAndDo should return nil when key not found")
    91  	s.Nil(err, "GetAndDo should return nil when function not applied")
    92  	s.False(ok, "GetAndDo should return false when key not found")
    93  	s.False(fnApplied, "GetAndDo should not apply function when key not exixts")
    94  
    95  	value = intType(1)
    96  	testMap.Put(key, &value)
    97  	interf, ok, err = testMap.GetAndDo(key, func(key interface{}, value interface{}) error {
    98  		fnApplied = true
    99  		intValue := value.(*intType)
   100  		*intValue++
   101  		return errors.New("some err")
   102  	})
   103  
   104  	value1 := interf.(*intType)
   105  	s.Equal(*(value1), intType(2))
   106  	s.NotNil(err, "GetAndDo should return non nil when function applied")
   107  	s.True(ok, "GetAndDo should return true when key found")
   108  	s.True(fnApplied, "GetAndDo should apply function when key exixts")
   109  }
   110  
   111  func (s *ConcurrentTxMapSuite) TestPutOrDo() {
   112  	testMap := NewShardedConcurrentTxMap(1, UUIDHashCode)
   113  	key := uuid.New()
   114  	var value intType
   115  	fnApplied := false
   116  
   117  	value = intType(1)
   118  	interf, ok, err := testMap.PutOrDo(key, &value, func(key interface{}, value interface{}) error {
   119  		fnApplied = true
   120  		return errors.New("some err")
   121  	})
   122  	valueRetuern := interf.(*intType)
   123  	s.Equal(value, *valueRetuern)
   124  	s.Nil(err, "PutOrDo should return nil when function not applied")
   125  	s.False(ok, "PutOrDo should return false when function not applied")
   126  	s.False(fnApplied, "PutOrDo should not apply function when key not exixts")
   127  
   128  	anotherValue := intType(111)
   129  	interf, ok, err = testMap.PutOrDo(key, &anotherValue, func(key interface{}, value interface{}) error {
   130  		fnApplied = true
   131  		intValue := value.(*intType)
   132  		*intValue++
   133  		return errors.New("some err")
   134  	})
   135  	valueRetuern = interf.(*intType)
   136  	s.Equal(value, *valueRetuern)
   137  	s.NotNil(err, "PutOrDo should return non nil when function applied")
   138  	s.True(ok, "PutOrDo should return true when function applied")
   139  	s.True(fnApplied, "PutOrDo should apply function when key exixts")
   140  }
   141  
   142  func (s *ConcurrentTxMapSuite) TestRemoveIf() {
   143  	testMap := NewShardedConcurrentTxMap(1, UUIDHashCode)
   144  	key := uuid.New()
   145  	value := intType(1)
   146  	testMap.Put(key, &value)
   147  
   148  	removed := testMap.RemoveIf(key, func(key interface{}, value interface{}) bool {
   149  		intValue := value.(*intType)
   150  		return *intValue == intType(2)
   151  	})
   152  	s.Equal(1, testMap.Len(), "TestRemoveIf should only entry if condition is met")
   153  	s.False(removed, "TestRemoveIf should return false if key is not deleted")
   154  
   155  	removed = testMap.RemoveIf(key, func(key interface{}, value interface{}) bool {
   156  		intValue := value.(*intType)
   157  		return *intValue == intType(1)
   158  	})
   159  	s.Equal(0, testMap.Len(), "TestRemoveIf should only entry if condition is met")
   160  	s.True(removed, "TestRemoveIf should return true if key is deleted")
   161  }
   162  
   163  func (s *ConcurrentTxMapSuite) TestGetAfterPut() {
   164  
   165  	countMap := make(map[string]int)
   166  	testMap := NewShardedConcurrentTxMap(1, UUIDHashCode)
   167  
   168  	for i := 0; i < 1024; i++ {
   169  		key := uuid.New()
   170  		countMap[key] = 0
   171  		testMap.Put(key, boolType(true))
   172  	}
   173  
   174  	for k := range countMap {
   175  		v, ok := testMap.Get(k)
   176  		boolValue := v.(boolType)
   177  		s.True(ok, "Get after put failed")
   178  		s.True(bool(boolValue), "Wrong value returned from map")
   179  	}
   180  
   181  	s.Equal(len(countMap), testMap.Len(), "Size() returned wrong value")
   182  
   183  	it := testMap.Iter()
   184  	for entry := range it.Entries() {
   185  		countMap[entry.Key.(string)]++
   186  	}
   187  	it.Close()
   188  
   189  	for _, v := range countMap {
   190  		s.Equal(1, v, "Iterator test failed")
   191  	}
   192  
   193  	for k := range countMap {
   194  		testMap.Remove(k)
   195  	}
   196  
   197  	s.Equal(0, testMap.Len(), "Map returned non-zero size after deleting all entries")
   198  }
   199  
   200  func (s *ConcurrentTxMapSuite) TestPutIfNotExist() {
   201  	testMap := NewShardedConcurrentTxMap(1, UUIDHashCode)
   202  	key := uuid.New()
   203  	ok := testMap.PutIfNotExist(key, boolType(true))
   204  	s.True(ok, "PutIfNotExist failed to insert item")
   205  	ok = testMap.PutIfNotExist(key, boolType(true))
   206  	s.False(ok, "PutIfNotExist invariant failed")
   207  }
   208  
   209  func (s *ConcurrentTxMapSuite) TestMapConcurrency() {
   210  	nKeys := 1024
   211  	keys := make([]string, nKeys)
   212  	for i := 0; i < nKeys; i++ {
   213  		keys[i] = uuid.New()
   214  	}
   215  
   216  	var total int32
   217  	var startWG sync.WaitGroup
   218  	var doneWG sync.WaitGroup
   219  	testMap := NewShardedConcurrentTxMap(1024, UUIDHashCode)
   220  
   221  	startWG.Add(1)
   222  
   223  	for i := 0; i < 10; i++ {
   224  
   225  		doneWG.Add(1)
   226  
   227  		go func() {
   228  			startWG.Wait()
   229  			for n := 0; n < nKeys; n++ {
   230  				val := intType(rand.Int())
   231  				if testMap.PutIfNotExist(keys[n], val) {
   232  					atomic.AddInt32(&total, int32(val))
   233  					_, ok := testMap.Get(keys[n])
   234  					s.True(ok, "Concurrency Get test failed")
   235  				}
   236  			}
   237  			doneWG.Done()
   238  		}()
   239  	}
   240  
   241  	startWG.Done()
   242  	doneWG.Wait()
   243  
   244  	s.Equal(nKeys, testMap.Len(), "Wrong concurrent map size")
   245  
   246  	var gotTotal int32
   247  	for i := 0; i < nKeys; i++ {
   248  		v, ok := testMap.Get(keys[i])
   249  		s.True(ok, "Get failed to find previously inserted key")
   250  		intVal := v.(intType)
   251  		gotTotal += int32(intVal)
   252  	}
   253  
   254  	s.Equal(total, gotTotal, "Concurrent put test failed, wrong sum of values inserted")
   255  }