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
16package runtime
17
18import (
19	"internal/cpu"
20	"internal/goarch"
21	"internal/runtime/atomic"
22	"runtime/internal/sys"
23	"unsafe"
24)
25
26type 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).
46type 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.
57func (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.
67func (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
115func (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.
136func (tab *traceMap) reset() {
137	tab.root.Store(nil)
138	tab.seq.Store(0)
139	tab.mem.drop()
140}
141