Helper files for the flatKV cache implementation

Makefile

@@ -157,6 +157,13 @@ lint:
 	go mod tidy
 	go mod verify
 
+# Run lint on the sei-db package. Much faster than running lint on the entire project.
+# Makes life easier for storage team when iterating on changes inside the sei-db package.
+dblint:
+	go run github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.8.0 run ./sei-db/...
+	go fmt ./sei-db/...
+	go vet ./sei-db/...
+
 build:
 	go build $(BUILD_FLAGS) -o ./build/seid ./cmd/seid
 

sei-db/db_engine/dbcache/cache.go

@@ -0,0 +1,47 @@
+package dbcache
+
+import (
+	"github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
+)
+
+// Cache describes a cache capable of being used by a FlatKV store.
+type Cache interface {
+
+	// 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) 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 updates to the cache.
+	BatchSet(updates []CacheUpdate) error
+}
+
+// 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
+}
+
+// IsDelete returns true if the update is a delete operation.
+func (u *CacheUpdate) IsDelete() bool {
+	return u.Value == nil
+}

sei-db/db_engine/dbcache/cached_batch.go

@@ -0,0 +1,55 @@
+package dbcache
+
+import (
+	"fmt"
+
+	"github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
+)
+
+// cachedBatch wraps a types.Batch and applies pending mutations to the cache
+// after a successful commit.
+type cachedBatch struct {
+	inner   types.Batch
+	cache   Cache
+	pending []CacheUpdate
+}
+
+var _ types.Batch = (*cachedBatch)(nil)
+
+func newCachedBatch(inner types.Batch, cache Cache) *cachedBatch {
+	return &cachedBatch{inner: inner, cache: cache}
+}
+
+func (cb *cachedBatch) Set(key, value []byte) error {
+	cb.pending = append(cb.pending, CacheUpdate{Key: key, Value: value})
+	return cb.inner.Set(key, value)
+}
+
+func (cb *cachedBatch) Delete(key []byte) error {
+	cb.pending = append(cb.pending, CacheUpdate{Key: key, Value: nil})
+	return cb.inner.Delete(key)
+}
+
+func (cb *cachedBatch) Commit(opts types.WriteOptions) error {
+	if err := cb.inner.Commit(opts); err != nil {
+		return err
+	}
+	if err := cb.cache.BatchSet(cb.pending); err != nil {
+		return fmt.Errorf("failed to update cache after commit: %w", err)
+	}
+	cb.pending = nil
+	return nil
+}
+
+func (cb *cachedBatch) Len() int {
+	return cb.inner.Len()
+}
+
+func (cb *cachedBatch) Reset() {
+	cb.inner.Reset()
+	cb.pending = nil
+}
+
+func (cb *cachedBatch) Close() error {
+	return cb.inner.Close()
+}

sei-db/db_engine/dbcache/cached_batch_test.go

