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/db_engine/pebbledb/batch.go

@@ -2,38 +2,55 @@ package pebbledb
 
 import (
 	"github.com/cockroachdb/pebble/v2"
+	"github.com/sei-protocol/sei-chain/sei-db/db_engine/pebbledb/flatcache"
 	"github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
 )
 
 // pebbleBatch wraps a Pebble batch for atomic writes.
 // Important: Callers must call Close() after Commit() to release batch resources,
 // even if Commit() succeeds. Failure to Close() will leak memory.
 type pebbleBatch struct {
-	b *pebble.Batch
+	b     *pebble.Batch
+	cache flatcache.Cache
+
+	// Writes are tracked so the cache can be updated after a successful commit.
+	pendingCacheUpdates []flatcache.CacheUpdate
 }
 
 var _ types.Batch = (*pebbleBatch)(nil)
 
-func newPebbleBatch(db *pebble.DB) *pebbleBatch {
-	return &pebbleBatch{b: db.NewBatch()}
+func newPebbleBatch(db *pebble.DB, cache flatcache.Cache) *pebbleBatch {
+	return &pebbleBatch{b: db.NewBatch(), cache: cache}
 }
 
 func (p *pebbleDB) NewBatch() types.Batch {
-	return newPebbleBatch(p.db)
+	return newPebbleBatch(p.db, p.cache)
 }
 
 func (pb *pebbleBatch) Set(key, value []byte) error {
-	// Durability options are applied on Commit.
+	pb.pendingCacheUpdates = append(pb.pendingCacheUpdates, flatcache.CacheUpdate{
+		Key:   key,
+		Value: value,
+	})
 	return pb.b.Set(key, value, nil)
 }
 
 func (pb *pebbleBatch) Delete(key []byte) error {
-	// Durability options are applied on Commit.
+	pb.pendingCacheUpdates = append(pb.pendingCacheUpdates, flatcache.CacheUpdate{
+		Key:      key,
+		IsDelete: true,
+	})
 	return pb.b.Delete(key, nil)
 }
 
 func (pb *pebbleBatch) Commit(opts types.WriteOptions) error {
-	return pb.b.Commit(toPebbleWriteOpts(opts))
+	err := pb.b.Commit(toPebbleWriteOpts(opts))
+	if err != nil {
+		return err
+	}
+	pb.cache.BatchSet(pb.pendingCacheUpdates)
+	pb.pendingCacheUpdates = nil
+	return nil
 }
 
 func (pb *pebbleBatch) Len() int {
@@ -42,6 +59,7 @@ func (pb *pebbleBatch) Len() int {
 
 func (pb *pebbleBatch) Reset() {
 	pb.b.Reset()
+	pb.pendingCacheUpdates = nil
 }
 
 func (pb *pebbleBatch) Close() error {

sei-db/db_engine/pebbledb/db.go

@@ -14,6 +14,7 @@ import (
 
 	errorutils "github.com/sei-protocol/sei-chain/sei-db/common/errors"
 	"github.com/sei-protocol/sei-chain/sei-db/common/metrics"
+	"github.com/sei-protocol/sei-chain/sei-db/db_engine/pebbledb/flatcache"
 	"github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
 )
 
@@ -23,6 +24,7 @@ const metricsScrapeInterval = 10 * time.Second
 type pebbleDB struct {
 	db            *pebble.DB
 	metricsCancel context.CancelFunc
+	cache         flatcache.Cache
 }
 
 var _ types.KeyValueDB = (*pebbleDB)(nil)
@@ -44,11 +46,12 @@ func Open(
 		}
 	}
 
-	cache := pebble.NewCache(1024 * 1024 * 512) // 512MB cache
-	defer cache.Unref()
+	// Internal pebbleDB cache, used to cache pages in memory. // TODO verify accuracy of this statement
+	pebbleCache := pebble.NewCache(1024 * 1024 * 512) // 512MB cache
+	defer pebbleCache.Unref()
 
 	popts := &pebble.Options{
-		Cache:    cache,
+		Cache:    pebbleCache,
 		Comparer: cmp,
 		// FormatMajorVersion is pinned to a specific version to prevent accidental
 		// breaking changes when updating the pebble dependency. Using FormatNewest
@@ -92,33 +95,77 @@ func Open(
 		return nil, err
 	}
 
+	readFunction := func(key []byte) []byte { // TODO error handling!
+		val, closer, err := db.Get(key)
+		if err != nil {
+			return nil
+		}
+		cloned := bytes.Clone(val)
+		_ = closer.Close()
+		return cloned
+	}
+
+	// A high level cache per key.
+	cache, err := flatcache.NewCache(
+		ctx,
+		readFunction,
+		8,
+		1024*1024*1024,
+		20,
+		64,
+		10*time.Second)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create flatcache: %w", err)
+	}
+
 	ctx, cancel := context.WithCancel(ctx)
 	if enableMetrics {
 		metrics.NewPebbleMetrics(ctx, db, filepath.Base(path), metricsScrapeInterval)
 	}
 
-	return &pebbleDB{db: db, metricsCancel: cancel}, nil
+	return &pebbleDB{
+		db:            db,
+		metricsCancel: cancel,
+		cache:         cache,
+	}, nil
 }
 
 func (p *pebbleDB) Get(key []byte) ([]byte, error) {
-	// Pebble returns a zero-copy view plus a closer; we copy and close internally.
-	val, closer, err := p.db.Get(key)
+	// // Pebble returns a zero-copy view plus a closer; we copy and close internally.
+	// val, closer, err := p.db.Get(key)
+	// if err != nil {
+	// 	if errors.Is(err, pebble.ErrNotFound) {
+	// 		return nil, errorutils.ErrNotFound
+	// 	}
+	// 	return nil, err
+	// }
+	// cloned := bytes.Clone(val)
+	// _ = closer.Close()
+
+	val, found, err := p.cache.Get(key, true)
 	if err != nil {
-		if errors.Is(err, pebble.ErrNotFound) {
-			return nil, errorutils.ErrNotFound
-		}
-		return nil, err
+		return nil, fmt.Errorf("failed to get value from cache: %w", err)
+	}
+	if !found {
+		return nil, errorutils.ErrNotFound
 	}
-	cloned := bytes.Clone(val)
-	_ = closer.Close()
-	return cloned, nil
+
+	return val, nil
+}
+
+func (p *pebbleDB) BatchGet(keys map[string]types.BatchGetResult) {
+	p.cache.BatchGet(keys)
 }
 
 func (p *pebbleDB) Set(key, value []byte, opts types.WriteOptions) error {
+	// TODO batch set!
+	p.cache.Set(key, value)
 	return p.db.Set(key, value, toPebbleWriteOpts(opts))
 }
 
 func (p *pebbleDB) Delete(key []byte, opts types.WriteOptions) error {
+	// TODO batch delete!
+	p.cache.Delete(key)
 	return p.db.Delete(key, toPebbleWriteOpts(opts))
 }
 

sei-db/db_engine/pebbledb/flatcache/cache.go

@@ -0,0 +1,45 @@
+package flatcache // TODO rename the flatcache package!
+
+import "github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
+
+// CacheUpdate describes a single key-value mutation to apply to the cache.
+type CacheUpdate struct {
+	// The key to update.
+	Key []byte
+	// The value to set. If nil, the key will be deleted.
+	Value []byte
+	// If true, the key will be deleted.
+	// If false, the key will be set to the given value.
+	IsDelete bool
+}
+
+// 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(
+		// The entry to fetch.
+		key []byte,
+		// If true, the LRU queue will be updated. If false, the LRU queue will not be updated.
+		// Useful for when an operation is performed multiple times in close succession on the same key,
+		// since it requires non-zero overhead to do so with little benefit.
+		updateLru bool,
+	) ([]byte, bool, error)
+
+	// Perform a batch read operation. Given a map of keys to read, performs the reads and updates the
+	// map with the results.
+	//
+	// It is not thread safe to read or mutate the map while this method is running.
+	BatchGet(keys map[string]types.BatchGetResult)
+
+	// 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 updates to the cache.
+	BatchSet(updates []CacheUpdate)
+}