feedback

Author: julienrbrt

pkg/store/cached_store.go

@@ -3,6 +3,7 @@ package store
 import (
 	"context"
 	"sync"
+	"time"
 
 	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/rs/zerolog"
@@ -18,6 +19,11 @@ const (
 	DefaultBlockDataCacheSize = 200_000
 
 	asyncWriteBufferSize = 8192
+
+	// batchWindow is the time the write goroutine waits after receiving the first
+	// op before flushing. This allows bursts of metadata writes (e.g. 3-4 per
+	// height in the submitter) to be coalesced into a single Badger WriteBatch.
+	batchWindow = 100 * time.Microsecond
 )
 
 type asyncWriteOp struct {
@@ -113,13 +119,19 @@ func (cs *CachedStore) startWriteLoop() {
 		defer close(cs.done)
 		for op := range cs.writeCh {
 			ops := []asyncWriteOp{op}
-		drain:
+
+			timer := time.NewTimer(batchWindow)
+			collect:
 			for {
 				select {
-				case op := <-cs.writeCh:
+				case op, ok := <-cs.writeCh:
+					if !ok {
+						timer.Stop()
+						break collect
+					}
 					ops = append(ops, op)
-				default:
-					break drain
+				case <-timer.C:
+					break collect
 				}
 			}
 
@@ -234,40 +246,33 @@ func (cs *CachedStore) PruneBlocks(ctx context.Context, height uint64) error {
 }
 
 // SetMetadata queues an asynchronous metadata write. The write is persisted
-// by the background goroutine. If the buffer is full or the store has been
-// stopped, the write falls back to synchronous execution on the underlying store.
+// by the background goroutine via BatchMetadata. If the store has been stopped,
+// the write falls back to synchronous execution on the underlying store.
 func (cs *CachedStore) SetMetadata(ctx context.Context, key string, value []byte) error {
 	cs.stopMu.RLock()
 	if cs.stopped {
 		cs.stopMu.RUnlock()
 		return cs.Store.SetMetadata(ctx, key, value)
 	}
-	select {
-	case cs.writeCh <- asyncWriteOp{key: key, value: value}:
-		cs.stopMu.RUnlock()
-		return nil
-	default:
-		cs.stopMu.RUnlock()
-		return cs.Store.SetMetadata(ctx, key, value)
-	}
+	cs.stopMu.RUnlock()
+
+	valueCopy := append([]byte(nil), value...)
+	cs.writeCh <- asyncWriteOp{key: key, value: valueCopy}
+	return nil
 }
 
-// DeleteMetadata queues an asynchronous metadata delete. If the buffer is full
-// or the store has been stopped, the delete falls back to synchronous execution.
+// DeleteMetadata queues an asynchronous metadata delete. If the store has been
+// stopped, the delete falls back to synchronous execution.
 func (cs *CachedStore) DeleteMetadata(ctx context.Context, key string) error {
 	cs.stopMu.RLock()
 	if cs.stopped {
 		cs.stopMu.RUnlock()
 		return cs.Store.DeleteMetadata(ctx, key)
 	}
-	select {
-	case cs.writeCh <- asyncWriteOp{key: key, isDelete: true}:
-		cs.stopMu.RUnlock()
-		return nil
-	default:
-		cs.stopMu.RUnlock()
-		return cs.Store.DeleteMetadata(ctx, key)
-	}
+	cs.stopMu.RUnlock()
+
+	cs.writeCh <- asyncWriteOp{key: key, isDelete: true}
+	return nil
 }
 
 // Close drains pending async writes, then closes the underlying store.

pkg/store/cached_store_test.go

@@ -2,9 +2,11 @@ package store
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
+	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -315,27 +317,43 @@ func TestCachedStore_AsyncDeleteMetadata(t *testing.T) {
 }
 
 func TestCachedStore_Close_FlushesPendingWrites(t *testing.T) {
+	ctx := context.Background()
+
 	kv, err := NewTestInMemoryKVStore()
 	require.NoError(t, err)
 
 	base := New(kv)
-	cs, err := NewCachedStore(base)
-	require.NoError(t, err)
 
-	ctx := context.Background()
+	writeCh := make(chan asyncWriteOp, asyncWriteBufferSize)
+	done := make(chan struct{})
+
+	cs := &CachedStore{
+		Store:   base,
+		writeCh: writeCh,
+		done:    done,
+		logger:  zerolog.Nop(),
+	}
+
+	cs.startWriteLoop()
+
 	const n = 100
 	for i := 0; i < n; i++ {
-		k := []byte{byte(i)}
-		require.NoError(t, cs.SetMetadata(ctx, string(k), k))
+		k := fmt.Sprintf("key-%d", i)
+		require.NoError(t, cs.SetMetadata(ctx, k, []byte(k)))
 	}
 
-	// Wait for the last key to land via pass-through read
-	require.Eventually(t, func() bool {
-		v, err := cs.GetMetadata(ctx, string([]byte{byte(n - 1)}))
-		return err == nil && len(v) == 1 && v[0] == byte(n-1)
-	}, 2*time.Second, 10*time.Millisecond)
+	cs.stopMu.Lock()
+	cs.stopped = true
+	close(writeCh)
+	cs.stopMu.Unlock()
+	<-done
 
-	require.NoError(t, cs.Close())
+	for i := 0; i < n; i++ {
+		k := fmt.Sprintf("key-%d", i)
+		v, err := base.GetMetadata(ctx, k)
+		require.NoError(t, err)
+		require.Equal(t, []byte(k), v)
+	}
 }
 
 func TestCachedStore_WriteAfterClose_FallsBack(t *testing.T) {