Add additional Batch.Replay invariant tests (ava-labs#2567)

database/test_database.go

@@ -5,8 +5,11 @@ package database
 
 import (
 	"bytes"
+	"io"
 	"testing"
 
+	"github.com/golang/mock/gomock"
+
 	"github.com/stretchr/testify/require"
 
 	"golang.org/x/exp/slices"
@@ -30,6 +33,7 @@ var Tests = []func(t *testing.T, db Database){
 	TestBatchReuse,
 	TestBatchRewrite,
 	TestBatchReplay,
+	TestBatchReplayPropagateError,
 	TestBatchInner,
 	TestBatchLargeSize,
 	TestIteratorSnapshot,
@@ -341,6 +345,7 @@ func TestBatchReset(t *testing.T, db Database) {
 
 	batch.Reset()
 
+	require.Zero(batch.Size())
 	require.NoError(batch.Write())
 
 	has, err := db.Has(key)
@@ -376,6 +381,7 @@ func TestBatchReuse(t *testing.T, db Database) {
 
 	batch.Reset()
 
+	require.Zero(batch.Size())
 	require.NoError(batch.Put(key2, value2))
 	require.NoError(batch.Write())
 
@@ -425,6 +431,9 @@ func TestBatchRewrite(t *testing.T, db Database) {
 // TestBatchReplay tests to make sure that batches will correctly replay their
 // contents.
 func TestBatchReplay(t *testing.T, db Database) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
 	require := require.New(t)
 
 	key1 := []byte("hello1")
@@ -438,29 +447,55 @@ func TestBatchReplay(t *testing.T, db Database) {
 
 	require.NoError(batch.Put(key1, value1))
 	require.NoError(batch.Put(key2, value2))
+	require.NoError(batch.Delete(key1))
+	require.NoError(batch.Delete(key2))
+	require.NoError(batch.Put(key1, value2))
+
+	for i := 0; i < 2; i++ {
+		mockBatch := NewMockBatch(ctrl)
+		gomock.InOrder(
+			mockBatch.EXPECT().Put(key1, value1).Times(1),
+			mockBatch.EXPECT().Put(key2, value2).Times(1),
+			mockBatch.EXPECT().Delete(key1).Times(1),
+			mockBatch.EXPECT().Delete(key2).Times(1),
+			mockBatch.EXPECT().Put(key1, value2).Times(1),
+		)
+
+		require.NoError(batch.Replay(mockBatch))
+	}
+}
 
-	secondBatch := db.NewBatch()
-	require.NotNil(secondBatch)
+// TestBatchReplayPropagateError tests to make sure that batches will correctly
+// propagate any returned error during Replay.
+func TestBatchReplayPropagateError(t *testing.T, db Database) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
 
-	require.NoError(batch.Replay(secondBatch))
-	require.NoError(secondBatch.Write())
+	require := require.New(t)
 
-	has, err := db.Has(key1)
-	require.NoError(err)
-	require.True(has)
+	key1 := []byte("hello1")
+	value1 := []byte("world1")
 
-	v, err := db.Get(key1)
-	require.NoError(err)
-	require.Equal(value1, v)
+	key2 := []byte("hello2")
+	value2 := []byte("world2")
 
-	thirdBatch := db.NewBatch()
-	require.NotNil(thirdBatch)
+	batch := db.NewBatch()
+	require.NotNil(batch)
 
-	require.NoError(thirdBatch.Delete(key1))
-	require.NoError(thirdBatch.Delete(key2))
-	require.NoError(db.Close())
-	require.Equal(ErrClosed, batch.Replay(db))
-	require.Equal(ErrClosed, thirdBatch.Replay(db))
+	require.NoError(batch.Put(key1, value1))
+	require.NoError(batch.Put(key2, value2))
+
+	mockBatch := NewMockBatch(ctrl)
+	gomock.InOrder(
+		mockBatch.EXPECT().Put(key1, value1).Return(ErrClosed).Times(1),
+	)
+	require.Equal(ErrClosed, batch.Replay(mockBatch))
+
+	mockBatch = NewMockBatch(ctrl)
+	gomock.InOrder(
+		mockBatch.EXPECT().Put(key1, value1).Return(io.ErrClosedPipe).Times(1),
+	)
+	require.Equal(io.ErrClosedPipe, batch.Replay(mockBatch))
 }
 
 // TestBatchInner tests to make sure that inner can be used to write to the

x/merkledb/batch.go

@@ -7,12 +7,12 @@ import (
 	"golang.org/x/exp/slices"
 
 	"github.com/ava-labs/avalanchego/database"
-	"github.com/ava-labs/avalanchego/utils/linkedhashmap"
 )
 
 var _ database.Batch = &batch{}
 
 type batchOp struct {
+	key    []byte
 	value  []byte
 	delete bool
 }
@@ -21,29 +21,27 @@ type batchOp struct {
 // when Write is called.
 type batch struct {
 	db   *Database
-	data linkedhashmap.LinkedHashmap[string, batchOp]
+	data []batchOp
 	size int
 }
 
 // on [Write], put key, value into the database
 func (b *batch) Put(key []byte, value []byte) error {
-	b.putOp(
-		key,
-		batchOp{
-			value: slices.Clone(value),
-		},
-	)
+	b.data = append(b.data, batchOp{
+		key:   slices.Clone(key),
+		value: slices.Clone(value),
+	})
+	b.size += len(key) + len(value)
 	return nil
 }
 
 // on [Write], delete key from database
 func (b *batch) Delete(key []byte) error {
-	b.putOp(
-		key,
-		batchOp{
-			delete: true,
-		},
-	)
+	b.data = append(b.data, batchOp{
+		key:    slices.Clone(key),
+		delete: true,
+	})
+	b.size += len(key)
 	return nil
 }
 
@@ -60,20 +58,17 @@ func (b *batch) Write() error {
 }
 
 func (b *batch) Reset() {
-	b.data = linkedhashmap.New[string, batchOp]()
+	b.data = nil
 	b.size = 0
 }
 
 func (b *batch) Replay(w database.KeyValueWriterDeleter) error {
-	it := b.data.NewIterator()
-	for it.Next() {
-		key := []byte(it.Key())
-		op := it.Value()
+	for _, op := range b.data {
 		if op.delete {
-			if err := w.Delete(key); err != nil {
+			if err := w.Delete(op.key); err != nil {
 				return err
 			}
-		} else if err := w.Put(key, op.value); err != nil {
+		} else if err := w.Put(op.key, op.value); err != nil {
 			return err
 		}
 	}
@@ -83,13 +78,3 @@ func (b *batch) Replay(w database.KeyValueWriterDeleter) error {
 func (b *batch) Inner() database.Batch {
 	return b
 }
-
-func (b *batch) putOp(key []byte, op batchOp) {
-	stringKey := string(key)
-	lenKey := len(key)
-	if existing, ok := b.data.Get(stringKey); ok {
-		b.size -= lenKey + len(existing.value)
-	}
-	b.data.Put(stringKey, op)
-	b.size += lenKey + len(op.value)
-}