canonmap_test.go

     1  // Copyright 2025 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 unique
     6  
     7  import (
     8  	"internal/abi"
     9  	"runtime"
    10  	"strconv"
    11  	"sync"
    12  	"testing"
    13  	"unsafe"
    14  )
    15  
    16  func TestCanonMap(t *testing.T) {
    17  	testCanonMap(t, func() *canonMap[string] {
    18  		return newCanonMap[string]()
    19  	})
    20  }
    21  
    22  func TestCanonMapBadHash(t *testing.T) {
    23  	testCanonMap(t, func() *canonMap[string] {
    24  		return newBadCanonMap[string]()
    25  	})
    26  }
    27  
    28  func TestCanonMapTruncHash(t *testing.T) {
    29  	testCanonMap(t, func() *canonMap[string] {
    30  		// Stub out the good hash function with a different terrible one
    31  		// (truncated hash). Everything should still work as expected.
    32  		// This is useful to test independently to catch issues with
    33  		// near collisions, where only the last few bits of the hash differ.
    34  		return newTruncCanonMap[string]()
    35  	})
    36  }
    37  
    38  func testCanonMap(t *testing.T, newMap func() *canonMap[string]) {
    39  	t.Run("LoadEmpty", func(t *testing.T) {
    40  		m := newMap()
    41  
    42  		for _, s := range testData {
    43  			expectMissing(t, s)(m.Load(s))
    44  		}
    45  	})
    46  	t.Run("LoadOrStore", func(t *testing.T) {
    47  		t.Run("Sequential", func(t *testing.T) {
    48  			m := newMap()
    49  
    50  			var refs []*string
    51  			for _, s := range testData {
    52  				expectMissing(t, s)(m.Load(s))
    53  				refs = append(refs, expectPresent(t, s)(m.LoadOrStore(s)))
    54  				expectPresent(t, s)(m.Load(s))
    55  				expectPresent(t, s)(m.LoadOrStore(s))
    56  			}
    57  			drainCleanupQueue(t)
    58  			for _, s := range testData {
    59  				expectPresent(t, s)(m.Load(s))
    60  				expectPresent(t, s)(m.LoadOrStore(s))
    61  			}
    62  			runtime.KeepAlive(refs)
    63  			refs = nil
    64  			drainCleanupQueue(t)
    65  			for _, s := range testData {
    66  				expectMissing(t, s)(m.Load(s))
    67  				expectPresent(t, s)(m.LoadOrStore(s))
    68  			}
    69  		})
    70  		t.Run("ConcurrentUnsharedKeys", func(t *testing.T) {
    71  			makeKey := func(s string, id int) string {
    72  				return s + "-" + strconv.Itoa(id)
    73  			}
    74  
    75  			// Expand and shrink the map multiple times to try to get
    76  			// insertions and cleanups to overlap.
    77  			m := newMap()
    78  			gmp := runtime.GOMAXPROCS(-1)
    79  			for try := range 3 {
    80  				var wg sync.WaitGroup
    81  				for i := range gmp {
    82  					wg.Add(1)
    83  					go func(id int) {
    84  						defer wg.Done()
    85  
    86  						var refs []*string
    87  						for _, s := range testData {
    88  							key := makeKey(s, id)
    89  							if try == 0 {
    90  								expectMissing(t, key)(m.Load(key))
    91  							}
    92  							refs = append(refs, expectPresent(t, key)(m.LoadOrStore(key)))
    93  							expectPresent(t, key)(m.Load(key))
    94  							expectPresent(t, key)(m.LoadOrStore(key))
    95  						}
    96  						for i, s := range testData {
    97  							key := makeKey(s, id)
    98  							expectPresent(t, key)(m.Load(key))
    99  							if got, want := expectPresent(t, key)(m.LoadOrStore(key)), refs[i]; got != want {
   100  								t.Errorf("canonical entry %p did not match ref %p", got, want)
   101  							}
   102  						}
   103  						// N.B. We avoid trying to test entry cleanup here
   104  						// because it's going to be very flaky, especially
   105  						// in the bad hash cases.
   106  					}(i)
   107  				}
   108  				wg.Wait()
   109  			}
   110  
   111  			// Run an extra GC cycle to de-flake. Sometimes the cleanups
   112  			// fail to run in time, despite drainCleanupQueue.
   113  			//
   114  			// TODO(mknyszek): Figure out why the extra GC is necessary,
   115  			// and what is transiently keeping the cleanups live.
   116  			// * I have confirmed that they are not completely stuck, and
   117  			//   they always eventually run.
   118  			// * I have also confirmed it's not asynchronous preemption
   119  			//   keeping them around (though that is a possibility).
   120  			// * I have confirmed that they are not simply sitting on
   121  			//   the queue, and that drainCleanupQueue is just failing
   122  			//   to actually empty the queue.
   123  			// * I have confirmed that it's not a write barrier that's
   124  			//   keeping it alive, nor is it a weak pointer dereference
   125  			//   (which shades the object during the GC).
   126  			// The corresponding objects do seem to be transiently truly
   127  			// reachable, but I have no idea by what path.
   128  			runtime.GC()
   129  
   130  			// Drain cleanups so everything is deleted.
   131  			drainCleanupQueue(t)
   132  
   133  			// Double-check that it's all gone.
   134  			for id := range gmp {
   135  				makeKey := func(s string) string {
   136  					return s + "-" + strconv.Itoa(id)
   137  				}
   138  				for _, s := range testData {
   139  					key := makeKey(s)
   140  					expectMissing(t, key)(m.Load(key))
   141  				}
   142  			}
   143  		})
   144  	})
   145  }
   146  
   147  func expectMissing[T comparable](t *testing.T, key T) func(got *T) {
   148  	t.Helper()
   149  	return func(got *T) {
   150  		t.Helper()
   151  
   152  		if got != nil {
   153  			t.Errorf("expected key %v to be missing from map, got %p", key, got)
   154  		}
   155  	}
   156  }
   157  
   158  func expectPresent[T comparable](t *testing.T, key T) func(got *T) *T {
   159  	t.Helper()
   160  	return func(got *T) *T {
   161  		t.Helper()
   162  
   163  		if got == nil {
   164  			t.Errorf("expected key %v to be present in map, got %p", key, got)
   165  		}
   166  		if got != nil && *got != key {
   167  			t.Errorf("key %v is present in map, but canonical version has the wrong value: got %v, want %v", key, *got, key)
   168  		}
   169  		return got
   170  	}
   171  }
   172  
   173  // newBadCanonMap creates a new canonMap for the provided key type
   174  // but with an intentionally bad hash function.
   175  func newBadCanonMap[T comparable]() *canonMap[T] {
   176  	// Stub out the good hash function with a terrible one.
   177  	// Everything should still work as expected.
   178  	m := newCanonMap[T]()
   179  	m.hash = func(_ unsafe.Pointer, _ uintptr) uintptr {
   180  		return 0
   181  	}
   182  	return m
   183  }
   184  
   185  // newTruncCanonMap creates a new canonMap for the provided key type
   186  // but with an intentionally bad hash function.
   187  func newTruncCanonMap[T comparable]() *canonMap[T] {
   188  	// Stub out the good hash function with a terrible one.
   189  	// Everything should still work as expected.
   190  	m := newCanonMap[T]()
   191  	var mx map[string]int
   192  	mapType := abi.TypeOf(mx).MapType()
   193  	hasher := mapType.Hasher
   194  	m.hash = func(p unsafe.Pointer, n uintptr) uintptr {
   195  		return hasher(p, n) & ((uintptr(1) << 4) - 1)
   196  	}
   197  	return m
   198  }
   199
View as plain text