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