tracemap.go

     1  // Copyright 2023 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  // Simple append-only thread-safe hash map for tracing.
     6  // Provides a mapping between variable-length data and a
     7  // unique ID. Subsequent puts of the same data will return
     8  // the same ID. The zero value is ready to use.
     9  //
    10  // Uses a region-based allocation scheme internally, and
    11  // reset clears the whole map.
    12  //
    13  // It avoids doing any high-level Go operations so it's safe
    14  // to use even in sensitive contexts.
    15  
    16  package runtime
    17  
    18  import (
    19  	"internal/cpu"
    20  	"internal/goarch"
    21  	"internal/runtime/atomic"
    22  	"internal/runtime/sys"
    23  	"unsafe"
    24  )
    25  
    26  type traceMap struct {
    27  	root atomic.UnsafePointer // *traceMapNode (can't use generics because it's notinheap)
    28  	_    cpu.CacheLinePad
    29  	seq  atomic.Uint64
    30  	_    cpu.CacheLinePad
    31  	mem  traceRegionAlloc
    32  }
    33  
    34  // traceMapNode is an implementation of a lock-free append-only hash-trie
    35  // (a trie of the hash bits).
    36  //
    37  // Key features:
    38  //   - 4-ary trie. Child nodes are indexed by the upper 2 (remaining) bits of the hash.
    39  //     For example, top level uses bits [63:62], next level uses [61:60] and so on.
    40  //   - New nodes are placed at the first empty level encountered.
    41  //   - When the first child is added to a node, the existing value is not moved into a child.
    42  //     This means that you must check the key at each level, not just at the leaf.
    43  //   - No deletion or rebalancing.
    44  //   - Intentionally devolves into a linked list on hash collisions (the hash bits will all
    45  //     get shifted out during iteration, and new nodes will just be appended to the 0th child).
    46  type traceMapNode struct {
    47  	_ sys.NotInHeap
    48  
    49  	children [4]atomic.UnsafePointer // *traceMapNode (can't use generics because it's notinheap)
    50  	hash     uintptr
    51  	id       uint64
    52  	data     []byte
    53  }
    54  
    55  // stealID steals an ID from the table, ensuring that it will not
    56  // appear in the table anymore.
    57  func (tab *traceMap) stealID() uint64 {
    58  	return tab.seq.Add(1)
    59  }
    60  
    61  // put inserts the data into the table.
    62  //
    63  // It's always safe for callers to noescape data because put copies its bytes.
    64  //
    65  // Returns a unique ID for the data and whether this is the first time
    66  // the data has been added to the map.
    67  func (tab *traceMap) put(data unsafe.Pointer, size uintptr) (uint64, bool) {
    68  	if size == 0 {
    69  		return 0, false
    70  	}
    71  	hash := memhash(data, 0, size)
    72  
    73  	var newNode *traceMapNode
    74  	m := &tab.root
    75  	hashIter := hash
    76  	for {
    77  		n := (*traceMapNode)(m.Load())
    78  		if n == nil {
    79  			// Try to insert a new map node. We may end up discarding
    80  			// this node if we fail to insert because it turns out the
    81  			// value is already in the map.
    82  			//
    83  			// The discard will only happen if two threads race on inserting
    84  			// the same value. Both might create nodes, but only one will
    85  			// succeed on insertion. If two threads race to insert two
    86  			// different values, then both nodes will *always* get inserted,
    87  			// because the equality checking below will always fail.
    88  			//
    89  			// Performance note: contention on insertion is likely to be
    90  			// higher for small maps, but since this data structure is
    91  			// append-only, either the map stays small because there isn't
    92  			// much activity, or the map gets big and races to insert on
    93  			// the same node are much less likely.
    94  			if newNode == nil {
    95  				newNode = tab.newTraceMapNode(data, size, hash, tab.seq.Add(1))
    96  			}
    97  			if m.CompareAndSwapNoWB(nil, unsafe.Pointer(newNode)) {
    98  				return newNode.id, true
    99  			}
   100  			// Reload n. Because pointers are only stored once,
   101  			// we must have lost the race, and therefore n is not nil
   102  			// anymore.
   103  			n = (*traceMapNode)(m.Load())
   104  		}
   105  		if n.hash == hash && uintptr(len(n.data)) == size {
   106  			if memequal(unsafe.Pointer(&n.data[0]), data, size) {
   107  				return n.id, false
   108  			}
   109  		}
   110  		m = &n.children[hashIter>>(8*goarch.PtrSize-2)]
   111  		hashIter <<= 2
   112  	}
   113  }
   114  
   115  func (tab *traceMap) newTraceMapNode(data unsafe.Pointer, size, hash uintptr, id uint64) *traceMapNode {
   116  	// Create data array.
   117  	sl := notInHeapSlice{
   118  		array: tab.mem.alloc(size),
   119  		len:   int(size),
   120  		cap:   int(size),
   121  	}
   122  	memmove(unsafe.Pointer(sl.array), data, size)
   123  
   124  	// Create metadata structure.
   125  	meta := (*traceMapNode)(unsafe.Pointer(tab.mem.alloc(unsafe.Sizeof(traceMapNode{}))))
   126  	*(*notInHeapSlice)(unsafe.Pointer(&meta.data)) = sl
   127  	meta.id = id
   128  	meta.hash = hash
   129  	return meta
   130  }
   131  
   132  // reset drops all allocated memory from the table and resets it.
   133  //
   134  // The caller must ensure that there are no put operations executing concurrently
   135  // with this function.
   136  func (tab *traceMap) reset() {
   137  	tab.root.Store(nil)
   138  	tab.seq.Store(0)
   139  	tab.mem.drop()
   140  }
   141
View as plain text