flatkv cache

sei-db/common/utils/chan_utils.go

@@ -0,0 +1,32 @@
+package utils
+
+import (
+	"context"
+	"fmt"
+)
+
+// TODO unit test before merge
+
+// Push to a channel, returning an error if the context is cancelled before the value is pushed.
+func InterruptiblePush[T any](ctx context.Context, ch chan T, value T) error {
+	select {
+	case <-ctx.Done():
+		return fmt.Errorf("context cancelled: %w", ctx.Err())
+	case ch <- value:
+		return nil
+	}
+}
+
+// Pull from a channel, returning an error if the context is cancelled before the value is pulled.
+func InterruptiblePull[T any](ctx context.Context, ch <-chan T) (T, error) {
+	var zero T
+	select {
+	case <-ctx.Done():
+		return zero, fmt.Errorf("context cancelled: %w", ctx.Err())
+	case value, ok := <-ch:
+		if !ok {
+			return zero, fmt.Errorf("channel closed")
+		}
+		return value, nil
+	}
+}

sei-db/state_db/sc/flatkv/flatcache/cache.go

@@ -0,0 +1,21 @@
+package flatcache
+
+import "github.com/sei-protocol/sei-chain/sei-db/proto"
+
+// Cache describes a cache kapable of being used by a FlatKV store.
+type Cache interface {
+
+	// TODO decide if we should support individual modifications
+
+	// Get returns the value for the given key, or (nil, false) if not found.
+	Get(key []byte) ([]byte, bool, error)
+
+	// Set sets the value for the given key.
+	Set(key []byte, value []byte)
+
+	// Delete deletes the value for the given key.
+	Delete(key []byte)
+
+	// BatchSet applies the given changesets to the cache.
+	BatchSet(cs []*proto.NamedChangeSet)
+}

sei-db/state_db/sc/flatkv/flatcache/cache_impl.go

