Implementing a Lock-Free B-Link Tree in Go for Concurrent Indexing — code snippets

snippet-1.go

package blink

import (
	"bytes"
	"sync"
	"sync/atomic"
	"unsafe"
)

// Key represents the indexing key.
type Key []byte

// Compare returns -1 if k < other, 0 if k == other, 1 if k > other.
func (k Key) Compare(other Key) int {
	return bytes.Compare(k, other)
}

// NodeState holds the immutable snapshot of a node's keys and pointers.
// This allows lock-free reads via atomic pointer swaps.
type NodeState struct {
	Keys    []Key
	Ptrs    []unsafe.Pointer // If leaf, values; if internal, child *Node pointers
	HighKey Key              // Keys in this node are <= HighKey (except if rightmost)
}

// Node represents a single node in the B-Link tree.
type Node struct {
	IsLeaf bool
	Mu     sync.Mutex     // Protects writes and splits
	State  unsafe.Pointer // Points to *NodeState; read lock-freely
	Right  unsafe.Pointer // Points to the right-sibling *Node; read lock-freely
}

// GetState returns the current NodeState pointer atomically.
func (n *Node) GetState() *NodeState {
	return (*NodeState)(atomic.LoadPointer(&n.State))
}

// GetRight returns the right-sibling Node pointer atomically.
func (n *Node) GetRight() *Node {
	return (*Node)(atomic.LoadPointer(&n.Right))
}

snippet-2.go

type BLinkTree struct {
	root    unsafe.Pointer // *Node
	maxKeys int
	rootMu  sync.Mutex     // Protects root replacement
}

func (t *BLinkTree) getRoot() *Node {
	return (*Node)(atomic.LoadPointer(&t.root))
}

// Search traverses the B-Link tree lock-freely to find the value for the given key.
func (t *BLinkTree) Search(key Key) (unsafe.Pointer, bool) {
	curr := t.getRoot()
	if curr == nil {
		return nil, false
	}

	for {
		// 1. Move right if the key is greater than the high key of the current node.
		// This handles concurrent splits without requiring parent backtracking.
		for {
			state := curr.GetState()
			if state.HighKey != nil && key.Compare(state.HighKey) > 0 {
				right := curr.GetRight()
				if right == nil {
					break
				}
				curr = right
				continue
			}
			break
		}

		state := curr.GetState()
		if curr.IsLeaf {
			idx, found := binarySearch(state.Keys, key)
			if found {
				return atomic.LoadPointer(&state.Ptrs[idx]), true
			}
			return nil, false
		}

		// Internal node traversal: find the appropriate child pointer.
		idx := findChildIndex(state.Keys, key)
		childPtr := atomic.LoadPointer(&state.Ptrs[idx])
		curr = (*Node)(childPtr)
	}
}

func binarySearch(keys []Key, key Key) (int, bool) {
	low, high := 0, len(keys)-1
	for low <= high {
		mid := (low + high) / 2
		cmp := key.Compare(keys[mid])
		if cmp == 0 {
			return mid, true
		} else if cmp < 0 {
			high = mid - 1
		} else {
			low = mid + 1
		}
	}
	return low, false
}

func findChildIndex(keys []Key, key Key) int {
	low, high := 0, len(keys)-1
	idx := len(keys) - 1
	for low <= high {
		mid := (low + high) / 2
		if keys[mid].Compare(key) >= 0 {
			idx = mid
			high = mid - 1
		} else {
			low = mid + 1
		}
	}
	return idx
}

snippet-3.go

// splitNode splits a full node, returns the new right-sibling Node, and the split key to insert into the parent.
func (t *BLinkTree) splitNode(node *Node, state *NodeState) (*Node, Key) {
	mid := len(state.Keys) / 2
	splitKey := state.Keys[mid-1]

	// Create the new right sibling node
	rightNode := &Node{
		IsLeaf: node.IsLeaf,
	}

	// Right sibling gets the second half of the elements
	rightState := &NodeState{
		Keys:    make([]Key, len(state.Keys)-mid),
		Ptrs:    make([]unsafe.Pointer, len(state.Ptrs)-mid),
		HighKey: state.HighKey, // Inherits the original node's high key
	}
	copy(rightState.Keys, state.Keys[mid:])
	copy(rightState.Ptrs, state.Ptrs[mid:])
	rightNode.State = unsafe.Pointer(rightState)

	// The new node points to the old node's right sibling
	rightNode.Right = node.Right

	// Left node (current) keeps the first half of the elements
	leftState := &NodeState{
		Keys:    make([]Key, mid),
		Ptrs:    make([]unsafe.Pointer, mid),
		HighKey: splitKey, // The split key becomes the new high key of the left node
	}
	copy(leftState.Keys, state.Keys[:mid])
	copy(leftState.Ptrs, state.Ptrs[:mid])

	// Crucial order of operations for memory visibility:
	// 1. Link the right node in first.
	// 2. Swapping the state pointer makes the split active.
	atomic.StorePointer(&node.Right, unsafe.Pointer(rightNode))
	atomic.StorePointer(&node.State, unsafe.Pointer(leftState))

	return rightNode, splitKey
}

