De-duplicate batched writes by key in cached store

Author: julienrbrt

pkg/store/cached_store.go

@@ -121,7 +121,7 @@ func (cs *CachedStore) startWriteLoop() {
 			ops := []asyncWriteOp{op}
 
 			timer := time.NewTimer(batchWindow)
-			collect:
+		collect:
 			for {
 				select {
 				case op, ok := <-cs.writeCh:
@@ -135,17 +135,22 @@ func (cs *CachedStore) startWriteLoop() {
 				}
 			}
 
+			last := make(map[string]asyncWriteOp, len(ops))
+			for _, o := range ops {
+				last[o.key] = o
+			}
+
 			var puts []MetadataKV
 			var deletes []string
-			for _, o := range ops {
+			for _, o := range last {
 				if o.isDelete {
 					deletes = append(deletes, o.key)
 				} else {
 					puts = append(puts, MetadataKV{Key: o.key, Value: o.value})
 				}
 			}
 
-			if err := cs.Store.BatchMetadata(context.Background(), puts, deletes); err != nil {
+			if err := cs.BatchMetadata(context.Background(), puts, deletes); err != nil {
 				for _, o := range ops {
 					cs.logger.Error().Err(err).Str("key", o.key).
 						Bool("delete", o.isDelete).

pkg/store/cached_store_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	ds "github.com/ipfs/go-datastore"
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -337,7 +338,7 @@ func TestCachedStore_Close_FlushesPendingWrites(t *testing.T) {
 	cs.startWriteLoop()
 
 	const n = 100
-	for i := 0; i < n; i++ {
+	for i := range n {
 		k := fmt.Sprintf("key-%d", i)
 		require.NoError(t, cs.SetMetadata(ctx, k, []byte(k)))
 	}
@@ -348,7 +349,7 @@ func TestCachedStore_Close_FlushesPendingWrites(t *testing.T) {
 	cs.stopMu.Unlock()
 	<-done
 
-	for i := 0; i < n; i++ {
+	for i := range n {
 		k := fmt.Sprintf("key-%d", i)
 		v, err := base.GetMetadata(ctx, k)
 		require.NoError(t, err)
@@ -372,3 +373,38 @@ func TestCachedStore_WriteAfterClose_FallsBack(t *testing.T) {
 	err = cs.SetMetadata(ctx, "after", []byte("sync"))
 	require.Error(t, err)
 }
+
+func TestCachedStore_CoalescesSameKeyOps(t *testing.T) {
+	ctx := context.Background()
+
+	kv, err := NewTestInMemoryKVStore()
+	require.NoError(t, err)
+
+	require.NoError(t, kv.Put(ctx, ds.NewKey(GetMetaKey("k")), []byte("original")))
+
+	base := New(kv)
+
+	writeCh := make(chan asyncWriteOp, asyncWriteBufferSize)
+	done := make(chan struct{})
+	cs := &CachedStore{
+		Store:   base,
+		writeCh: writeCh,
+		done:    done,
+		logger:  zerolog.Nop(),
+	}
+	cs.startWriteLoop()
+
+	require.NoError(t, cs.SetMetadata(ctx, "k", []byte("v1")))
+	require.NoError(t, cs.DeleteMetadata(ctx, "k"))
+	require.NoError(t, cs.SetMetadata(ctx, "k", []byte("v2")))
+
+	cs.stopMu.Lock()
+	cs.stopped = true
+	close(writeCh)
+	cs.stopMu.Unlock()
+	<-done
+
+	v, err := base.GetMetadata(ctx, "k")
+	require.NoError(t, err)
+	require.Equal(t, []byte("v2"), v, "last write (Set) should win over delete")
+}