@@ -0,0 +1,168 @@
+package flatcache
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/sei-protocol/sei-chain/sei-db/proto"
+	iavl "github.com/sei-protocol/sei-chain/sei-iavl/proto"
+)
+
+var _ Cache = (*cache)(nil)
+
+// A standard implementation of a flatcache.
+type cache struct {
+	ctx context.Context
+
+	// A utility for assigning keys to shard indices.
+	shardManager *shardManager
+
+	// The shards in the cache.
+	shards []*shard
+
+	// A scheduler for asyncronous reads.
+	readScheduler *readScheduler
+
+	// The interval at which to run garbage collection.
+	garbageCollectionInterval time.Duration
+}
+
+// Creates a new Cache.
+func NewCache(
+	ctx context.Context,
+	// A function that reads a value from the database.
+	readFunc func(key []byte) []byte,
+	// The number of shards in the cache. Must be a power of two and greater than 0.
+	shardCount int,
+	// The maximum size of the cache, in bytes.
+	maxSize int,
+	// The number of background goroutines to read values from the database.
+	readWorkerCount int,
+	// The max size of the read queue.
+	readQueueSize int,
+	// The interval at which to run garbage collection.
+	garbageCollectionInterval time.Duration,
+) (Cache, error) {
+	if shardCount <= 0 || (shardCount&(shardCount-1)) != 0 {
+		return nil, ErrNumShardsNotPowerOfTwo
+	}
+	if maxSize <= 0 {
+		return nil, fmt.Errorf("maxSize must be greater than 0")
+	}
+	if readWorkerCount <= 0 {
+		return nil, fmt.Errorf("readWorkerCount must be greater than 0")
+	}
+	if readQueueSize <= 0 {
+		return nil, fmt.Errorf("readQueueSize must be greater than 0")
+	}
+
+	shardManager, err := NewShardManager(uint64(shardCount))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create shard manager: %w", err)
+	}
+	if garbageCollectionInterval <= 0 {
+		return nil, fmt.Errorf("garbageCollectionInterval must be greater than 0")
+	}
+
+	readScheduler := NewReadScheduler(ctx, readFunc, readWorkerCount, readQueueSize)
+
+	sizePerShard := maxSize / shardCount
+	if sizePerShard <= 0 {
+		return nil, fmt.Errorf("maxSize must be greater than shardCount")
+	}
+
+	shards := make([]*shard, shardCount)
+	for i := 0; i < shardCount; i++ {
+		shards[i], err = NewShard(ctx, readScheduler, sizePerShard)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create shard: %w", err)
+		}
+	}
+
+	c := &cache{
+		ctx:                       ctx,
+		shardManager:              shardManager,
+		shards:                    shards,
+		readScheduler:             readScheduler,
+		garbageCollectionInterval: garbageCollectionInterval,
+	}
+
+	go c.runGarbageCollection()
+
+	return c, nil
+}
+
+func (c *cache) BatchSet(cs []*proto.NamedChangeSet) {
+
+	// First, sort entries by shard index.
+	// This allows us to set all values in a single shard with only a single lock acquisition.
+	shardMap := make(map[uint64][]*iavl.KVPair)
+	for _, ncs := range cs {
+		for _, entry := range ncs.Changeset.Pairs {
+			shardMap[c.shardManager.Shard(entry.Key)] = append(shardMap[c.shardManager.Shard(entry.Key)], entry)
+		}
+	}
+
+	// This is probably quite fast, but if it isn't it can be parallelized.
+	for shardIndex, shardEntries := range shardMap {
+		shard := c.shards[shardIndex]
+		shard.BatchSet(shardEntries)
+	}
+}
+
+func (c *cache) Delete(key []byte) {
+	shardIndex := c.shardManager.Shard(key)
+	shard := c.shards[shardIndex]
+	shard.Delete(key)
+}
+
+func (c *cache) Get(key []byte) ([]byte, bool, error) {
+	shardIndex := c.shardManager.Shard(key)
+	shard := c.shards[shardIndex]
+
+	value, ok, err := shard.Get(key)
+	if err != nil {
+		return nil, false, fmt.Errorf("failed to get value from shard: %w", err)
+	}
+	if !ok {
+		return nil, false, nil
+	}
+	return value, ok, nil
+}
+
+func (c *cache) Set(key []byte, value []byte) {
+	shardIndex := c.shardManager.Shard(key)
+	shard := c.shards[shardIndex]
+	shard.Set(key, value)
+}
+
+// TODO add GC metrics
+
+// Periodically runs garbage collection in the background.
+func (c *cache) runGarbageCollection() {
+
+	// Spread out work evenly across all shards, so that we visit each shard roughly once per interval.
+	gcSubInterval := c.garbageCollectionInterval / time.Duration(len(c.shards))
+	if gcSubInterval == 0 {
+		// technically possible if the number of shards is very large and the interval is very small
+		gcSubInterval = 1
+	}
+	ticker := time.NewTicker(gcSubInterval)
+	defer ticker.Stop()
+
+	nextShardIndex := 0
+
+	for {
+		select {
+		case <-c.ctx.Done():
+			return
+		case <-ticker.C:
+			shardIndex := nextShardIndex
+			nextShardIndex = (nextShardIndex + 1) % len(c.shards)
+			c.shards[shardIndex].RunGarbageCollection()
+		}
+	}
+}
+
+// TODO create a warming mechanism

sei-db/state_db/sc/flatkv/flatcache/lru_queue.go

@@ -0,0 +1,83 @@
+package flatcache
+
+import "container/list"
+
+// Implements a queue-like abstraction with LRU semantics. Not thread safe.
+type lruQueue struct {
+	order     *list.List
+	entries   map[string]*list.Element
+	totalSize int
+}
+
+type lruQueueEntry struct {
+	key  []byte
+	size int
+}
+
+// Create a new LRU queue.
+func NewLRUQueue() *lruQueue {
+	return &lruQueue{
+		order:   list.New(),
+		entries: make(map[string]*list.Element),
+	}
+}
+
+// Add a new entry to the LRU queue. Can also be used to update an existing value with a new weight.
+func (lru *lruQueue) Push(
+	// the key in the cache that was recently interacted with
+	key []byte,
+	// the size of the key + value
+	size int,
+) {
+	keyString := string(key) // TODO revisit and maybe do unsafe copies
+	if elem, ok := lru.entries[keyString]; ok {
+		entry := elem.Value.(*lruQueueEntry)
+		lru.totalSize += size - entry.size
+		entry.size = size
+		lru.order.MoveToBack(elem)
+		return
+	}
+
+	keyCopy := append([]byte(nil), key...) // TODO don't do this
+	elem := lru.order.PushBack(&lruQueueEntry{
+		key:  keyCopy,
+		size: size,
+	})
+	lru.entries[keyString] = elem
+	lru.totalSize += size
+}
+
+// Signal that an entry has been interated with, moving it to the to the back of the queue
+// (i.e. making it so it doesn't get popped soon).
+func (lru *lruQueue) Touch(key []byte) {
+	elem, ok := lru.entries[string(key)]
+	if !ok {
+		return
+	}
+	lru.order.MoveToBack(elem)
+}
+
+// Returns the total size of all entries in the LRU queue.
+func (lru *lruQueue) GetTotalSize() int {
+	return lru.totalSize
+}
+
+// Returns a count of the number of entries in the LRU queue, where each entry counts for 1 regardless of size.
+func (lru *lruQueue) GetCount() int {
+	return len(lru.entries)
+}
+
+// Pops a single element out of the queue. The element removed is the entry least recently passed to Update().
+// Panics if the queue is empty.
+func (lru *lruQueue) PopLeastRecentlyUsed() []byte {
+	elem := lru.order.Front()
+	if elem == nil {
+		panic("cannot pop from empty LRU queue")
+	}
+
+	lru.order.Remove(elem)
+	entry := elem.Value.(*lruQueueEntry)
+	delete(lru.entries, string(entry.key))
+	lru.totalSize -= entry.size
+	return entry.key
+}