snippet-4.go

// Insert inserts a key-value pair into the B-Link tree.
func (t *BLinkTree) Insert(key Key, value unsafe.Pointer) {
	for {
		path, leaf := t.findLeafForWrite(key)
		leaf.Mu.Lock()

		// Verify if the leaf is still the correct node (right-moving correction)
		leaf = t.moveRightForWrite(leaf, key)

		state := leaf.GetState()
		idx, found := binarySearch(state.Keys, key)
		if found {
			// Key already exists, update the value
			newState := &NodeState{
				Keys:    state.Keys,
				Ptrs:    make([]unsafe.Pointer, len(state.Ptrs)),
				HighKey: state.HighKey,
			}
			copy(newState.Ptrs, state.Ptrs)
			newState.Ptrs[idx] = value
			atomic.StorePointer(&leaf.State, unsafe.Pointer(newState))
			leaf.Mu.Unlock()
			return
		}

		if len(state.Keys) < t.maxKeys {
			// Room for insertion: insert key and value in order
			newState := insertIntoState(state, idx, key, value)
			atomic.StorePointer(&leaf.State, unsafe.Pointer(newState))
			leaf.Mu.Unlock()
			return
		}

		// Node is full, split is required
		rightNode, splitKey := t.splitNode(leaf, state)
		
		// Insert key into the appropriate split node
		if key.Compare(splitKey) <= 0 {
			t.insertInNodeLocked(leaf, key, value)
		} else {
			t.insertInNodeLocked(rightNode, key, value)
		}

		leaf.Mu.Unlock()

		// Propagate split to parent
		t.insertIntoParent(path, splitKey, unsafe.Pointer(rightNode))
		return
	}
}

func (t *BLinkTree) moveRightForWrite(curr *Node, key Key) *Node {
	for {
		state := curr.GetState()
		if state.HighKey != nil && key.Compare(state.HighKey) > 0 {
			right := curr.GetRight()
			if right == nil {
				break
			}
			curr.Mu.Unlock()
			right.Mu.Lock()
			curr = right
			continue
		}
		break
	}
	return curr
}

func insertIntoState(state *NodeState, idx int, key Key, value unsafe.Pointer) *NodeState {
	newKeys := make([]Key, len(state.Keys)+1)
	newPtrs := make([]unsafe.Pointer, len(state.Ptrs)+1)
	
	copy(newKeys[:idx], state.Keys[:idx])
	newKeys[idx] = key
	copy(newKeys[idx+1:], state.Keys[idx:])
	
	copy(newPtrs[:idx], state.Ptrs[:idx])
	newPtrs[idx] = value
	copy(newPtrs[idx+1:], state.Ptrs[idx:])
	
	return &NodeState{
		Keys:    newKeys,
		Ptrs:    newPtrs,
		HighKey: state.HighKey,
	}
}

func (t *BLinkTree) insertInNodeLocked(node *Node, key Key, val unsafe.Pointer) {
	state := node.GetState()
	idx, _ := binarySearch(state.Keys, key)
	newState := insertIntoState(state, idx, key, val)
	atomic.StorePointer(&node.State, unsafe.Pointer(newState))
}

func (t *BLinkTree) findLeafForWrite(key Key) ([]*Node, *Node) {
	var path []*Node
	curr := t.getRoot()
	if curr == nil {
		return nil, nil
	}

	for {
		// Traverse right first if needed
		for {
			state := curr.GetState()
			if state.HighKey != nil && key.Compare(state.HighKey) > 0 {
				right := curr.GetRight()
				if right == nil {
					break
				}
				curr = right
				continue
			}
			break
		}

		if curr.IsLeaf {
			return path, curr
		}

		path = append(path, curr)
		state := curr.GetState()
		idx := findChildIndex(state.Keys, key)
		curr = (*Node)(atomic.LoadPointer(&state.Ptrs[idx]))
	}
}

