lbcd/database/internal/treap/immutable.go
Jonathan Moody 0ab37e1541 Distinguish between Mutable/Immutable treap node recycling, and make Mutable
treap recycle nodes consistently. Rework Immutable treap node recycling
attempting to make it safer in the presence of code that takes snapshots
(dbCacheSnapshot) of the treap. Add special mutable PutM and DeleteM
methods which DB transaction can use to apply changes more efficiently
without creating lots of garbage memory.
2022-06-03 11:47:27 -04:00

541 lines
16 KiB
Go

// Copyright (c) 2015-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.
package treap
import (
"bytes"
"math/rand"
"sync"
)
// cloneTreapNode returns a shallow copy of the passed node.
func cloneTreapNode(node *treapNode) *treapNode {
clone := getTreapNode(node.key, node.value, node.priority, node.generation+1)
clone.left = node.left
clone.right = node.right
return clone
}
// Immutable represents a treap data structure which is used to hold ordered
// key/value pairs using a combination of binary search tree and heap semantics.
// It is a self-organizing and randomized data structure that doesn't require
// complex operations to maintain balance. Search, insert, and delete
// operations are all O(log n). In addition, it provides O(1) snapshots for
// multi-version concurrency control (MVCC).
//
// All operations which result in modifying the treap return a new version of
// the treap with only the modified nodes updated. All unmodified nodes are
// shared with the previous version. This is extremely useful in concurrent
// applications since the caller only has to atomically replace the treap
// pointer with the newly returned version after performing any mutations. All
// readers can simply use their existing pointer as a snapshot since the treap
// it points to is immutable. This effectively provides O(1) snapshot
// capability with efficient memory usage characteristics since the old nodes
// only remain allocated until there are no longer any references to them.
type Immutable struct {
root *treapNode
count int
// totalSize is the best estimate of the total size of of all data in
// the treap including the keys, values, and node sizes.
totalSize uint64
// generation number starts at 0 after NewImmutable(), and
// is incremented with every Put()/Delete().
generation int
// snap is a pointer to a node in snapshot history linked list.
// A value nil means no snapshots are outstanding.
snap *SnapRecord
}
// newImmutable returns a new immutable treap given the passed parameters.
func newImmutable(root *treapNode, count int, totalSize uint64, generation int, snap *SnapRecord) *Immutable {
return &Immutable{root: root, count: count, totalSize: totalSize, generation: generation, snap: snap}
}
// Len returns the number of items stored in the treap.
func (t *Immutable) Len() int {
return t.count
}
// Size returns a best estimate of the total number of bytes the treap is
// consuming including all of the fields used to represent the nodes as well as
// the size of the keys and values. Shared values are not detected, so the
// returned size assumes each value is pointing to different memory.
func (t *Immutable) Size() uint64 {
return t.totalSize
}
// get returns the treap node that contains the passed key. It will return nil
// when the key does not exist.
func (t *Immutable) get(key []byte) *treapNode {
for node := t.root; node != nil; {
// Traverse left or right depending on the result of the
// comparison.
compareResult := bytes.Compare(key, node.key)
if compareResult < 0 {
node = node.left
continue
}
if compareResult > 0 {
node = node.right
continue
}
// The key exists.
return node
}
// A nil node was reached which means the key does not exist.
return nil
}
// Has returns whether or not the passed key exists.
func (t *Immutable) Has(key []byte) bool {
if node := t.get(key); node != nil {
return true
}
return false
}
// Get returns the value for the passed key. The function will return nil when
// the key does not exist.
func (t *Immutable) Get(key []byte) []byte {
if node := t.get(key); node != nil {
return node.value
}
return nil
}
// put inserts the passed key/value pair.
func (t *Immutable) put(key, value []byte, bumpGen int) (tp *Immutable, old parentStack) {
// Use an empty byte slice for the value when none was provided. This
// ultimately allows key existence to be determined from the value since
// an empty byte slice is distinguishable from nil.
if value == nil {
value = emptySlice
}
// The node is the root of the tree if there isn't already one.
if t.root == nil {
root := getTreapNode(key, value, rand.Int(), t.generation+bumpGen)
return newImmutable(root, 1, nodeSize(root), t.generation+bumpGen, t.snap), parentStack{}
}
// Find the binary tree insertion point and construct a replaced list of
// parents while doing so. This is done because this is an immutable
// data structure so regardless of where in the treap the new key/value
// pair ends up, all ancestors up to and including the root need to be
// replaced.
//
// When the key matches an entry already in the treap, replace the node
// with a new one that has the new value set and return.
var parents parentStack
var oldParents parentStack
var compareResult int
for node := t.root; node != nil; {
// Clone the node and link its parent to it if needed.
oldParents.Push(node)
nodeCopy := cloneTreapNode(node)
if oldParent := parents.At(0); oldParent != nil {
if oldParent.left == node {
oldParent.left = nodeCopy
} else {
oldParent.right = nodeCopy
}
}
parents.Push(nodeCopy)
// Traverse left or right depending on the result of comparing
// the keys.
compareResult = bytes.Compare(key, node.key)
if compareResult < 0 {
node = node.left
continue
}
if compareResult > 0 {
node = node.right
continue
}
// The key already exists, so update its value.
nodeCopy.value = value
// Return new immutable treap with the replaced node and
// ancestors up to and including the root of the tree.
newRoot := parents.At(parents.Len() - 1)
newTotalSize := t.totalSize - uint64(len(node.value)) +
uint64(len(value))
return newImmutable(newRoot, t.count, newTotalSize, t.generation+bumpGen, t.snap), oldParents
}
// Link the new node into the binary tree in the correct position.
node := getTreapNode(key, value, rand.Int(), t.generation+bumpGen)
parent := parents.At(0)
if compareResult < 0 {
parent.left = node
} else {
parent.right = node
}
// Perform any rotations needed to maintain the min-heap and replace
// the ancestors up to and including the tree root.
newRoot := parents.At(parents.Len() - 1)
for parents.Len() > 0 {
// There is nothing left to do when the node's priority is
// greater than or equal to its parent's priority.
parent = parents.Pop()
if node.priority >= parent.priority {
break
}
// Perform a right rotation if the node is on the left side or
// a left rotation if the node is on the right side.
if parent.left == node {
node.right, parent.left = parent, node.right
} else {
node.left, parent.right = parent, node.left
}
// Either set the new root of the tree when there is no
// grandparent or relink the grandparent to the node based on
// which side the old parent the node is replacing was on.
grandparent := parents.At(0)
if grandparent == nil {
newRoot = node
} else if grandparent.left == parent {
grandparent.left = node
} else {
grandparent.right = node
}
}
return newImmutable(newRoot, t.count+1, t.totalSize+nodeSize(node), t.generation+bumpGen, t.snap), oldParents
}
// Put is the immutable variant of put. Generation number is bumped, and old
// nodes become garbage unless referenced elswhere.
func (t *Immutable) Put(key, value []byte) *Immutable {
tp, _ := t.put(key, value, 1)
return tp
}
// PutM is the mutable variant of put. Generation number is NOT bumped, and old
// nodes are recycled if possible. This is only safe/useful in scenarios where
// multiple Put/Delete() ops are applied to a unique treap and no snapshots/aliases
// of the intermediate treap states are created or desired. For example:
//
// for i := range keys {
// t = t.Put(keys[i])
// }
//
// ...may be replaced with:
//
// for i := range keys {
// PutM(t, keys[i], nil)
// }
//
// If "excluded" is provided, that snapshot is ignored when counting
// snapshot records.
//
func PutM(dest **Immutable, key, value []byte, excluded *SnapRecord) {
tp, old := (*dest).put(key, value, 0)
// Examine old nodes and recycle if possible.
snapRecordMutex.Lock()
defer snapRecordMutex.Unlock()
snapCount := (*dest).snapCount(excluded)
for old.Len() > 0 {
node := old.Pop()
if node.generation == tp.generation && snapCount == 0 {
putTreapNode(node)
}
}
*dest = tp
}
// del removes the passed key from the treap and returns the resulting treap
// if it exists. The original immutable treap is returned if the key does not
// exist.
func (t *Immutable) del(key []byte, bumpGen int) (d *Immutable, old parentStack) {
// Find the node for the key while constructing a list of parents while
// doing so.
var oldParents parentStack
var delNode *treapNode
for node := t.root; node != nil; {
oldParents.Push(node)
// Traverse left or right depending on the result of the
// comparison.
compareResult := bytes.Compare(key, node.key)
if compareResult < 0 {
node = node.left
continue
}
if compareResult > 0 {
node = node.right
continue
}
// The key exists.
delNode = node
break
}
// There is nothing to do if the key does not exist.
if delNode == nil {
return t, parentStack{}
}
// When the only node in the tree is the root node and it is the one
// being deleted, there is nothing else to do besides removing it.
parent := oldParents.At(1)
if parent == nil && delNode.left == nil && delNode.right == nil {
return newImmutable(nil, 0, 0, t.generation+bumpGen, t.snap), oldParents
}
// Construct a replaced list of parents and the node to delete itself.
// This is done because this is an immutable data structure and
// therefore all ancestors of the node that will be deleted, up to and
// including the root, need to be replaced.
var newParents parentStack
for i := oldParents.Len(); i > 0; i-- {
node := oldParents.At(i - 1)
nodeCopy := cloneTreapNode(node)
if oldParent := newParents.At(0); oldParent != nil {
if oldParent.left == node {
oldParent.left = nodeCopy
} else {
oldParent.right = nodeCopy
}
}
newParents.Push(nodeCopy)
}
delNode = newParents.Pop()
parent = newParents.At(0)
// Perform rotations to move the node to delete to a leaf position while
// maintaining the min-heap while replacing the modified children.
var child *treapNode
newRoot := newParents.At(newParents.Len() - 1)
for delNode.left != nil || delNode.right != nil {
// Choose the child with the higher priority.
var isLeft bool
if delNode.left == nil {
child = delNode.right
} else if delNode.right == nil {
child = delNode.left
isLeft = true
} else if delNode.left.priority >= delNode.right.priority {
child = delNode.left
isLeft = true
} else {
child = delNode.right
}
// Rotate left or right depending on which side the child node
// is on. This has the effect of moving the node to delete
// towards the bottom of the tree while maintaining the
// min-heap.
child = cloneTreapNode(child)
if isLeft {
child.right, delNode.left = delNode, child.right
} else {
child.left, delNode.right = delNode, child.left
}
// Either set the new root of the tree when there is no
// grandparent or relink the grandparent to the node based on
// which side the old parent the node is replacing was on.
//
// Since the node to be deleted was just moved down a level, the
// new grandparent is now the current parent and the new parent
// is the current child.
if parent == nil {
newRoot = child
} else if parent.left == delNode {
parent.left = child
} else {
parent.right = child
}
// The parent for the node to delete is now what was previously
// its child.
parent = child
}
// Delete the node, which is now a leaf node, by disconnecting it from
// its parent.
if parent.right == delNode {
parent.right = nil
} else {
parent.left = nil
}
return newImmutable(newRoot, t.count-1, t.totalSize-nodeSize(delNode), t.generation+bumpGen, t.snap), oldParents
}
// Delete is the immutable variant of del. Generation number is bumped, and old
// nodes become garbage unless referenced elswhere.
func (t *Immutable) Delete(key []byte) *Immutable {
tp, _ := t.del(key, 1)
return tp
}
// DeleteM is the mutable variant of del. Generation number is NOT bumped, and old
// nodes are recycled if possible. This is only safe/useful in scenarios where
// multiple Put/Delete() ops are applied to a unique treap and no snapshots/aliases
// of the intermediate treap states are created or desired. For example:
//
// for i := range keys {
// t = t.Delete(keys[i])
// }
//
// ...may be replaced with:
//
// for i := range keys {
// DeleteM(t, keys[i], nil)
// }
//
// If "excluded" is provided, that snapshot is ignored when counting
// snapshot records.
//
func DeleteM(dest **Immutable, key []byte, excluded *SnapRecord) {
tp, old := (*dest).del(key, 0)
// Examine old nodes and recycle if possible.
snapRecordMutex.Lock()
defer snapRecordMutex.Unlock()
snapCount := (*dest).snapCount(excluded)
for old.Len() > 0 {
node := old.Pop()
if node.generation == tp.generation && snapCount == 0 {
putTreapNode(node)
}
}
*dest = tp
}
// ForEach invokes the passed function with every key/value pair in the treap
// in ascending order.
func (t *Immutable) ForEach(fn func(k, v []byte) bool) {
// Add the root node and all children to the left of it to the list of
// nodes to traverse and loop until they, and all of their child nodes,
// have been traversed.
var parents parentStack
for node := t.root; node != nil; node = node.left {
parents.Push(node)
}
for parents.Len() > 0 {
node := parents.Pop()
if !fn(node.key, node.value) {
return
}
// Extend the nodes to traverse by all children to the left of
// the current node's right child.
for node := node.right; node != nil; node = node.left {
parents.Push(node)
}
}
}
// NewImmutable returns a new empty immutable treap ready for use. See the
// documentation for the Immutable structure for more details.
func NewImmutable() *Immutable {
return &Immutable{}
}
// SnapRecord assists in tracking/releasing outstanding snapshots.
type SnapRecord struct {
prev *SnapRecord
next *SnapRecord
}
var snapRecordMutex sync.Mutex
// Snapshot makes a SnapRecord and linkis it into the snapshot history of a treap.
func (t *Immutable) Snapshot() *SnapRecord {
snapRecordMutex.Lock()
defer snapRecordMutex.Unlock()
// Link this record so it follows the existing t.snap record, if any.
prev := t.snap
var next *SnapRecord = nil
if prev != nil {
next = prev.next
}
t.snap = &SnapRecord{prev: prev, next: next}
if prev != nil {
prev.next = t.snap
}
return t.snap
}
// Release of SnapRecord unlinks that record from the snapshot history of a treap.
func (r *SnapRecord) Release() {
snapRecordMutex.Lock()
defer snapRecordMutex.Unlock()
// Unlink this record.
if r.prev != nil {
r.prev.next = r.next
}
if r.next != nil {
r.next.prev = r.prev
}
}
// snapCount returns the number of snapshots outstanding which were created
// but not released. When snapshots are absent, mutable PutM()/DeleteM() can
// recycle nodes more aggressively. The record "exclude" is not counted.
func (t *Immutable) snapCount(exclude *SnapRecord) int {
// snapRecordMutex should be locked already
sum := 0
if t.snap == nil {
// No snapshots.
return sum
}
// Count snapshots taken BEFORE creation of this instance.
for h := t.snap; h != nil; h = h.prev {
if h != exclude {
sum++
}
}
// Count snapshots taken AFTER creation of this instance.
for h := t.snap.next; h != nil; h = h.next {
if h != exclude {
sum++
}
}
return sum
}
func (t *Immutable) Recycle() {
snapCount := t.snapCount(nil) - 1
var parents parentStack
for node := t.root; node != nil; node = node.left {
parents.Push(node)
}
for parents.Len() > 0 {
node := parents.Pop()
// Extend the nodes to traverse by all children to the left of
// the current node's right child.
for n := node.right; n != nil; n = n.left {
parents.Push(n)
}
if node.generation == t.generation && snapCount == 0 {
putTreapNode(node)
}
}
}