Source file src/internal/sync/hashtriemap_test.go

     1  // Copyright 2024 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package sync_test
     6  
     7  import (
     8  	"fmt"
     9  	isync "internal/sync"
    10  	"math"
    11  	"runtime"
    12  	"strconv"
    13  	"sync"
    14  	"testing"
    15  	"weak"
    16  )
    17  
    18  func TestHashTrieMap(t *testing.T) {
    19  	testHashTrieMap(t, func() *isync.HashTrieMap[string, int] {
    20  		var m isync.HashTrieMap[string, int]
    21  		return &m
    22  	})
    23  }
    24  
    25  func TestHashTrieMapBadHash(t *testing.T) {
    26  	testHashTrieMap(t, func() *isync.HashTrieMap[string, int] {
    27  		return isync.NewBadHashTrieMap[string, int]()
    28  	})
    29  }
    30  
    31  func TestHashTrieMapTruncHash(t *testing.T) {
    32  	testHashTrieMap(t, func() *isync.HashTrieMap[string, int] {
    33  		// Stub out the good hash function with a different terrible one
    34  		// (truncated hash). Everything should still work as expected.
    35  		// This is useful to test independently to catch issues with
    36  		// near collisions, where only the last few bits of the hash differ.
    37  		return isync.NewTruncHashTrieMap[string, int]()
    38  	})
    39  }
    40  
    41  func testHashTrieMap(t *testing.T, newMap func() *isync.HashTrieMap[string, int]) {
    42  	t.Run("LoadEmpty", func(t *testing.T) {
    43  		m := newMap()
    44  
    45  		for _, s := range testData {
    46  			expectMissing(t, s, 0)(m.Load(s))
    47  		}
    48  	})
    49  	t.Run("LoadOrStore", func(t *testing.T) {
    50  		m := newMap()
    51  
    52  		for i, s := range testData {
    53  			expectMissing(t, s, 0)(m.Load(s))
    54  			expectStored(t, s, i)(m.LoadOrStore(s, i))
    55  			expectPresent(t, s, i)(m.Load(s))
    56  			expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
    57  		}
    58  		for i, s := range testData {
    59  			expectPresent(t, s, i)(m.Load(s))
    60  			expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
    61  		}
    62  	})
    63  	t.Run("All", func(t *testing.T) {
    64  		m := newMap()
    65  
    66  		testAll(t, m, testDataMap(testData[:]), func(_ string, _ int) bool {
    67  			return true
    68  		})
    69  	})
    70  	t.Run("Clear", func(t *testing.T) {
    71  		t.Run("Simple", func(t *testing.T) {
    72  			m := newMap()
    73  
    74  			for i, s := range testData {
    75  				expectMissing(t, s, 0)(m.Load(s))
    76  				expectStored(t, s, i)(m.LoadOrStore(s, i))
    77  				expectPresent(t, s, i)(m.Load(s))
    78  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
    79  			}
    80  			m.Clear()
    81  			for _, s := range testData {
    82  				expectMissing(t, s, 0)(m.Load(s))
    83  			}
    84  		})
    85  		t.Run("Concurrent", func(t *testing.T) {
    86  			m := newMap()
    87  
    88  			// Load up the map.
    89  			for i, s := range testData {
    90  				expectMissing(t, s, 0)(m.Load(s))
    91  				expectStored(t, s, i)(m.LoadOrStore(s, i))
    92  			}
    93  			gmp := runtime.GOMAXPROCS(-1)
    94  			var wg sync.WaitGroup
    95  			for i := range gmp {
    96  				wg.Add(1)
    97  				go func(id int) {
    98  					defer wg.Done()
    99  
   100  					for _, s := range testData {
   101  						// Try a couple things to interfere with the clear.
   102  						expectNotDeleted(t, s, math.MaxInt)(m.CompareAndDelete(s, math.MaxInt))
   103  						m.CompareAndSwap(s, i, i+1) // May succeed or fail; we don't care.
   104  					}
   105  				}(i)
   106  			}
   107  
   108  			// Concurrently clear the map.
   109  			runtime.Gosched()
   110  			m.Clear()
   111  
   112  			// Wait for workers to finish.
   113  			wg.Wait()
   114  
   115  			// It should all be empty now.
   116  			for _, s := range testData {
   117  				expectMissing(t, s, 0)(m.Load(s))
   118  			}
   119  		})
   120  	})
   121  	t.Run("CompareAndDelete", func(t *testing.T) {
   122  		t.Run("All", func(t *testing.T) {
   123  			m := newMap()
   124  
   125  			for range 3 {
   126  				for i, s := range testData {
   127  					expectMissing(t, s, 0)(m.Load(s))
   128  					expectStored(t, s, i)(m.LoadOrStore(s, i))
   129  					expectPresent(t, s, i)(m.Load(s))
   130  					expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   131  				}
   132  				for i, s := range testData {
   133  					expectPresent(t, s, i)(m.Load(s))
   134  					expectNotDeleted(t, s, math.MaxInt)(m.CompareAndDelete(s, math.MaxInt))
   135  					expectDeleted(t, s, i)(m.CompareAndDelete(s, i))
   136  					expectNotDeleted(t, s, i)(m.CompareAndDelete(s, i))
   137  					expectMissing(t, s, 0)(m.Load(s))
   138  				}
   139  				for _, s := range testData {
   140  					expectMissing(t, s, 0)(m.Load(s))
   141  				}
   142  			}
   143  		})
   144  		t.Run("One", func(t *testing.T) {
   145  			m := newMap()
   146  
   147  			for i, s := range testData {
   148  				expectMissing(t, s, 0)(m.Load(s))
   149  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   150  				expectPresent(t, s, i)(m.Load(s))
   151  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   152  			}
   153  			expectNotDeleted(t, testData[15], math.MaxInt)(m.CompareAndDelete(testData[15], math.MaxInt))
   154  			expectDeleted(t, testData[15], 15)(m.CompareAndDelete(testData[15], 15))
   155  			expectNotDeleted(t, testData[15], 15)(m.CompareAndDelete(testData[15], 15))
   156  			for i, s := range testData {
   157  				if i == 15 {
   158  					expectMissing(t, s, 0)(m.Load(s))
   159  				} else {
   160  					expectPresent(t, s, i)(m.Load(s))
   161  				}
   162  			}
   163  		})
   164  		t.Run("Multiple", func(t *testing.T) {
   165  			m := newMap()
   166  
   167  			for i, s := range testData {
   168  				expectMissing(t, s, 0)(m.Load(s))
   169  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   170  				expectPresent(t, s, i)(m.Load(s))
   171  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   172  			}
   173  			for _, i := range []int{1, 105, 6, 85} {
   174  				expectNotDeleted(t, testData[i], math.MaxInt)(m.CompareAndDelete(testData[i], math.MaxInt))
   175  				expectDeleted(t, testData[i], i)(m.CompareAndDelete(testData[i], i))
   176  				expectNotDeleted(t, testData[i], i)(m.CompareAndDelete(testData[i], i))
   177  			}
   178  			for i, s := range testData {
   179  				if i == 1 || i == 105 || i == 6 || i == 85 {
   180  					expectMissing(t, s, 0)(m.Load(s))
   181  				} else {
   182  					expectPresent(t, s, i)(m.Load(s))
   183  				}
   184  			}
   185  		})
   186  		t.Run("Iterate", func(t *testing.T) {
   187  			m := newMap()
   188  
   189  			testAll(t, m, testDataMap(testData[:]), func(s string, i int) bool {
   190  				expectDeleted(t, s, i)(m.CompareAndDelete(s, i))
   191  				return true
   192  			})
   193  			for _, s := range testData {
   194  				expectMissing(t, s, 0)(m.Load(s))
   195  			}
   196  		})
   197  		t.Run("ConcurrentUnsharedKeys", func(t *testing.T) {
   198  			m := newMap()
   199  
   200  			gmp := runtime.GOMAXPROCS(-1)
   201  			var wg sync.WaitGroup
   202  			for i := range gmp {
   203  				wg.Add(1)
   204  				go func(id int) {
   205  					defer wg.Done()
   206  
   207  					makeKey := func(s string) string {
   208  						return s + "-" + strconv.Itoa(id)
   209  					}
   210  					for _, s := range testData {
   211  						key := makeKey(s)
   212  						expectMissing(t, key, 0)(m.Load(key))
   213  						expectStored(t, key, id)(m.LoadOrStore(key, id))
   214  						expectPresent(t, key, id)(m.Load(key))
   215  						expectLoaded(t, key, id)(m.LoadOrStore(key, 0))
   216  					}
   217  					for _, s := range testData {
   218  						key := makeKey(s)
   219  						expectPresent(t, key, id)(m.Load(key))
   220  						expectDeleted(t, key, id)(m.CompareAndDelete(key, id))
   221  						expectMissing(t, key, 0)(m.Load(key))
   222  					}
   223  					for _, s := range testData {
   224  						key := makeKey(s)
   225  						expectMissing(t, key, 0)(m.Load(key))
   226  					}
   227  				}(i)
   228  			}
   229  			wg.Wait()
   230  		})
   231  		t.Run("ConcurrentSharedKeys", func(t *testing.T) {
   232  			m := newMap()
   233  
   234  			// Load up the map.
   235  			for i, s := range testData {
   236  				expectMissing(t, s, 0)(m.Load(s))
   237  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   238  			}
   239  			gmp := runtime.GOMAXPROCS(-1)
   240  			var wg sync.WaitGroup
   241  			for i := range gmp {
   242  				wg.Add(1)
   243  				go func(id int) {
   244  					defer wg.Done()
   245  
   246  					for i, s := range testData {
   247  						expectNotDeleted(t, s, math.MaxInt)(m.CompareAndDelete(s, math.MaxInt))
   248  						m.CompareAndDelete(s, i)
   249  						expectMissing(t, s, 0)(m.Load(s))
   250  					}
   251  					for _, s := range testData {
   252  						expectMissing(t, s, 0)(m.Load(s))
   253  					}
   254  				}(i)
   255  			}
   256  			wg.Wait()
   257  		})
   258  	})
   259  	t.Run("CompareAndSwap", func(t *testing.T) {
   260  		t.Run("All", func(t *testing.T) {
   261  			m := newMap()
   262  
   263  			for i, s := range testData {
   264  				expectMissing(t, s, 0)(m.Load(s))
   265  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   266  				expectPresent(t, s, i)(m.Load(s))
   267  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   268  			}
   269  			for j := range 3 {
   270  				for i, s := range testData {
   271  					expectPresent(t, s, i+j)(m.Load(s))
   272  					expectNotSwapped(t, s, math.MaxInt, i+j+1)(m.CompareAndSwap(s, math.MaxInt, i+j+1))
   273  					expectSwapped(t, s, i, i+j+1)(m.CompareAndSwap(s, i+j, i+j+1))
   274  					expectNotSwapped(t, s, i+j, i+j+1)(m.CompareAndSwap(s, i+j, i+j+1))
   275  					expectPresent(t, s, i+j+1)(m.Load(s))
   276  				}
   277  			}
   278  			for i, s := range testData {
   279  				expectPresent(t, s, i+3)(m.Load(s))
   280  			}
   281  		})
   282  		t.Run("One", func(t *testing.T) {
   283  			m := newMap()
   284  
   285  			for i, s := range testData {
   286  				expectMissing(t, s, 0)(m.Load(s))
   287  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   288  				expectPresent(t, s, i)(m.Load(s))
   289  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   290  			}
   291  			expectNotSwapped(t, testData[15], math.MaxInt, 16)(m.CompareAndSwap(testData[15], math.MaxInt, 16))
   292  			expectSwapped(t, testData[15], 15, 16)(m.CompareAndSwap(testData[15], 15, 16))
   293  			expectNotSwapped(t, testData[15], 15, 16)(m.CompareAndSwap(testData[15], 15, 16))
   294  			for i, s := range testData {
   295  				if i == 15 {
   296  					expectPresent(t, s, 16)(m.Load(s))
   297  				} else {
   298  					expectPresent(t, s, i)(m.Load(s))
   299  				}
   300  			}
   301  		})
   302  		t.Run("Multiple", func(t *testing.T) {
   303  			m := newMap()
   304  
   305  			for i, s := range testData {
   306  				expectMissing(t, s, 0)(m.Load(s))
   307  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   308  				expectPresent(t, s, i)(m.Load(s))
   309  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   310  			}
   311  			for _, i := range []int{1, 105, 6, 85} {
   312  				expectNotSwapped(t, testData[i], math.MaxInt, i+1)(m.CompareAndSwap(testData[i], math.MaxInt, i+1))
   313  				expectSwapped(t, testData[i], i, i+1)(m.CompareAndSwap(testData[i], i, i+1))
   314  				expectNotSwapped(t, testData[i], i, i+1)(m.CompareAndSwap(testData[i], i, i+1))
   315  			}
   316  			for i, s := range testData {
   317  				if i == 1 || i == 105 || i == 6 || i == 85 {
   318  					expectPresent(t, s, i+1)(m.Load(s))
   319  				} else {
   320  					expectPresent(t, s, i)(m.Load(s))
   321  				}
   322  			}
   323  		})
   324  
   325  		t.Run("ConcurrentUnsharedKeys", func(t *testing.T) {
   326  			m := newMap()
   327  
   328  			gmp := runtime.GOMAXPROCS(-1)
   329  			var wg sync.WaitGroup
   330  			for i := range gmp {
   331  				wg.Add(1)
   332  				go func(id int) {
   333  					defer wg.Done()
   334  
   335  					makeKey := func(s string) string {
   336  						return s + "-" + strconv.Itoa(id)
   337  					}
   338  					for _, s := range testData {
   339  						key := makeKey(s)
   340  						expectMissing(t, key, 0)(m.Load(key))
   341  						expectStored(t, key, id)(m.LoadOrStore(key, id))
   342  						expectPresent(t, key, id)(m.Load(key))
   343  						expectLoaded(t, key, id)(m.LoadOrStore(key, 0))
   344  					}
   345  					for _, s := range testData {
   346  						key := makeKey(s)
   347  						expectPresent(t, key, id)(m.Load(key))
   348  						expectSwapped(t, key, id, id+1)(m.CompareAndSwap(key, id, id+1))
   349  						expectPresent(t, key, id+1)(m.Load(key))
   350  					}
   351  					for _, s := range testData {
   352  						key := makeKey(s)
   353  						expectPresent(t, key, id+1)(m.Load(key))
   354  					}
   355  				}(i)
   356  			}
   357  			wg.Wait()
   358  		})
   359  		t.Run("ConcurrentUnsharedKeysWithDelete", func(t *testing.T) {
   360  			m := newMap()
   361  
   362  			gmp := runtime.GOMAXPROCS(-1)
   363  			var wg sync.WaitGroup
   364  			for i := range gmp {
   365  				wg.Add(1)
   366  				go func(id int) {
   367  					defer wg.Done()
   368  
   369  					makeKey := func(s string) string {
   370  						return s + "-" + strconv.Itoa(id)
   371  					}
   372  					for _, s := range testData {
   373  						key := makeKey(s)
   374  						expectMissing(t, key, 0)(m.Load(key))
   375  						expectStored(t, key, id)(m.LoadOrStore(key, id))
   376  						expectPresent(t, key, id)(m.Load(key))
   377  						expectLoaded(t, key, id)(m.LoadOrStore(key, 0))
   378  					}
   379  					for _, s := range testData {
   380  						key := makeKey(s)
   381  						expectPresent(t, key, id)(m.Load(key))
   382  						expectSwapped(t, key, id, id+1)(m.CompareAndSwap(key, id, id+1))
   383  						expectPresent(t, key, id+1)(m.Load(key))
   384  						expectDeleted(t, key, id+1)(m.CompareAndDelete(key, id+1))
   385  						expectNotSwapped(t, key, id+1, id+2)(m.CompareAndSwap(key, id+1, id+2))
   386  						expectNotDeleted(t, key, id+1)(m.CompareAndDelete(key, id+1))
   387  						expectMissing(t, key, 0)(m.Load(key))
   388  					}
   389  					for _, s := range testData {
   390  						key := makeKey(s)
   391  						expectMissing(t, key, 0)(m.Load(key))
   392  					}
   393  				}(i)
   394  			}
   395  			wg.Wait()
   396  		})
   397  		t.Run("ConcurrentSharedKeys", func(t *testing.T) {
   398  			m := newMap()
   399  
   400  			// Load up the map.
   401  			for i, s := range testData {
   402  				expectMissing(t, s, 0)(m.Load(s))
   403  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   404  			}
   405  			gmp := runtime.GOMAXPROCS(-1)
   406  			var wg sync.WaitGroup
   407  			for i := range gmp {
   408  				wg.Add(1)
   409  				go func(id int) {
   410  					defer wg.Done()
   411  
   412  					for i, s := range testData {
   413  						expectNotSwapped(t, s, math.MaxInt, i+1)(m.CompareAndSwap(s, math.MaxInt, i+1))
   414  						m.CompareAndSwap(s, i, i+1)
   415  						expectPresent(t, s, i+1)(m.Load(s))
   416  					}
   417  					for i, s := range testData {
   418  						expectPresent(t, s, i+1)(m.Load(s))
   419  					}
   420  				}(i)
   421  			}
   422  			wg.Wait()
   423  		})
   424  	})
   425  	t.Run("Swap", func(t *testing.T) {
   426  		t.Run("All", func(t *testing.T) {
   427  			m := newMap()
   428  
   429  			for i, s := range testData {
   430  				expectMissing(t, s, 0)(m.Load(s))
   431  				expectNotLoadedFromSwap(t, s, i)(m.Swap(s, i))
   432  				expectPresent(t, s, i)(m.Load(s))
   433  				expectLoadedFromSwap(t, s, i, i)(m.Swap(s, i))
   434  			}
   435  			for j := range 3 {
   436  				for i, s := range testData {
   437  					expectPresent(t, s, i+j)(m.Load(s))
   438  					expectLoadedFromSwap(t, s, i+j, i+j+1)(m.Swap(s, i+j+1))
   439  					expectPresent(t, s, i+j+1)(m.Load(s))
   440  				}
   441  			}
   442  			for i, s := range testData {
   443  				expectLoadedFromSwap(t, s, i+3, i+3)(m.Swap(s, i+3))
   444  			}
   445  		})
   446  		t.Run("One", func(t *testing.T) {
   447  			m := newMap()
   448  
   449  			for i, s := range testData {
   450  				expectMissing(t, s, 0)(m.Load(s))
   451  				expectNotLoadedFromSwap(t, s, i)(m.Swap(s, i))
   452  				expectPresent(t, s, i)(m.Load(s))
   453  				expectLoadedFromSwap(t, s, i, i)(m.Swap(s, i))
   454  			}
   455  			expectLoadedFromSwap(t, testData[15], 15, 16)(m.Swap(testData[15], 16))
   456  			for i, s := range testData {
   457  				if i == 15 {
   458  					expectPresent(t, s, 16)(m.Load(s))
   459  				} else {
   460  					expectPresent(t, s, i)(m.Load(s))
   461  				}
   462  			}
   463  		})
   464  		t.Run("Multiple", func(t *testing.T) {
   465  			m := newMap()
   466  
   467  			for i, s := range testData {
   468  				expectMissing(t, s, 0)(m.Load(s))
   469  				expectNotLoadedFromSwap(t, s, i)(m.Swap(s, i))
   470  				expectPresent(t, s, i)(m.Load(s))
   471  				expectLoadedFromSwap(t, s, i, i)(m.Swap(s, i))
   472  			}
   473  			for _, i := range []int{1, 105, 6, 85} {
   474  				expectLoadedFromSwap(t, testData[i], i, i+1)(m.Swap(testData[i], i+1))
   475  			}
   476  			for i, s := range testData {
   477  				if i == 1 || i == 105 || i == 6 || i == 85 {
   478  					expectPresent(t, s, i+1)(m.Load(s))
   479  				} else {
   480  					expectPresent(t, s, i)(m.Load(s))
   481  				}
   482  			}
   483  		})
   484  		t.Run("ConcurrentUnsharedKeys", func(t *testing.T) {
   485  			m := newMap()
   486  
   487  			gmp := runtime.GOMAXPROCS(-1)
   488  			var wg sync.WaitGroup
   489  			for i := range gmp {
   490  				wg.Add(1)
   491  				go func(id int) {
   492  					defer wg.Done()
   493  
   494  					makeKey := func(s string) string {
   495  						return s + "-" + strconv.Itoa(id)
   496  					}
   497  					for _, s := range testData {
   498  						key := makeKey(s)
   499  						expectMissing(t, key, 0)(m.Load(key))
   500  						expectNotLoadedFromSwap(t, key, id)(m.Swap(key, id))
   501  						expectPresent(t, key, id)(m.Load(key))
   502  						expectLoadedFromSwap(t, key, id, id)(m.Swap(key, id))
   503  					}
   504  					for _, s := range testData {
   505  						key := makeKey(s)
   506  						expectPresent(t, key, id)(m.Load(key))
   507  						expectLoadedFromSwap(t, key, id, id+1)(m.Swap(key, id+1))
   508  						expectPresent(t, key, id+1)(m.Load(key))
   509  					}
   510  					for _, s := range testData {
   511  						key := makeKey(s)
   512  						expectPresent(t, key, id+1)(m.Load(key))
   513  					}
   514  				}(i)
   515  			}
   516  			wg.Wait()
   517  		})
   518  		t.Run("ConcurrentUnsharedKeysWithDelete", func(t *testing.T) {
   519  			m := newMap()
   520  
   521  			gmp := runtime.GOMAXPROCS(-1)
   522  			var wg sync.WaitGroup
   523  			for i := range gmp {
   524  				wg.Add(1)
   525  				go func(id int) {
   526  					defer wg.Done()
   527  
   528  					makeKey := func(s string) string {
   529  						return s + "-" + strconv.Itoa(id)
   530  					}
   531  					for _, s := range testData {
   532  						key := makeKey(s)
   533  						expectMissing(t, key, 0)(m.Load(key))
   534  						expectNotLoadedFromSwap(t, key, id)(m.Swap(key, id))
   535  						expectPresent(t, key, id)(m.Load(key))
   536  						expectLoadedFromSwap(t, key, id, id)(m.Swap(key, id))
   537  					}
   538  					for _, s := range testData {
   539  						key := makeKey(s)
   540  						expectPresent(t, key, id)(m.Load(key))
   541  						expectLoadedFromSwap(t, key, id, id+1)(m.Swap(key, id+1))
   542  						expectPresent(t, key, id+1)(m.Load(key))
   543  						expectDeleted(t, key, id+1)(m.CompareAndDelete(key, id+1))
   544  						expectNotLoadedFromSwap(t, key, id+2)(m.Swap(key, id+2))
   545  						expectPresent(t, key, id+2)(m.Load(key))
   546  					}
   547  					for _, s := range testData {
   548  						key := makeKey(s)
   549  						expectPresent(t, key, id+2)(m.Load(key))
   550  					}
   551  				}(i)
   552  			}
   553  			wg.Wait()
   554  		})
   555  		t.Run("ConcurrentSharedKeys", func(t *testing.T) {
   556  			m := newMap()
   557  
   558  			// Load up the map.
   559  			for i, s := range testData {
   560  				expectMissing(t, s, 0)(m.Load(s))
   561  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   562  			}
   563  			gmp := runtime.GOMAXPROCS(-1)
   564  			var wg sync.WaitGroup
   565  			for i := range gmp {
   566  				wg.Add(1)
   567  				go func(id int) {
   568  					defer wg.Done()
   569  
   570  					for i, s := range testData {
   571  						m.Swap(s, i+1)
   572  						expectPresent(t, s, i+1)(m.Load(s))
   573  					}
   574  					for i, s := range testData {
   575  						expectPresent(t, s, i+1)(m.Load(s))
   576  					}
   577  				}(i)
   578  			}
   579  			wg.Wait()
   580  		})
   581  	})
   582  	t.Run("LoadAndDelete", func(t *testing.T) {
   583  		t.Run("All", func(t *testing.T) {
   584  			m := newMap()
   585  
   586  			for range 3 {
   587  				for i, s := range testData {
   588  					expectMissing(t, s, 0)(m.Load(s))
   589  					expectStored(t, s, i)(m.LoadOrStore(s, i))
   590  					expectPresent(t, s, i)(m.Load(s))
   591  					expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   592  				}
   593  				for i, s := range testData {
   594  					expectPresent(t, s, i)(m.Load(s))
   595  					expectLoadedFromDelete(t, s, i)(m.LoadAndDelete(s))
   596  					expectMissing(t, s, 0)(m.Load(s))
   597  					expectNotLoadedFromDelete(t, s, 0)(m.LoadAndDelete(s))
   598  				}
   599  				for _, s := range testData {
   600  					expectMissing(t, s, 0)(m.Load(s))
   601  				}
   602  			}
   603  		})
   604  		t.Run("One", func(t *testing.T) {
   605  			m := newMap()
   606  
   607  			for i, s := range testData {
   608  				expectMissing(t, s, 0)(m.Load(s))
   609  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   610  				expectPresent(t, s, i)(m.Load(s))
   611  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   612  			}
   613  			expectPresent(t, testData[15], 15)(m.Load(testData[15]))
   614  			expectLoadedFromDelete(t, testData[15], 15)(m.LoadAndDelete(testData[15]))
   615  			expectMissing(t, testData[15], 0)(m.Load(testData[15]))
   616  			expectNotLoadedFromDelete(t, testData[15], 0)(m.LoadAndDelete(testData[15]))
   617  			for i, s := range testData {
   618  				if i == 15 {
   619  					expectMissing(t, s, 0)(m.Load(s))
   620  				} else {
   621  					expectPresent(t, s, i)(m.Load(s))
   622  				}
   623  			}
   624  		})
   625  		t.Run("Multiple", func(t *testing.T) {
   626  			m := newMap()
   627  
   628  			for i, s := range testData {
   629  				expectMissing(t, s, 0)(m.Load(s))
   630  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   631  				expectPresent(t, s, i)(m.Load(s))
   632  				expectLoaded(t, s, i)(m.LoadOrStore(s, 0))
   633  			}
   634  			for _, i := range []int{1, 105, 6, 85} {
   635  				expectPresent(t, testData[i], i)(m.Load(testData[i]))
   636  				expectLoadedFromDelete(t, testData[i], i)(m.LoadAndDelete(testData[i]))
   637  				expectMissing(t, testData[i], 0)(m.Load(testData[i]))
   638  				expectNotLoadedFromDelete(t, testData[i], 0)(m.LoadAndDelete(testData[i]))
   639  			}
   640  			for i, s := range testData {
   641  				if i == 1 || i == 105 || i == 6 || i == 85 {
   642  					expectMissing(t, s, 0)(m.Load(s))
   643  				} else {
   644  					expectPresent(t, s, i)(m.Load(s))
   645  				}
   646  			}
   647  		})
   648  		t.Run("Iterate", func(t *testing.T) {
   649  			m := newMap()
   650  
   651  			testAll(t, m, testDataMap(testData[:]), func(s string, i int) bool {
   652  				expectLoadedFromDelete(t, s, i)(m.LoadAndDelete(s))
   653  				return true
   654  			})
   655  			for _, s := range testData {
   656  				expectMissing(t, s, 0)(m.Load(s))
   657  			}
   658  		})
   659  		t.Run("ConcurrentUnsharedKeys", func(t *testing.T) {
   660  			m := newMap()
   661  
   662  			gmp := runtime.GOMAXPROCS(-1)
   663  			var wg sync.WaitGroup
   664  			for i := range gmp {
   665  				wg.Add(1)
   666  				go func(id int) {
   667  					defer wg.Done()
   668  
   669  					makeKey := func(s string) string {
   670  						return s + "-" + strconv.Itoa(id)
   671  					}
   672  					for _, s := range testData {
   673  						key := makeKey(s)
   674  						expectMissing(t, key, 0)(m.Load(key))
   675  						expectStored(t, key, id)(m.LoadOrStore(key, id))
   676  						expectPresent(t, key, id)(m.Load(key))
   677  						expectLoaded(t, key, id)(m.LoadOrStore(key, 0))
   678  					}
   679  					for _, s := range testData {
   680  						key := makeKey(s)
   681  						expectPresent(t, key, id)(m.Load(key))
   682  						expectLoadedFromDelete(t, key, id)(m.LoadAndDelete(key))
   683  						expectMissing(t, key, 0)(m.Load(key))
   684  					}
   685  					for _, s := range testData {
   686  						key := makeKey(s)
   687  						expectMissing(t, key, 0)(m.Load(key))
   688  					}
   689  				}(i)
   690  			}
   691  			wg.Wait()
   692  		})
   693  		t.Run("ConcurrentSharedKeys", func(t *testing.T) {
   694  			m := newMap()
   695  
   696  			// Load up the map.
   697  			for i, s := range testData {
   698  				expectMissing(t, s, 0)(m.Load(s))
   699  				expectStored(t, s, i)(m.LoadOrStore(s, i))
   700  			}
   701  			gmp := runtime.GOMAXPROCS(-1)
   702  			var wg sync.WaitGroup
   703  			for i := range gmp {
   704  				wg.Add(1)
   705  				go func(id int) {
   706  					defer wg.Done()
   707  
   708  					for _, s := range testData {
   709  						m.LoadAndDelete(s)
   710  						expectMissing(t, s, 0)(m.Load(s))
   711  					}
   712  					for _, s := range testData {
   713  						expectMissing(t, s, 0)(m.Load(s))
   714  					}
   715  				}(i)
   716  			}
   717  			wg.Wait()
   718  		})
   719  	})
   720  }
   721  
   722  func testAll[K, V comparable](t *testing.T, m *isync.HashTrieMap[K, V], testData map[K]V, yield func(K, V) bool) {
   723  	for k, v := range testData {
   724  		expectStored(t, k, v)(m.LoadOrStore(k, v))
   725  	}
   726  	visited := make(map[K]int)
   727  	m.All()(func(key K, got V) bool {
   728  		want, ok := testData[key]
   729  		if !ok {
   730  			t.Errorf("unexpected key %v in map", key)
   731  			return false
   732  		}
   733  		if got != want {
   734  			t.Errorf("expected key %v to have value %v, got %v", key, want, got)
   735  			return false
   736  		}
   737  		visited[key]++
   738  		return yield(key, got)
   739  	})
   740  	for key, n := range visited {
   741  		if n > 1 {
   742  			t.Errorf("visited key %v more than once", key)
   743  		}
   744  	}
   745  }
   746  
   747  func expectPresent[K, V comparable](t *testing.T, key K, want V) func(got V, ok bool) {
   748  	t.Helper()
   749  	return func(got V, ok bool) {
   750  		t.Helper()
   751  
   752  		if !ok {
   753  			t.Errorf("expected key %v to be present in map", key)
   754  		}
   755  		if ok && got != want {
   756  			t.Errorf("expected key %v to have value %v, got %v", key, want, got)
   757  		}
   758  	}
   759  }
   760  
   761  func expectMissing[K, V comparable](t *testing.T, key K, want V) func(got V, ok bool) {
   762  	t.Helper()
   763  	if want != *new(V) {
   764  		// This is awkward, but the want argument is necessary to smooth over type inference.
   765  		// Just make sure the want argument always looks the same.
   766  		panic("expectMissing must always have a zero value variable")
   767  	}
   768  	return func(got V, ok bool) {
   769  		t.Helper()
   770  
   771  		if ok {
   772  			t.Errorf("expected key %v to be missing from map, got value %v", key, got)
   773  		}
   774  		if !ok && got != want {
   775  			t.Errorf("expected missing key %v to be paired with the zero value; got %v", key, got)
   776  		}
   777  	}
   778  }
   779  
   780  func expectLoaded[K, V comparable](t *testing.T, key K, want V) func(got V, loaded bool) {
   781  	t.Helper()
   782  	return func(got V, loaded bool) {
   783  		t.Helper()
   784  
   785  		if !loaded {
   786  			t.Errorf("expected key %v to have been loaded, not stored", key)
   787  		}
   788  		if got != want {
   789  			t.Errorf("expected key %v to have value %v, got %v", key, want, got)
   790  		}
   791  	}
   792  }
   793  
   794  func expectStored[K, V comparable](t *testing.T, key K, want V) func(got V, loaded bool) {
   795  	t.Helper()
   796  	return func(got V, loaded bool) {
   797  		t.Helper()
   798  
   799  		if loaded {
   800  			t.Errorf("expected inserted key %v to have been stored, not loaded", key)
   801  		}
   802  		if got != want {
   803  			t.Errorf("expected inserted key %v to have value %v, got %v", key, want, got)
   804  		}
   805  	}
   806  }
   807  
   808  func expectDeleted[K, V comparable](t *testing.T, key K, old V) func(deleted bool) {
   809  	t.Helper()
   810  	return func(deleted bool) {
   811  		t.Helper()
   812  
   813  		if !deleted {
   814  			t.Errorf("expected key %v with value %v to be in map and deleted", key, old)
   815  		}
   816  	}
   817  }
   818  
   819  func expectNotDeleted[K, V comparable](t *testing.T, key K, old V) func(deleted bool) {
   820  	t.Helper()
   821  	return func(deleted bool) {
   822  		t.Helper()
   823  
   824  		if deleted {
   825  			t.Errorf("expected key %v with value %v to not be in map and thus not deleted", key, old)
   826  		}
   827  	}
   828  }
   829  
   830  func expectSwapped[K, V comparable](t *testing.T, key K, old, new V) func(swapped bool) {
   831  	t.Helper()
   832  	return func(swapped bool) {
   833  		t.Helper()
   834  
   835  		if !swapped {
   836  			t.Errorf("expected key %v with value %v to be in map and swapped for %v", key, old, new)
   837  		}
   838  	}
   839  }
   840  
   841  func expectNotSwapped[K, V comparable](t *testing.T, key K, old, new V) func(swapped bool) {
   842  	t.Helper()
   843  	return func(swapped bool) {
   844  		t.Helper()
   845  
   846  		if swapped {
   847  			t.Errorf("expected key %v with value %v to not be in map or not swapped for %v", key, old, new)
   848  		}
   849  	}
   850  }
   851  
   852  func expectLoadedFromSwap[K, V comparable](t *testing.T, key K, want, new V) func(got V, loaded bool) {
   853  	t.Helper()
   854  	return func(got V, loaded bool) {
   855  		t.Helper()
   856  
   857  		if !loaded {
   858  			t.Errorf("expected key %v to be in map and for %v to have been swapped for %v", key, want, new)
   859  		} else if want != got {
   860  			t.Errorf("key %v had its value %v swapped for %v, but expected it to have value %v", key, got, new, want)
   861  		}
   862  	}
   863  }
   864  
   865  func expectNotLoadedFromSwap[K, V comparable](t *testing.T, key K, new V) func(old V, loaded bool) {
   866  	t.Helper()
   867  	return func(old V, loaded bool) {
   868  		t.Helper()
   869  
   870  		if loaded {
   871  			t.Errorf("expected key %v to not be in map, but found value %v for it", key, old)
   872  		}
   873  	}
   874  }
   875  
   876  func expectLoadedFromDelete[K, V comparable](t *testing.T, key K, want V) func(got V, loaded bool) {
   877  	t.Helper()
   878  	return func(got V, loaded bool) {
   879  		t.Helper()
   880  
   881  		if !loaded {
   882  			t.Errorf("expected key %v to be in map to be deleted", key)
   883  		} else if want != got {
   884  			t.Errorf("key %v was deleted with value %v, but expected it to have value %v", key, got, want)
   885  		}
   886  	}
   887  }
   888  
   889  func expectNotLoadedFromDelete[K, V comparable](t *testing.T, key K, _ V) func(old V, loaded bool) {
   890  	t.Helper()
   891  	return func(old V, loaded bool) {
   892  		t.Helper()
   893  
   894  		if loaded {
   895  			t.Errorf("expected key %v to not be in map, but found value %v for it", key, old)
   896  		}
   897  	}
   898  }
   899  
   900  func testDataMap(data []string) map[string]int {
   901  	m := make(map[string]int)
   902  	for i, s := range data {
   903  		m[s] = i
   904  	}
   905  	return m
   906  }
   907  
   908  var (
   909  	testDataSmall [8]string
   910  	testData      [128]string
   911  	testDataLarge [128 << 10]string
   912  )
   913  
   914  func init() {
   915  	for i := range testDataSmall {
   916  		testDataSmall[i] = fmt.Sprintf("%b", i)
   917  	}
   918  	for i := range testData {
   919  		testData[i] = fmt.Sprintf("%b", i)
   920  	}
   921  	for i := range testDataLarge {
   922  		testDataLarge[i] = fmt.Sprintf("%b", i)
   923  	}
   924  }
   925  
   926  // TestConcurrentCache tests HashTrieMap in a scenario where it is used as
   927  // the basis of a memory-efficient concurrent cache. We're specifically
   928  // looking to make sure that CompareAndSwap and CompareAndDelete are
   929  // atomic with respect to one another. When competing for the same
   930  // key-value pair, they must not both succeed.
   931  //
   932  // This test is a regression test for issue #70970.
   933  func TestConcurrentCache(t *testing.T) {
   934  	type dummy [32]byte
   935  
   936  	var m isync.HashTrieMap[int, weak.Pointer[dummy]]
   937  
   938  	type cleanupArg struct {
   939  		key   int
   940  		value weak.Pointer[dummy]
   941  	}
   942  	cleanup := func(arg cleanupArg) {
   943  		m.CompareAndDelete(arg.key, arg.value)
   944  	}
   945  	get := func(m *isync.HashTrieMap[int, weak.Pointer[dummy]], key int) *dummy {
   946  		nv := new(dummy)
   947  		nw := weak.Make(nv)
   948  		for {
   949  			w, loaded := m.LoadOrStore(key, nw)
   950  			if !loaded {
   951  				runtime.AddCleanup(nv, cleanup, cleanupArg{key, nw})
   952  				return nv
   953  			}
   954  			if v := w.Value(); v != nil {
   955  				return v
   956  			}
   957  
   958  			// Weak pointer was reclaimed, try to replace it with nw.
   959  			if m.CompareAndSwap(key, w, nw) {
   960  				runtime.AddCleanup(nv, cleanup, cleanupArg{key, nw})
   961  				return nv
   962  			}
   963  		}
   964  	}
   965  
   966  	const N = 100_000
   967  	const P = 5_000
   968  
   969  	var wg sync.WaitGroup
   970  	wg.Add(N)
   971  	for i := range N {
   972  		go func() {
   973  			defer wg.Done()
   974  			a := get(&m, i%P)
   975  			b := get(&m, i%P)
   976  			if a != b {
   977  				t.Errorf("consecutive cache reads returned different values: a != b (%p vs %p)\n", a, b)
   978  			}
   979  		}()
   980  	}
   981  	wg.Wait()
   982  }
   983  

View as plain text