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
View as plain text