Source file src/unique/handle.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 unique 6 7 import ( 8 "internal/abi" 9 "internal/concurrent" 10 "internal/weak" 11 "runtime" 12 "sync" 13 _ "unsafe" 14 ) 15 16 // Handle is a globally unique identity for some value of type T. 17 // 18 // Two handles compare equal exactly if the two values used to create the handles 19 // would have also compared equal. The comparison of two handles is trivial and 20 // typically much more efficient than comparing the values used to create them. 21 type Handle[T comparable] struct { 22 value *T 23 } 24 25 // Value returns a shallow copy of the T value that produced the Handle. 26 func (h Handle[T]) Value() T { 27 return *h.value 28 } 29 30 // Make returns a globally unique handle for a value of type T. Handles 31 // are equal if and only if the values used to produce them are equal. 32 func Make[T comparable](value T) Handle[T] { 33 // Find the map for type T. 34 typ := abi.TypeOf(value) 35 ma, ok := uniqueMaps.Load(typ) 36 if !ok { 37 // This is a good time to initialize cleanup, since we must go through 38 // this path on the first use of Make, and it's not on the hot path. 39 setupMake.Do(registerCleanup) 40 ma = addUniqueMap[T](typ) 41 } 42 m := ma.(*uniqueMap[T]) 43 44 // Keep around any values we allocate for insertion. There 45 // are a few different ways we can race with other threads 46 // and create values that we might discard. By keeping 47 // the first one we make around, we can avoid generating 48 // more than one per racing thread. 49 var ( 50 toInsert *T // Keep this around to keep it alive. 51 toInsertWeak weak.Pointer[T] 52 ) 53 newValue := func() weak.Pointer[T] { 54 if toInsert == nil { 55 toInsert = new(T) 56 *toInsert = clone(value, &m.cloneSeq) 57 toInsertWeak = weak.Make(toInsert) 58 } 59 return toInsertWeak 60 } 61 var ptr *T 62 for { 63 // Check the map. 64 wp, ok := m.Load(value) 65 if !ok { 66 // Try to insert a new value into the map. 67 wp, _ = m.LoadOrStore(value, newValue()) 68 } 69 // Now that we're sure there's a value in the map, let's 70 // try to get the pointer we need out of it. 71 ptr = wp.Strong() 72 if ptr != nil { 73 break 74 } 75 // The weak pointer is nil, so the old value is truly dead. 76 // Try to remove it and start over. 77 m.CompareAndDelete(value, wp) 78 } 79 runtime.KeepAlive(toInsert) 80 return Handle[T]{ptr} 81 } 82 83 var ( 84 // uniqueMaps is an index of type-specific concurrent maps used for unique.Make. 85 // 86 // The two-level map might seem odd at first since the HashTrieMap could have "any" 87 // as its key type, but the issue is escape analysis. We do not want to force lookups 88 // to escape the argument, and using a type-specific map allows us to avoid that where 89 // possible (for example, for strings and plain-ol'-data structs). We also get the 90 // benefit of not cramming every different type into a single map, but that's certainly 91 // not enough to outweigh the cost of two map lookups. What is worth it though, is saving 92 // on those allocations. 93 uniqueMaps = concurrent.NewHashTrieMap[*abi.Type, any]() // any is always a *uniqueMap[T]. 94 95 // cleanupFuncs are functions that clean up dead weak pointers in type-specific 96 // maps in uniqueMaps. We express cleanup this way because there's no way to iterate 97 // over the sync.Map and call functions on the type-specific data structures otherwise. 98 // These cleanup funcs each close over one of these type-specific maps. 99 // 100 // cleanupMu protects cleanupNotify and is held across the entire cleanup. Used for testing. 101 // cleanupNotify is a test-only mechanism that allow tests to wait for the cleanup to run. 102 cleanupMu sync.Mutex 103 cleanupFuncsMu sync.Mutex 104 cleanupFuncs []func() 105 cleanupNotify []func() // One-time notifications when cleanups finish. 106 ) 107 108 type uniqueMap[T comparable] struct { 109 *concurrent.HashTrieMap[T, weak.Pointer[T]] 110 cloneSeq 111 } 112 113 func addUniqueMap[T comparable](typ *abi.Type) *uniqueMap[T] { 114 // Create a map for T and try to register it. We could 115 // race with someone else, but that's fine; it's one 116 // small, stray allocation. The number of allocations 117 // this can create is bounded by a small constant. 118 m := &uniqueMap[T]{ 119 HashTrieMap: concurrent.NewHashTrieMap[T, weak.Pointer[T]](), 120 cloneSeq: makeCloneSeq(typ), 121 } 122 a, loaded := uniqueMaps.LoadOrStore(typ, m) 123 if !loaded { 124 // Add a cleanup function for the new map. 125 cleanupFuncsMu.Lock() 126 cleanupFuncs = append(cleanupFuncs, func() { 127 // Delete all the entries whose weak references are nil and clean up 128 // deleted entries. 129 m.Enumerate(func(key T, wp weak.Pointer[T]) bool { 130 if wp.Strong() == nil { 131 m.CompareAndDelete(key, wp) 132 } 133 return true 134 }) 135 }) 136 cleanupFuncsMu.Unlock() 137 } 138 return a.(*uniqueMap[T]) 139 } 140 141 // setupMake is used to perform initial setup for unique.Make. 142 var setupMake sync.Once 143 144 // startBackgroundCleanup sets up a background goroutine to occasionally call cleanupFuncs. 145 func registerCleanup() { 146 runtime_registerUniqueMapCleanup(func() { 147 // Lock for cleanup. 148 cleanupMu.Lock() 149 150 // Grab funcs to run. 151 cleanupFuncsMu.Lock() 152 cf := cleanupFuncs 153 cleanupFuncsMu.Unlock() 154 155 // Run cleanup. 156 for _, f := range cf { 157 f() 158 } 159 160 // Run cleanup notifications. 161 for _, f := range cleanupNotify { 162 f() 163 } 164 cleanupNotify = nil 165 166 // Finished. 167 cleanupMu.Unlock() 168 }) 169 } 170 171 // Implemented in runtime. 172 173 //go:linkname runtime_registerUniqueMapCleanup 174 func runtime_registerUniqueMapCleanup(cleanup func()) 175