sei-db/state_db/sc/flatkv/flatcache/lru_queue_test.go

@@ -0,0 +1,82 @@
+package flatcache
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestLRUQueueTracksSizeCountAndOrder(t *testing.T) {
+	lru := NewLRUQueue()
+
+	lru.Push([]byte("a"), 3)
+	lru.Push([]byte("b"), 5)
+	lru.Push([]byte("c"), 7)
+
+	if got := lru.GetCount(); got != 3 {
+		t.Fatalf("GetCount() = %d, want 3", got)
+	}
+
+	if got := lru.GetTotalSize(); got != 15 {
+		t.Fatalf("GetTotalSize() = %d, want 15", got)
+	}
+
+	lru.Touch([]byte("a"))
+
+	if got := lru.PopLeastRecentlyUsed(); !bytes.Equal(got, []byte("b")) {
+		t.Fatalf("first pop = %q, want %q", got, []byte("b"))
+	}
+
+	if got := lru.PopLeastRecentlyUsed(); !bytes.Equal(got, []byte("c")) {
+		t.Fatalf("second pop = %q, want %q", got, []byte("c"))
+	}
+
+	if got := lru.PopLeastRecentlyUsed(); !bytes.Equal(got, []byte("a")) {
+		t.Fatalf("third pop = %q, want %q", got, []byte("a"))
+	}
+
+	if got := lru.GetCount(); got != 0 {
+		t.Fatalf("GetCount() after pops = %d, want 0", got)
+	}
+
+	if got := lru.GetTotalSize(); got != 0 {
+		t.Fatalf("GetTotalSize() after pops = %d, want 0", got)
+	}
+}
+
+func TestLRUQueuePushUpdatesExistingEntry(t *testing.T) {
+	lru := NewLRUQueue()
+
+	lru.Push([]byte("a"), 3)
+	lru.Push([]byte("b"), 5)
+	lru.Push([]byte("a"), 11)
+
+	if got := lru.GetCount(); got != 2 {
+		t.Fatalf("GetCount() = %d, want 2", got)
+	}
+
+	if got := lru.GetTotalSize(); got != 16 {
+		t.Fatalf("GetTotalSize() = %d, want 16", got)
+	}
+
+	if got := lru.PopLeastRecentlyUsed(); !bytes.Equal(got, []byte("b")) {
+		t.Fatalf("first pop = %q, want %q", got, []byte("b"))
+	}
+
+	if got := lru.PopLeastRecentlyUsed(); !bytes.Equal(got, []byte("a")) {
+		t.Fatalf("second pop = %q, want %q", got, []byte("a"))
+	}
+}
+
+func TestLRUQueueCopiesInsertedKey(t *testing.T) {
+	lru := NewLRUQueue()
+
+	key := []byte("a")
+	lru.Push(key, 1)
+	key[0] = 'z'
+
+	if got := lru.PopLeastRecentlyUsed(); !bytes.Equal(got, []byte("a")) {
+		t.Fatalf("pop after mutating caller key = %q, want %q", got, []byte("a"))
+	}
+}
+
+// TODO expand these tests