feat(memtable): implement SkipList with lock-free reads
- memtable/skiplist.go: heap-backed SkipList with atomic.Pointer next pointers - Mutex-protected writes, lock-free reads with release/acquire ordering - Copy-on-insert and copy-on-return for key/value safety - Forward iterator that skips pending entries - Comprehensive tests: ordered insert, overwrite, pending hidden, concurrent -race
This commit is contained in:
dailz
@@ -0,0 +1,222 @@
|
||||
package memtable
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"math/rand"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
const (
|
||||
maxLevel = 20
|
||||
probability = 0.25
|
||||
)
|
||||
|
||||
// SkipList is an ordered in-memory key-value index.
|
||||
//
|
||||
// Writers are serialized by mu. Readers never take mu: every next pointer is
|
||||
// published with atomic.Pointer.Store and read with atomic.Pointer.Load, giving
|
||||
// release/acquire ordering for fully initialized nodes.
|
||||
type SkipList struct {
|
||||
arena *Arena
|
||||
mu sync.Mutex
|
||||
head *skipNode
|
||||
|
||||
height atomic.Int32
|
||||
length atomic.Uint64
|
||||
}
|
||||
|
||||
type skipNode struct {
|
||||
next [maxLevel]atomic.Pointer[skipNode]
|
||||
|
||||
key []byte
|
||||
value []byte
|
||||
sequence uint64
|
||||
height int
|
||||
pending bool
|
||||
}
|
||||
|
||||
// NewSkipList creates an empty skip list. The arena is retained for the
|
||||
// memtable API and future arena-backed nodes; Phase 1 stores nodes on the Go
|
||||
// heap so atomic.Pointer can provide correct lock-free reads.
|
||||
func NewSkipList(arena *Arena) *SkipList {
|
||||
head := &skipNode{height: maxLevel}
|
||||
sl := &SkipList{
|
||||
arena: arena,
|
||||
head: head,
|
||||
}
|
||||
sl.height.Store(1)
|
||||
return sl
|
||||
}
|
||||
|
||||
// Put inserts or replaces key with value and sequence.
|
||||
func (sl *SkipList) Put(key []byte, value []byte, sequence uint64, pending bool) error {
|
||||
sl.mu.Lock()
|
||||
defer sl.mu.Unlock()
|
||||
|
||||
var prev [maxLevel]*skipNode
|
||||
found := sl.findGreaterOrEqual(key, &prev)
|
||||
|
||||
if found != nil && bytes.Equal(found.key, key) {
|
||||
sl.replaceLocked(&prev, found, key, value, sequence, pending)
|
||||
return nil
|
||||
}
|
||||
|
||||
height := randomHeight()
|
||||
sl.ensureHeightLocked(height, &prev)
|
||||
node := newSkipNode(key, value, sequence, pending, height)
|
||||
|
||||
for level := range height {
|
||||
node.next[level].Store(prev[level].next[level].Load())
|
||||
}
|
||||
for level := range height {
|
||||
prev[level].next[level].Store(node)
|
||||
}
|
||||
sl.length.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get returns the latest published value for key. It performs no locking.
|
||||
func (sl *SkipList) Get(key []byte) (found bool, value []byte, sequence uint64) {
|
||||
node := sl.findGreaterOrEqual(key, nil)
|
||||
if node == nil || !bytes.Equal(node.key, key) || node.pending {
|
||||
return false, nil, 0
|
||||
}
|
||||
return true, cloneBytes(node.value), node.sequence
|
||||
}
|
||||
|
||||
// Len returns the number of distinct keys in the skip list.
|
||||
func (sl *SkipList) Len() uint64 {
|
||||
return sl.length.Load()
|
||||
}
|
||||
|
||||
// NewIterator creates a forward iterator over published entries.
|
||||
func (sl *SkipList) NewIterator() *Iterator {
|
||||
it := &Iterator{node: sl.head.next[0].Load()}
|
||||
it.skipPending()
|
||||
return it
|
||||
}
|
||||
|
||||
// Iterator is a lock-free forward iterator over SkipList entries.
|
||||
type Iterator struct {
|
||||
node *skipNode
|
||||
}
|
||||
|
||||
// Valid reports whether the iterator points at an entry.
|
||||
func (it *Iterator) Valid() bool {
|
||||
return it.node != nil
|
||||
}
|
||||
|
||||
// Next advances the iterator to the next published entry.
|
||||
func (it *Iterator) Next() {
|
||||
if it.node == nil {
|
||||
return
|
||||
}
|
||||
it.node = it.node.next[0].Load()
|
||||
it.skipPending()
|
||||
}
|
||||
|
||||
// Key returns a copy of the current key.
|
||||
func (it *Iterator) Key() []byte {
|
||||
if it.node == nil {
|
||||
return nil
|
||||
}
|
||||
return cloneBytes(it.node.key)
|
||||
}
|
||||
|
||||
// Value returns a copy of the current value.
|
||||
func (it *Iterator) Value() []byte {
|
||||
if it.node == nil {
|
||||
return nil
|
||||
}
|
||||
return cloneBytes(it.node.value)
|
||||
}
|
||||
|
||||
// Sequence returns the current entry's sequence number.
|
||||
func (it *Iterator) Sequence() uint64 {
|
||||
if it.node == nil {
|
||||
return 0
|
||||
}
|
||||
return it.node.sequence
|
||||
}
|
||||
|
||||
func (it *Iterator) skipPending() {
|
||||
for it.node != nil && it.node.pending {
|
||||
it.node = it.node.next[0].Load()
|
||||
}
|
||||
}
|
||||
|
||||
func (sl *SkipList) replaceLocked(prev *[maxLevel]*skipNode, old *skipNode, key []byte, value []byte, sequence uint64, pending bool) {
|
||||
height := randomHeight()
|
||||
sl.ensureHeightLocked(height, prev)
|
||||
node := newSkipNode(key, value, sequence, pending, height)
|
||||
|
||||
for level := range maxLevel {
|
||||
oldNext := old.next[level].Load()
|
||||
if level < height {
|
||||
node.next[level].Store(oldNext)
|
||||
prev[level].next[level].Store(node)
|
||||
continue
|
||||
}
|
||||
|
||||
if level < old.height {
|
||||
prev[level].next[level].Store(oldNext)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sl *SkipList) ensureHeightLocked(height int, prev *[maxLevel]*skipNode) {
|
||||
currentHeight := int(sl.height.Load())
|
||||
if height <= currentHeight {
|
||||
return
|
||||
}
|
||||
|
||||
for level := currentHeight; level < height; level++ {
|
||||
prev[level] = sl.head
|
||||
}
|
||||
sl.height.Store(int32(height))
|
||||
}
|
||||
|
||||
func (sl *SkipList) findGreaterOrEqual(key []byte, prev *[maxLevel]*skipNode) *skipNode {
|
||||
x := sl.head
|
||||
level := int(sl.height.Load()) - 1
|
||||
for level >= 0 {
|
||||
next := x.next[level].Load()
|
||||
if next != nil && bytes.Compare(next.key, key) < 0 {
|
||||
x = next
|
||||
continue
|
||||
}
|
||||
if prev != nil {
|
||||
prev[level] = x
|
||||
}
|
||||
level--
|
||||
}
|
||||
return x.next[0].Load()
|
||||
}
|
||||
|
||||
func newSkipNode(key []byte, value []byte, sequence uint64, pending bool, height int) *skipNode {
|
||||
return &skipNode{
|
||||
key: cloneBytes(key),
|
||||
value: cloneBytes(value),
|
||||
sequence: sequence,
|
||||
height: height,
|
||||
pending: pending,
|
||||
}
|
||||
}
|
||||
|
||||
func randomHeight() int {
|
||||
height := 1
|
||||
for height < maxLevel && rand.Float64() < probability {
|
||||
height++
|
||||
}
|
||||
return height
|
||||
}
|
||||
|
||||
func cloneBytes(src []byte) []byte {
|
||||
if src == nil {
|
||||
return nil
|
||||
}
|
||||
dst := make([]byte, len(src))
|
||||
copy(dst, src)
|
||||
return dst
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
package memtable
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSkipListEmpty(t *testing.T) {
|
||||
sl := NewSkipList(NewArena(1024))
|
||||
|
||||
found, value, sequence := sl.Get([]byte("missing"))
|
||||
require.False(t, found)
|
||||
require.Nil(t, value)
|
||||
require.Zero(t, sequence)
|
||||
}
|
||||
|
||||
func TestSkipListOrdered(t *testing.T) {
|
||||
sl := NewSkipList(NewArena(64 << 20))
|
||||
|
||||
keys := []string{"c", "a", "e", "b", "d"}
|
||||
for i, key := range keys {
|
||||
require.NoError(t, sl.Put([]byte(key), []byte("value-"+key), uint64(i+1), false))
|
||||
}
|
||||
|
||||
for i, key := range keys {
|
||||
found, value, sequence := sl.Get([]byte(key))
|
||||
require.True(t, found, "key %q", key)
|
||||
require.Equal(t, []byte("value-"+key), value)
|
||||
require.Equal(t, uint64(i+1), sequence)
|
||||
}
|
||||
|
||||
it := sl.NewIterator()
|
||||
var got []string
|
||||
for it.Valid() {
|
||||
got = append(got, string(it.Key()))
|
||||
it.Next()
|
||||
}
|
||||
require.Equal(t, []string{"a", "b", "c", "d", "e"}, got)
|
||||
}
|
||||
|
||||
func TestSkipListOverwrite(t *testing.T) {
|
||||
sl := NewSkipList(NewArena(1024))
|
||||
|
||||
require.NoError(t, sl.Put([]byte("key"), []byte("old"), 1, false))
|
||||
require.NoError(t, sl.Put([]byte("key"), []byte("new"), 2, false))
|
||||
|
||||
found, value, sequence := sl.Get([]byte("key"))
|
||||
require.True(t, found)
|
||||
require.Equal(t, []byte("new"), value)
|
||||
require.Equal(t, uint64(2), sequence)
|
||||
|
||||
it := sl.NewIterator()
|
||||
require.True(t, it.Valid())
|
||||
require.Equal(t, []byte("key"), it.Key())
|
||||
require.Equal(t, []byte("new"), it.Value())
|
||||
it.Next()
|
||||
require.False(t, it.Valid())
|
||||
}
|
||||
|
||||
func TestSkipListPendingEntriesAreHidden(t *testing.T) {
|
||||
sl := NewSkipList(NewArena(1024))
|
||||
|
||||
require.NoError(t, sl.Put([]byte("a"), []byte("pending"), 1, true))
|
||||
require.NoError(t, sl.Put([]byte("b"), []byte("published"), 2, false))
|
||||
|
||||
found, _, _ := sl.Get([]byte("a"))
|
||||
require.False(t, found)
|
||||
|
||||
it := sl.NewIterator()
|
||||
require.True(t, it.Valid())
|
||||
require.Equal(t, []byte("b"), it.Key())
|
||||
it.Next()
|
||||
require.False(t, it.Valid())
|
||||
}
|
||||
|
||||
func TestSkipListConcurrent(t *testing.T) {
|
||||
sl := NewSkipList(NewArena(64 << 20))
|
||||
|
||||
const (
|
||||
writers = 4
|
||||
readers = 4
|
||||
keysPerWriter = 100
|
||||
readsPerReader = 1000
|
||||
)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
errCh := make(chan error, writers+readers)
|
||||
|
||||
for writer := range writers {
|
||||
wg.Go(func() {
|
||||
<-start
|
||||
for i := range keysPerWriter {
|
||||
key := fmt.Appendf(nil, "writer-%d-key-%03d", writer, i)
|
||||
value := fmt.Appendf(nil, "value-%d-%03d", writer, i)
|
||||
sequence := uint64(writer*keysPerWriter + i + 1)
|
||||
if err := sl.Put(key, value, sequence, false); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
for reader := range readers {
|
||||
wg.Go(func() {
|
||||
<-start
|
||||
rng := rand.New(rand.NewSource(int64(reader)))
|
||||
for range readsPerReader {
|
||||
writer := rng.Intn(writers)
|
||||
keyID := rng.Intn(keysPerWriter)
|
||||
key := fmt.Appendf(nil, "writer-%d-key-%03d", writer, keyID)
|
||||
found, value, sequence := sl.Get(key)
|
||||
if found {
|
||||
expectedValue := fmt.Appendf(nil, "value-%d-%03d", writer, keyID)
|
||||
expectedSequence := uint64(writer*keysPerWriter + keyID + 1)
|
||||
if err := compareEntry(value, sequence, expectedValue, expectedSequence); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
for writer := range writers {
|
||||
for i := range keysPerWriter {
|
||||
key := fmt.Appendf(nil, "writer-%d-key-%03d", writer, i)
|
||||
found, value, sequence := sl.Get(key)
|
||||
require.True(t, found, "key %q", key)
|
||||
require.Equal(t, fmt.Appendf(nil, "value-%d-%03d", writer, i), value)
|
||||
require.Equal(t, uint64(writer*keysPerWriter+i+1), sequence)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func compareEntry(value []byte, sequence uint64, expectedValue []byte, expectedSequence uint64) error {
|
||||
if string(value) != string(expectedValue) {
|
||||
return fmt.Errorf("value = %q, want %q", value, expectedValue)
|
||||
}
|
||||
if sequence != expectedSequence {
|
||||
return fmt.Errorf("sequence = %d, want %d", sequence, expectedSequence)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user