snippet-5.go

// insertIntoParent propagates a split up the tree.
func (t *BLinkTree) insertIntoParent(path []*Node, key Key, child unsafe.Pointer) {
	if len(path) == 0 {
		t.createNewRoot(key, child)
		return
	}

	parent := path[len(path)-1]
	parentPath := path[:len(path)-1]

	parent.Mu.Lock()
	parent = t.moveRightForWrite(parent, key)
	state := parent.GetState()

	if len(state.Keys) < t.maxKeys {
		t.insertInNodeLocked(parent, key, child)
		parent.Mu.Unlock()
		return
	}

	// Parent is full, must split parent node
	rightParent, splitKey := t.splitNode(parent, state)

	if key.Compare(splitKey) <= 0 {
		t.insertInNodeLocked(parent, key, child)
	} else {
		t.insertInNodeLocked(rightParent, key, child)
	}

	parent.Mu.Unlock()

	// Recursive call to propagate split upwards
	t.insertIntoParent(parentPath, splitKey, unsafe.Pointer(rightParent))
}

func (t *BLinkTree) createNewRoot(key Key, child unsafe.Pointer) {
	t.rootMu.Lock()
	defer t.rootMu.Unlock()

	currRoot := t.getRoot()
	if t.getRoot() != currRoot {
		// Root changed while acquiring lock, retry insertion with new path
		t.insertIntoParent([]*Node{t.getRoot()}, key, child)
		return
	}

	newRoot := &Node{
		IsLeaf: false,
	}
	
	state := &NodeState{
		Keys:    []Key{key},
		Ptrs:    []unsafe.Pointer{unsafe.Pointer(currRoot), child},
		HighKey: nil, // Root has no high key
	}
	newRoot.State = unsafe.Pointer(state)

	atomic.StorePointer(&t.root, unsafe.Pointer(newRoot))
}

snippet-6.go

package blink

import (
	"sync"
	"sync/atomic"
	"unsafe"
)

// EpochReclaimer manages safe memory reclamation for concurrent lock-free reads.
type EpochReclaimer struct {
	epoch       uint64
	activeReads int64
	limboLists  [3][]unsafe.Pointer // 3 epochs are sufficient for EBR (current, previous, old)
	mu          sync.Mutex
}

func NewEpochReclaimer() *EpochReclaimer {
	return &EpochReclaimer{}
}

// Enter flags that a reader is entering a critical section.
func (er *EpochReclaimer) Enter() uint64 {
	atomic.AddInt64(&er.activeReads, 1)
	return atomic.LoadUint64(&er.epoch)
}

// Exit flags that a reader has exited a critical section.
func (er *EpochReclaimer) Exit() {
	atomic.AddInt64(&er.activeReads, -1)
}

// Retire places a pointer in the current epoch's limbo list.
func (er *EpochReclaimer) Retire(ptr unsafe.Pointer) {
	er.mu.Lock()
	defer er.mu.Unlock()
	
	currEpoch := atomic.LoadUint64(&er.epoch)
	idx := currEpoch % 3
	er.limboLists[idx] = append(er.limboLists[idx], ptr)

	// If no active readers, advance the epoch and reclaim old memory
	if atomic.LoadInt64(&er.activeReads) == 0 {
		er.advanceEpochLocked()
	}
}

func (er *EpochReclaimer) advanceEpochLocked() {
	nextEpoch := atomic.AddUint64(&er.epoch, 1)
	reclaimIdx := (nextEpoch - 2) % 3
	for _, ptr := range er.limboLists[reclaimIdx] {
		state := (*NodeState)(ptr)
		releaseState(state)
	}
	er.limboLists[reclaimIdx] = er.limboLists[reclaimIdx][:0]
}

var statePool = sync.Pool{
	New: func() interface{} {
		return &NodeState{}
	},
}

func releaseState(state *NodeState) {
	// Zero out pointers to avoid retaining referenced objects in memory
	for i := range state.Ptrs {
		state.Ptrs[i] = nil
	}
	statePool.Put(state)
}