@@ -0,0 +1,204 @@
+package dbcache
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
+)
+
+// ---------------------------------------------------------------------------
+// mock batch
+// ---------------------------------------------------------------------------
+
+type mockBatch struct {
+	sets       []CacheUpdate
+	deletes    [][]byte
+	committed  bool
+	closed     bool
+	resetCount int
+	commitErr  error
+}
+
+func (m *mockBatch) Set(key, value []byte) error {
+	m.sets = append(m.sets, CacheUpdate{Key: key, Value: value})
+	return nil
+}
+
+func (m *mockBatch) Delete(key []byte) error {
+	m.deletes = append(m.deletes, key)
+	return nil
+}
+
+func (m *mockBatch) Commit(opts types.WriteOptions) error {
+	if m.commitErr != nil {
+		return m.commitErr
+	}
+	m.committed = true
+	return nil
+}
+
+func (m *mockBatch) Len() int {
+	return len(m.sets) + len(m.deletes)
+}
+
+func (m *mockBatch) Reset() {
+	m.sets = nil
+	m.deletes = nil
+	m.committed = false
+	m.resetCount++
+}
+
+func (m *mockBatch) Close() error {
+	m.closed = true
+	return nil
+}
+
+// ---------------------------------------------------------------------------
+// mock cache
+// ---------------------------------------------------------------------------
+
+type mockCache struct {
+	data        map[string][]byte
+	batchSetErr error
+}
+
+func newMockCache() *mockCache {
+	return &mockCache{data: make(map[string][]byte)}
+}
+
+func (mc *mockCache) Get(key []byte, _ bool) ([]byte, bool, error) {
+	v, ok := mc.data[string(key)]
+	return v, ok, nil
+}
+
+func (mc *mockCache) BatchGet(keys map[string]types.BatchGetResult) error {
+	for k := range keys {
+		v, ok := mc.data[k]
+		if ok {
+			keys[k] = types.BatchGetResult{Value: v}
+		}
+	}
+	return nil
+}
+
+func (mc *mockCache) Set(key, value []byte) {
+	mc.data[string(key)] = value
+}
+
+func (mc *mockCache) Delete(key []byte) {
+	delete(mc.data, string(key))
+}
+
+func (mc *mockCache) BatchSet(updates []CacheUpdate) error {
+	if mc.batchSetErr != nil {
+		return mc.batchSetErr
+	}
+	for _, u := range updates {
+		if u.IsDelete() {
+			delete(mc.data, string(u.Key))
+		} else {
+			mc.data[string(u.Key)] = u.Value
+		}
+	}
+	return nil
+}
+
+// ---------------------------------------------------------------------------
+// tests
+// ---------------------------------------------------------------------------
+
+func TestCachedBatchCommitUpdatesCacheOnSuccess(t *testing.T) {
+	inner := &mockBatch{}
+	cache := newMockCache()
+	cb := newCachedBatch(inner, cache)
+
+	require.NoError(t, cb.Set([]byte("a"), []byte("1")))
+	require.NoError(t, cb.Set([]byte("b"), []byte("2")))
+	require.NoError(t, cb.Commit(types.WriteOptions{}))
+
+	require.True(t, inner.committed)
+	v, ok := cache.data["a"]
+	require.True(t, ok)
+	require.Equal(t, []byte("1"), v)
+	v, ok = cache.data["b"]
+	require.True(t, ok)
+	require.Equal(t, []byte("2"), v)
+}
+
+func TestCachedBatchCommitDoesNotUpdateCacheOnInnerFailure(t *testing.T) {
+	inner := &mockBatch{commitErr: errors.New("disk full")}
+	cache := newMockCache()
+	cb := newCachedBatch(inner, cache)
+
+	require.NoError(t, cb.Set([]byte("a"), []byte("1")))
+	err := cb.Commit(types.WriteOptions{})
+
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "disk full")
+	_, ok := cache.data["a"]
+	require.False(t, ok, "cache should not be updated when inner commit fails")
+}
+
+func TestCachedBatchCommitReturnsCacheError(t *testing.T) {
+	inner := &mockBatch{}
+	cache := newMockCache()
+	cache.batchSetErr = errors.New("cache broken")
+	cb := newCachedBatch(inner, cache)
+
+	require.NoError(t, cb.Set([]byte("a"), []byte("1")))
+	err := cb.Commit(types.WriteOptions{})
+
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "cache broken")
+	require.True(t, inner.committed, "inner batch should have committed")
+}
+
+func TestCachedBatchDeleteMarksKeyForRemoval(t *testing.T) {
+	inner := &mockBatch{}
+	cache := newMockCache()
+	cache.Set([]byte("x"), []byte("old"))
+	cb := newCachedBatch(inner, cache)
+
+	require.NoError(t, cb.Delete([]byte("x")))
+	require.NoError(t, cb.Commit(types.WriteOptions{}))
+
+	_, ok := cache.data["x"]
+	require.False(t, ok, "key should be deleted from cache")
+}
+
+func TestCachedBatchResetClearsPending(t *testing.T) {
+	inner := &mockBatch{}
+	cache := newMockCache()
+	cb := newCachedBatch(inner, cache)
+
+	require.NoError(t, cb.Set([]byte("a"), []byte("1")))
+	require.NoError(t, cb.Set([]byte("b"), []byte("2")))
+	cb.Reset()
+
+	require.NoError(t, cb.Commit(types.WriteOptions{}))
+
+	require.Empty(t, cache.data, "cache should have no entries after reset + commit")
+}
+
+func TestCachedBatchLenDelegatesToInner(t *testing.T) {
+	inner := &mockBatch{}
+	cache := newMockCache()
+	cb := newCachedBatch(inner, cache)
+
+	require.Equal(t, 0, cb.Len())
+	require.NoError(t, cb.Set([]byte("a"), []byte("1")))
+	require.NoError(t, cb.Delete([]byte("b")))
+	require.Equal(t, 2, cb.Len())
+}
+
+func TestCachedBatchCloseDelegatesToInner(t *testing.T) {
+	inner := &mockBatch{}
+	cache := newMockCache()
+	cb := newCachedBatch(inner, cache)
+
+	require.NoError(t, cb.Close())
+	require.True(t, inner.closed)
+}