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.TypeFor[T]()
    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() (T, 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 *toInsert, 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  			k, v := newValue()
    68  			wp, _ = m.LoadOrStore(k, v)
    69  		}
    70  		// Now that we're sure there's a value in the map, let's
    71  		// try to get the pointer we need out of it.
    72  		ptr = wp.Strong()
    73  		if ptr != nil {
    74  			break
    75  		}
    76  		// The weak pointer is nil, so the old value is truly dead.
    77  		// Try to remove it and start over.
    78  		m.CompareAndDelete(value, wp)
    79  	}
    80  	runtime.KeepAlive(toInsert)
    81  	return Handle[T]{ptr}
    82  }
    83  
    84  var (
    85  	// uniqueMaps is an index of type-specific concurrent maps used for unique.Make.
    86  	//
    87  	// The two-level map might seem odd at first since the HashTrieMap could have "any"
    88  	// as its key type, but the issue is escape analysis. We do not want to force lookups
    89  	// to escape the argument, and using a type-specific map allows us to avoid that where
    90  	// possible (for example, for strings and plain-ol'-data structs). We also get the
    91  	// benefit of not cramming every different type into a single map, but that's certainly
    92  	// not enough to outweigh the cost of two map lookups. What is worth it though, is saving
    93  	// on those allocations.
    94  	uniqueMaps = concurrent.NewHashTrieMap[*abi.Type, any]() // any is always a *uniqueMap[T].
    95  
    96  	// cleanupFuncs are functions that clean up dead weak pointers in type-specific
    97  	// maps in uniqueMaps. We express cleanup this way because there's no way to iterate
    98  	// over the sync.Map and call functions on the type-specific data structures otherwise.
    99  	// These cleanup funcs each close over one of these type-specific maps.
   100  	//
   101  	// cleanupMu protects cleanupNotify and is held across the entire cleanup. Used for testing.
   102  	// cleanupNotify is a test-only mechanism that allow tests to wait for the cleanup to run.
   103  	cleanupMu      sync.Mutex
   104  	cleanupFuncsMu sync.Mutex
   105  	cleanupFuncs   []func()
   106  	cleanupNotify  []func() // One-time notifications when cleanups finish.
   107  )
   108  
   109  type uniqueMap[T comparable] struct {
   110  	*concurrent.HashTrieMap[T, weak.Pointer[T]]
   111  	cloneSeq
   112  }
   113  
   114  func addUniqueMap[T comparable](typ *abi.Type) *uniqueMap[T] {
   115  	// Create a map for T and try to register it. We could
   116  	// race with someone else, but that's fine; it's one
   117  	// small, stray allocation. The number of allocations
   118  	// this can create is bounded by a small constant.
   119  	m := &uniqueMap[T]{
   120  		HashTrieMap: concurrent.NewHashTrieMap[T, weak.Pointer[T]](),
   121  		cloneSeq:    makeCloneSeq(typ),
   122  	}
   123  	a, loaded := uniqueMaps.LoadOrStore(typ, m)
   124  	if !loaded {
   125  		// Add a cleanup function for the new map.
   126  		cleanupFuncsMu.Lock()
   127  		cleanupFuncs = append(cleanupFuncs, func() {
   128  			// Delete all the entries whose weak references are nil and clean up
   129  			// deleted entries.
   130  			m.All()(func(key T, wp weak.Pointer[T]) bool {
   131  				if wp.Strong() == nil {
   132  					m.CompareAndDelete(key, wp)
   133  				}
   134  				return true
   135  			})
   136  		})
   137  		cleanupFuncsMu.Unlock()
   138  	}
   139  	return a.(*uniqueMap[T])
   140  }
   141  
   142  // setupMake is used to perform initial setup for unique.Make.
   143  var setupMake sync.Once
   144  
   145  // startBackgroundCleanup sets up a background goroutine to occasionally call cleanupFuncs.
   146  func registerCleanup() {
   147  	runtime_registerUniqueMapCleanup(func() {
   148  		// Lock for cleanup.
   149  		cleanupMu.Lock()
   150  
   151  		// Grab funcs to run.
   152  		cleanupFuncsMu.Lock()
   153  		cf := cleanupFuncs
   154  		cleanupFuncsMu.Unlock()
   155  
   156  		// Run cleanup.
   157  		for _, f := range cf {
   158  			f()
   159  		}
   160  
   161  		// Run cleanup notifications.
   162  		for _, f := range cleanupNotify {
   163  			f()
   164  		}
   165  		cleanupNotify = nil
   166  
   167  		// Finished.
   168  		cleanupMu.Unlock()
   169  	})
   170  }
   171  
   172  // Implemented in runtime.
   173  
   174  //go:linkname runtime_registerUniqueMapCleanup
   175  func runtime_registerUniqueMapCleanup(cleanup func())
   176
View as plain text