Add pebble database implementation

database/pebble/batch.go

@@ -0,0 +1,108 @@
+// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package pebble
+
+import (
+	"fmt"
+	"sync/atomic"
+
+	"github.com/cockroachdb/pebble"
+
+	"github.com/ava-labs/avalanchego/database"
+)
+
+var _ database.Batch = (*batch)(nil)
+
+// Not safe for concurrent use.
+type batch struct {
+	batch *pebble.Batch
+	db    *Database
+	size  int
+
+	// True iff [batch] has been written to the database
+	// since the last time [Reset] was called.
+	written atomic.Bool
+}
+
+func (db *Database) NewBatch() database.Batch {
+	return &batch{
+		db:    db,
+		batch: db.pebbleDB.NewBatch(),
+	}
+}
+
+func (b *batch) Put(key, value []byte) error {
+	b.size += len(key) + len(value) + pebbleByteOverHead
+	return b.batch.Set(key, value, pebble.Sync)
+}
+
+func (b *batch) Delete(key []byte) error {
+	b.size += len(key) + pebbleByteOverHead
+	return b.batch.Delete(key, pebble.Sync)
+}
+
+func (b *batch) Size() int {
+	return b.size
+}
+
+// Assumes [b.db.lock] is not held.
+func (b *batch) Write() error {
+	b.db.lock.RLock()
+	defer b.db.lock.RUnlock()
+
+	// Committing to a closed database makes pebble panic
+	// so make sure [b.db] isn't closed.
+	if b.db.closed {
+		return database.ErrClosed
+	}
+
+	if b.written.CompareAndSwap(false, true) {
+		// This batch has not been written to the database yet.
+		return updateError(b.batch.Commit(pebble.Sync))
+	}
+
+	// pebble doesn't support writing a batch twice so we have to clone
+	// [b] and commit the clone.
+	batchClone := b.db.pebbleDB.NewBatch()
+
+	// Copy the batch.
+	if err := batchClone.Apply(b.batch, nil); err != nil {
+		return err
+	}
+
+	// Commit the new batch.
+	return updateError(batchClone.Commit(pebble.Sync))
+}
+
+func (b *batch) Reset() {
+	b.batch.Reset()
+	b.written.Store(false)
+	b.size = 0
+}
+
+func (b *batch) Replay(w database.KeyValueWriterDeleter) error {
+	reader := b.batch.Reader()
+	for {
+		kind, k, v, ok := reader.Next()
+		if !ok {
+			return nil
+		}
+		switch kind {
+		case pebble.InternalKeyKindSet:
+			if err := w.Put(k, v); err != nil {
+				return err
+			}
+		case pebble.InternalKeyKindDelete:
+			if err := w.Delete(k); err != nil {
+				return err
+			}
+		default:
+			return fmt.Errorf("%w: %v", ErrInvalidOperation, kind)
+		}
+	}
+}
+
+func (b *batch) Inner() database.Batch {
+	return b
+}

database/pebble/batch_test.go

@@ -0,0 +1,47 @@
+// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package pebble
+
+import (
+	"os"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/avalanchego/utils/logging"
+)
+
+// Note: TestInterface tests other batch functionality.
+func TestBatch(t *testing.T) {
+	require := require.New(t)
+	dirName := os.TempDir()
+	defer os.Remove(dirName)
+
+	db, err := New(dirName, DefaultConfig, logging.NoLog{}, "", prometheus.NewRegistry())
+	require.NoError(err)
+
+	batchIntf := db.NewBatch()
+	batch, ok := batchIntf.(*batch)
+	require.True(ok)
+
+	require.False(batch.written.Load())
+
+	key1, value1 := []byte("key1"), []byte("value1")
+	require.NoError(batch.Put(key1, value1))
+	require.Equal(len(key1)+len(value1)+pebbleByteOverHead, batch.Size())
+
+	require.NoError(batch.Write())
+
+	require.True(batch.written.Load())
+
+	got, err := db.Get(key1)
+	require.NoError(err)
+	require.Equal(value1, got)
+
+	batch.Reset()
+	require.False(batch.written.Load())
+	require.Zero(batch.Size())
+}