perf: updated skiplist find algorithm to be more efficient and cache aware

pkg/memtable/bench_test.go

@@ -90,6 +90,190 @@ func BenchmarkMemTableGet(b *testing.B) {
 	}
 }
 
+func BenchmarkMemTableDelete(b *testing.B) {
+	mt := NewMemTable()
+
+	// Prepare keys ahead of time
+	keys := make([][]byte, b.N)
+	for i := 0; i < b.N; i++ {
+		keys[i] = []byte(fmt.Sprintf("key-%d", i))
+	}
+
+	// Insert entries first
+	for i := 0; i < b.N; i++ {
+		value := []byte(fmt.Sprintf("value-%d", i))
+		mt.Put(keys[i], value, uint64(i))
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		mt.Delete(keys[i], uint64(i+b.N))
+	}
+}
+
+func BenchmarkImmutableMemTableGet(b *testing.B) {
+	mt := NewMemTable()
+
+	// Insert entries first
+	const numEntries = 100000
+	keys := make([][]byte, numEntries)
+	for i := 0; i < numEntries; i++ {
+		key := []byte(fmt.Sprintf("key-%d", i))
+		value := []byte(fmt.Sprintf("value-%d", i))
+		keys[i] = key
+		mt.Put(key, value, uint64(i))
+	}
+
+	// Mark memtable as immutable
+	mt.SetImmutable()
+
+	// Create random keys for lookup
+	lookupKeys := make([][]byte, b.N)
+	r := rand.New(rand.NewSource(42)) // Use fixed seed for reproducibility
+	for i := 0; i < b.N; i++ {
+		idx := r.Intn(numEntries)
+		lookupKeys[i] = keys[idx]
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		mt.Get(lookupKeys[i])
+	}
+}
+
+func BenchmarkConcurrentMemTableGet(b *testing.B) {
+	// This benchmark tests concurrent read performance on a mutable memtable
+	mt := NewMemTable()
+	
+	// Insert entries first
+	const numEntries = 100000
+	keys := make([][]byte, numEntries)
+	for i := 0; i < numEntries; i++ {
+		key := []byte(fmt.Sprintf("key-%d", i))
+		value := []byte(fmt.Sprintf("value-%d", i))
+		keys[i] = key
+		mt.Put(key, value, uint64(i))
+	}
+	
+	// Create random keys for lookup
+	r := rand.New(rand.NewSource(42)) // Use fixed seed for reproducibility
+	lookupKeys := make([][]byte, b.N)
+	for i := 0; i < b.N; i++ {
+		idx := r.Intn(numEntries)
+		lookupKeys[i] = keys[idx]
+	}
+	
+	// Set up for parallel benchmark
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		// Each goroutine needs its own random sequence
+		localRand := rand.New(rand.NewSource(rand.Int63()))
+		i := 0
+		for pb.Next() {
+			// Pick a random key from our prepared list
+			idx := localRand.Intn(len(lookupKeys))
+			mt.Get(lookupKeys[idx])
+			i++
+		}
+	})
+}
+
+func BenchmarkConcurrentImmutableMemTableGet(b *testing.B) {
+	// This benchmark tests concurrent read performance on an immutable memtable
+	mt := NewMemTable()
+	
+	// Insert entries first
+	const numEntries = 100000
+	keys := make([][]byte, numEntries)
+	for i := 0; i < numEntries; i++ {
+		key := []byte(fmt.Sprintf("key-%d", i))
+		value := []byte(fmt.Sprintf("value-%d", i))
+		keys[i] = key
+		mt.Put(key, value, uint64(i))
+	}
+	
+	// Mark memtable as immutable
+	mt.SetImmutable()
+	
+	// Create random keys for lookup
+	r := rand.New(rand.NewSource(42)) // Use fixed seed for reproducibility
+	lookupKeys := make([][]byte, b.N)
+	for i := 0; i < b.N; i++ {
+		idx := r.Intn(numEntries)
+		lookupKeys[i] = keys[idx]
+	}
+	
+	// Set up for parallel benchmark
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		// Each goroutine needs its own random sequence
+		localRand := rand.New(rand.NewSource(rand.Int63()))
+		i := 0
+		for pb.Next() {
+			// Pick a random key from our prepared list
+			idx := localRand.Intn(len(lookupKeys))
+			mt.Get(lookupKeys[idx])
+			i++
+		}
+	})
+}
+
+func BenchmarkMixedWorkload(b *testing.B) {
+	// Skip very long benchmarks if testing with -short flag
+	if testing.Short() {
+		b.Skip("Skipping mixed workload benchmark in short mode")
+	}
+	
+	// This benchmark tests a mixed workload with concurrent reads and writes
+	mt := NewMemTable()
+	
+	// Pre-populate with some data
+	const initialEntries = 50000
+	keys := make([][]byte, initialEntries)
+	for i := 0; i < initialEntries; i++ {
+		key := []byte(fmt.Sprintf("key-%d", i))
+		value := []byte(fmt.Sprintf("value-%d", i))
+		keys[i] = key
+		mt.Put(key, value, uint64(i))
+	}
+	
+	// Prepare random operations
+	readRatio := 0.8 // 80% reads, 20% writes
+	
+	b.ResetTimer()
+	
+	// Run the benchmark in parallel mode
+	b.RunParallel(func(pb *testing.PB) {
+		// Each goroutine gets its own random number generator
+		r := rand.New(rand.NewSource(rand.Int63()))
+		localCount := 0
+		
+		// Continue until the benchmark is done
+		for pb.Next() {
+			// Determine operation: read or write
+			op := r.Float64()
+			if op < readRatio {
+				// Read operation
+				idx := r.Intn(initialEntries)
+				mt.Get(keys[idx])
+			} else {
+				// Write operation (alternating put and delete)
+				if localCount%2 == 0 {
+					// Put
+					newKey := []byte(fmt.Sprintf("key-new-%d", localCount))
+					newValue := []byte(fmt.Sprintf("value-new-%d", localCount))
+					mt.Put(newKey, newValue, uint64(initialEntries+localCount))
+				} else {
+					// Delete (use an existing key)
+					idx := r.Intn(initialEntries)
+					mt.Delete(keys[idx], uint64(initialEntries+localCount))
+				}
+			}
+			localCount++
+		}
+	})
+}
+
 func BenchmarkMemPoolGet(b *testing.B) {
 	cfg := createTestConfig()
 	cfg.MemTableSize = 1024 * 1024 * 32 // 32MB for benchmark

pkg/memtable/memtable.go

@@ -32,7 +32,7 @@ func (m *MemTable) Put(key, value []byte, seqNum uint64) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if m.immutable.Load() {
+	if m.IsImmutable() {
 		// Don't modify immutable memtables
 		return
 	}
@@ -51,7 +51,7 @@ func (m *MemTable) Delete(key []byte, seqNum uint64) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if m.immutable.Load() {
+	if m.IsImmutable() {
 		// Don't modify immutable memtables
 		return
 	}
@@ -70,28 +70,52 @@ func (m *MemTable) Delete(key []byte, seqNum uint64) {
 // Returns (nil, false) if the key does not exist
 // Returns (value, true) if the key exists and has a value
 func (m *MemTable) Get(key []byte) ([]byte, bool) {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
+	// Use atomic check for immutability first
+	if m.IsImmutable() {
+		// For immutable memtables, we can bypass the write lock completely
+		e := m.skipList.Find(key)
+		if e == nil {
+			return nil, false
+		}
 
-	e := m.skipList.Find(key)
-	if e == nil {
-		return nil, false
+		// Check if this is a deletion marker
+		if e.valueType == TypeDeletion {
+			return nil, true // Key exists but was deleted
+		}
+
+		return e.value, true
+	} else {
+		// For mutable memtables, we still need read lock protection
+		// as the structure could be modified during reads
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+
+		e := m.skipList.Find(key)
+		if e == nil {
+			return nil, false
+		}
+
+		// Check if this is a deletion marker
+		if e.valueType == TypeDeletion {
+			return nil, true // Key exists but was deleted
+		}
+
+		return e.value, true
 	}
-
-	// Check if this is a deletion marker
-	if e.valueType == TypeDeletion {
-		return nil, true // Key exists but was deleted
-	}
-
-	return e.value, true
 }
 
 // Contains checks if the key exists in the MemTable
 func (m *MemTable) Contains(key []byte) bool {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-
-	return m.skipList.Find(key) != nil
+	// For immutable memtables, we can bypass the RWLock completely
+	if m.IsImmutable() {
+		return m.skipList.Find(key) != nil
+	} else {
+		// For mutable memtables, we still need read lock protection
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		
+		return m.skipList.Find(key) != nil
+	}
 }
 
 // ApproximateSize returns the approximate size of the MemTable in bytes
@@ -117,14 +141,29 @@ func (m *MemTable) Age() float64 {
 
 // NewIterator returns an iterator for the MemTable
 func (m *MemTable) NewIterator() *Iterator {
-	return m.skipList.NewIterator()
+	// For immutable memtables, we can bypass the lock
+	if m.IsImmutable() {
+		return m.skipList.NewIterator()
+	} else {
+		// For mutable memtables, we need read lock to ensure stability during iteration
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		
+		return m.skipList.NewIterator()
+	}
 }
 
 // GetNextSequenceNumber returns the next sequence number to use
 func (m *MemTable) GetNextSequenceNumber() uint64 {
-	m.mu.RLock()
-	defer m.mu.RUnlock()
-	return m.nextSeqNum
+	// For immutable memtables, nextSeqNum won't change
+	if m.IsImmutable() {
+		return m.nextSeqNum
+	} else {
+		// For mutable memtables, we need read lock
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		return m.nextSeqNum
+	}
 }
 
 // ProcessWALEntry processes a WAL entry and applies it to the MemTable

pkg/memtable/skiplist.go

@@ -14,7 +14,7 @@ const (
 	MaxHeight = 12
 
 	// BranchingFactor determines the probability of increasing the height
-	BranchingFactor = 4
+	BranchingFactor = 2
 
 	// DefaultCacheLineSize aligns nodes to cache lines for better performance
 	DefaultCacheLineSize = 64
@@ -124,12 +124,15 @@ func NewSkipList() *SkipList {
 }
 
 // randomHeight generates a random height for a new node
+// Uses a geometric distribution with p=0.5 for better balanced trees
 func (s *SkipList) randomHeight() int {
 	s.rndMtx.Lock()
 	defer s.rndMtx.Unlock()
 
+	// Use a geometric distribution with p=0.5
+	// Each level has 50% chance of promotion to next level
 	height := 1
-	for height < MaxHeight && s.rnd.Intn(BranchingFactor) == 0 {
+	for height < MaxHeight && s.rnd.Int31n(BranchingFactor) == 0 {
 		height++
 	}
 	return height
@@ -155,22 +158,24 @@ func (s *SkipList) Insert(e *entry) {
 		}
 	}
 
-	// Find where to insert at each level
+	// Find where to insert at each level - using the efficient search algorithm
 	current := s.head
 	for level := currHeight - 1; level >= 0; level-- {
-		// Find the insertion point at this level
-		for next := current.getNext(level); next != nil; next = current.getNext(level) {
-			if next.entry.compareWithEntry(e) >= 0 {
-				break
-			}
+		// Move right until we find a node >= the key we're inserting
+		next := current.getNext(level)
+		for next != nil && next.entry.compareWithEntry(e) < 0 {
 			current = next
+			next = current.getNext(level)
 		}
+		// Remember the node before the insertion point at this level
 		prev[level] = current
 	}
 
-	// Insert the node at each level
+	// Insert the node at each level - from bottom up
 	for level := 0; level < height; level++ {
+		// Link the new node's next pointer to the next node at this level
 		node.setNext(level, prev[level].getNext(level))
+		// Link the previous node's next pointer to the new node
 		prev[level].setNext(level, node)
 	}
 
@@ -181,46 +186,37 @@ func (s *SkipList) Insert(e *entry) {
 // Find looks for an entry with the specified key
 // If multiple entries have the same key, the most recent one is returned
 func (s *SkipList) Find(key []byte) *entry {
-	var result *entry
 	current := s.head
 	height := s.getCurrentHeight()
 
-	// Start from the highest level for efficient search
+	// Start at the highest level, and work our way down
+	// At each level, move right as far as possible without overshooting
 	for level := height - 1; level >= 0; level-- {
-		// Scan forward until we find a key greater than or equal to the target
-		for next := current.getNext(level); next != nil; next = current.getNext(level) {
-			cmp := next.entry.compare(key)
-			if cmp > 0 {
-				// Key at next is greater than target, go down a level
-				break
-			} else if cmp == 0 {
-				// Found a match, check if it's newer than our current result
-				if result == nil || next.entry.seqNum > result.seqNum {
-					result = next.entry
-				}
-				// Continue at this level to see if there are more entries with same key
-				current = next
-			} else {
-				// Key at next is less than target, move forward
-				current = next
-			}
+		next := current.getNext(level)
+		for next != nil && next.entry.compare(key) < 0 {
+			current = next
+			next = current.getNext(level)
 		}
+		// When we exit this loop, current.next[level] is either nil or >= key
 	}
 
-	// For level 0, do one more sweep to ensure we get the newest entry
-	current = s.head
-	for next := current.getNext(0); next != nil; next = next.getNext(0) {
-		cmp := next.entry.compare(key)
-		if cmp > 0 {
-			// Past the key
-			break
-		} else if cmp == 0 {
-			// Found a match, update result if it's newer
-			if result == nil || next.entry.seqNum > result.seqNum {
-				result = next.entry
-			}
+	// We're now at level 0 with current just before the potential target
+	// Check next node to see if it's our target key
+	candidate := current.getNext(0)
+	if candidate == nil || candidate.entry.compare(key) != 0 {
+		// Key doesn't exist in the list
+		return nil
+	}
+
+	// We found the key, but there might be multiple entries with the same key
+	// Find the one with the highest sequence number (most recent)
+	var result *entry = candidate.entry
+	for candidate != nil && candidate.entry.compare(key) == 0 {
+		// Update result if this entry has a higher sequence number
+		if candidate.entry.seqNum > result.seqNum {
+			result = candidate.entry
 		}
-		current = next
+		candidate = candidate.getNext(0)
 	}
 
 	return result
@@ -269,17 +265,18 @@ func (it *Iterator) Seek(key []byte) {
 	current := it.list.head
 	height := it.list.getCurrentHeight()
 
-	// Search algorithm similar to Find
+	// Start at the highest level, and work our way down
+	// At each level, move right as far as possible without overshooting
 	for level := height - 1; level >= 0; level-- {
-		for next := current.getNext(level); next != nil; next = current.getNext(level) {
-			if next.entry.compare(key) >= 0 {
-				break
-			}
+		next := current.getNext(level)
+		for next != nil && next.entry.compare(key) < 0 {
 			current = next
+			next = current.getNext(level)
 		}
+		// When we exit this loop, current.next[level] is either nil or >= key
 	}
 
-	// Move to the next node, which should be >= target
+	// Move to the next node at level 0, which should be >= target
 	it.current = current.getNext(0)
 }