github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/gtl/rcu_map.go.tpl (about)

     1  package PACKAGE
     2  
     3  import (
     4  	"sync/atomic"
     5  	"unsafe"
     6  )
     7  
     8  // ZZMap is a concurrent map. A reader can access the map without lock,
     9  // regardless of background updates.  The writer side must coordinate using an
    10  // external mutex if there are multiple writers. This map is linearizable.
    11  //
    12  // Example:
    13  //
    14  //   m := NewZZMap(10)
    15  //   go func() {  // writer
    16  //     m.Store("foo", "bar")
    17  //   }()
    18  //   go func() {  // reader
    19  //     val, ok := m.Load("foo")
    20  //   }
    21  type ZZMap struct {
    22  	p unsafe.Pointer // *zzMapState
    23  }
    24  
    25  // ZZMapState represents a fixed-size chained hash table. It can store up to
    26  // maxCapacity key/value pairs.  Beyond that, the caller must create a new
    27  // ZZMapState with a larger capacity.
    28  type zzMapState struct {
    29  	log2Len     uint             // ==log2(len(table))
    30  	mask        uint64           // == ^(log2Len-1)
    31  	table       []unsafe.Pointer // *zzMapNode
    32  	n           int              // # of objects currently stored in the table
    33  	maxCapacity int              // max # of object that can be stored
    34  }
    35  
    36  // ZZMapNode represents a hash bucket.
    37  type zzMapNode struct {
    38  	key   KEY
    39  	value VALUE
    40  
    41  	// next points to the next element in the same hash bucket
    42  	next unsafe.Pointer // *zzMapNode
    43  }
    44  
    45  func newZZMapState(log2Len uint) *zzMapState {
    46  	len := int(1 << log2Len)
    47  	table := &zzMapState{
    48  		log2Len:     log2Len,
    49  		mask:        uint64(log2Len - 1),
    50  		table:       make([]unsafe.Pointer, 1<<log2Len),
    51  		maxCapacity: int(float64(len) * 0.8),
    52  	}
    53  	if table.maxCapacity < len {
    54  		table.maxCapacity = len
    55  	}
    56  	return table
    57  }
    58  
    59  // NewZZMap creates a new map. Arg initialLenHint suggests the the initial
    60  // capacity.  If you plan to store 100 keys, then pass 100 as the value. If you
    61  // don't know the capacity, pass 0 as initialLenHint.
    62  func NewZZMap(initialLenHint int) *ZZMap {
    63  	log2Len := uint(3) // 8 nodes
    64  	for (1 << log2Len) < initialLenHint {
    65  		if log2Len > 31 {
    66  			// TODO(saito) We could make the table to grow larger than 32 bits, but
    67  			// doing so will break 32bit builds.
    68  			panic(initialLenHint)
    69  		}
    70  		log2Len++
    71  	}
    72  	m := ZZMap{p: unsafe.Pointer(newZZMapState(log2Len))}
    73  	return &m
    74  }
    75  
    76  // Load finds a value with the given key. Returns false if not found.
    77  func (m *ZZMap) Load(key KEY) (VALUE, bool) {
    78  	hash := HASH(key)
    79  	table := (*zzMapState)(atomic.LoadPointer(&m.p))
    80  	b := int(hash & table.mask)
    81  	node := (*zzMapNode)(atomic.LoadPointer(&table.table[b]))
    82  	for node != nil {
    83  		if node.key == key {
    84  			return node.value, true
    85  		}
    86  		node = (*zzMapNode)(atomic.LoadPointer(&node.next))
    87  	}
    88  	var dummy VALUE
    89  	return dummy, false
    90  }
    91  
    92  // store returns false iff the table needs resizing.
    93  func (t *zzMapState) store(key KEY, value VALUE) bool {
    94  	var (
    95  		hash     = HASH(key)
    96  		b        = int(hash & t.mask)
    97  		node     = (*zzMapNode)(t.table[b])
    98  		probeLen = 0
    99  		prevNode *zzMapNode
   100  	)
   101  	for node != nil {
   102  		if node.key == key {
   103  			newNode := *node
   104  			newNode.value = value
   105  			if prevNode == nil {
   106  				atomic.StorePointer(&t.table[b], unsafe.Pointer(&newNode))
   107  			} else {
   108  				atomic.StorePointer(&prevNode.next, unsafe.Pointer(&newNode))
   109  			}
   110  			return true
   111  		}
   112  		prevNode = node
   113  		node = (*zzMapNode)(node.next)
   114  		probeLen++
   115  		if probeLen >= 4 && t.n >= t.maxCapacity {
   116  			return false
   117  		}
   118  	}
   119  	newNode := zzMapNode{key: key, value: value}
   120  	if prevNode == nil {
   121  		atomic.StorePointer(&t.table[b], unsafe.Pointer(&newNode))
   122  	} else {
   123  		atomic.StorePointer(&prevNode.next, unsafe.Pointer(&newNode))
   124  	}
   125  	t.n++
   126  	return true
   127  }
   128  
   129  // Store stores the value for the given key. If the key is already in the map,
   130  // it updates the mapping to the given value.
   131  //
   132  // Caution: if Store() is going to be called concurrently, it must be serialized
   133  // externally.
   134  func (m *ZZMap) Store(key KEY, value VALUE) {
   135  	table := (*zzMapState)(atomic.LoadPointer(&m.p))
   136  	if table.store(key, value) {
   137  		return
   138  	}
   139  	log2Len := table.log2Len + 1
   140  	if log2Len > 31 {
   141  		panic(log2Len)
   142  	}
   143  	newTable := newZZMapState(log2Len)
   144  	// Copy the contents of the old table over to the new table.
   145  	for _, p := range table.table {
   146  		node := (*zzMapNode)(p)
   147  		for node != nil {
   148  			if !newTable.store(node.key, node.value) {
   149  				panic(node)
   150  			}
   151  			node = (*zzMapNode)(node.next)
   152  		}
   153  	}
   154  	if !newTable.store(key, value) {
   155  		panic(key)
   156  	}
   157  	atomic.StorePointer(&m.p, unsafe.Pointer(newTable))
   158  }