x/sync / x/merkledb -- add SyncableDB interface (ava-labs#1555)

joshua-kim · May 30, 2023 · 243e313 · 243e313
1 parent b66e25e
commit 243e313
Show file tree

Hide file tree

Showing 9 changed files with 71 additions and 43 deletions.
diff --git a/x/merkledb/db.go b/x/merkledb/db.go
@@ -10,8 +10,6 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/prometheus/client_golang/prometheus"
-
 	"go.opentelemetry.io/otel/attribute"
 
 	oteltrace "go.opentelemetry.io/otel/trace"
@@ -27,6 +25,7 @@ import (
 	"github.com/ava-labs/avalanchego/utils"
 	"github.com/ava-labs/avalanchego/utils/math"
 	"github.com/ava-labs/avalanchego/utils/set"
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 const (
@@ -53,22 +52,7 @@ var (
 	errSameRoot = errors.New("start and end root are the same")
 )
 
-type Config struct {
-	// The number of changes to the database that we store in memory in order to
-	// serve change proofs.
-	HistoryLength int
-	NodeCacheSize int
-	// If [Reg] is nil, metrics are collected locally but not exported through
-	// Prometheus.
-	// This may be useful for testing.
-	Reg    prometheus.Registerer
-	Tracer trace.Tracer
-}
-
-type MerkleDB interface {
-	database.Database
-	Trie
-
+type ChangeProofer interface {
 	// GetChangeProof returns a proof for a subset of the key/value changes in key range
 	// [start, end] that occurred between [startRootID] and [endRootID].
 	// Returns at most [maxLength] key/value pairs.
@@ -105,7 +89,9 @@ type MerkleDB interface {
 
 	// CommitChangeProof commits the key/value pairs within the [proof] to the db.
 	CommitChangeProof(ctx context.Context, proof *ChangeProof) error
+}
 
+type RangeProofer interface {
 	// GetRangeProofAtRoot returns a proof for the key/value pairs in this trie within the range
 	// [start, end] when the root of the trie was [rootID].
 	GetRangeProofAtRoot(
@@ -121,6 +107,27 @@ type MerkleDB interface {
 	CommitRangeProof(ctx context.Context, start []byte, proof *RangeProof) error
 }
 
+type MerkleDB interface {
+	database.Database
+	Trie
+	MerkleRootGetter
+	ProofGetter
+	ChangeProofer
+	RangeProofer
+}
+
+type Config struct {
+	// The number of changes to the database that we store in memory in order to
+	// serve change proofs.
+	HistoryLength int
+	NodeCacheSize int
+	// If [Reg] is nil, metrics are collected locally but not exported through
+	// Prometheus.
+	// This may be useful for testing.
+	Reg    prometheus.Registerer
+	Tracer trace.Tracer
+}
+
 // merkleDB can only be edited by committing changes from a trieView.
 type merkleDB struct {
 	// Must be held when reading/writing fields.

diff --git a/x/merkledb/trie.go b/x/merkledb/trie.go
@@ -13,7 +13,21 @@ import (
 
 var errNoNewRoot = errors.New("there was no updated root in change list")
 
+type MerkleRootGetter interface {
+	// GetMerkleRoot returns the merkle root of the Trie
+	GetMerkleRoot(ctx context.Context) (ids.ID, error)
+}
+
+type ProofGetter interface {
+	// GetProof generates a proof of the value associated with a particular key,
+	// or a proof of its absence from the trie
+	GetProof(ctx context.Context, bytesPath []byte) (*Proof, error)
+}
+
 type ReadOnlyTrie interface {
+	MerkleRootGetter
+	ProofGetter
+
 	// GetValue gets the value associated with the specified key
 	// database.ErrNotFound if the key is not present
 	GetValue(ctx context.Context, key []byte) ([]byte, error)
@@ -26,15 +40,9 @@ type ReadOnlyTrie interface {
 	// database.ErrNotFound if the key is not present
 	getValue(key path, lock bool) ([]byte, error)
 
-	// GetMerkleRoot returns the merkle root of the Trie
-	GetMerkleRoot(ctx context.Context) (ids.ID, error)
-
 	// get an editable copy of the node with the given key path
 	getEditableNode(key path) (*node, error)
 
-	// GetProof generates a proof of the value associated with a particular key, or a proof of its absence from the trie
-	GetProof(ctx context.Context, bytesPath []byte) (*Proof, error)
-
 	// GetRangeProof generates a proof of up to maxLength smallest key/values with keys between start and end
 	GetRangeProof(ctx context.Context, start, end []byte, maxLength int) (*RangeProof, error)
 

diff --git a/x/sync/client.go b/x/sync/client.go
@@ -45,7 +45,7 @@ type Client interface {
 	// GetChangeProof synchronously sends the given request, returning a parsed ChangesResponse or error
 	// [verificationDB] is the local db that has all key/values in it for the proof's startroot within the proof's key range
 	// Note: this verifies the response including the change proof.
-	GetChangeProof(ctx context.Context, request *syncpb.ChangeProofRequest, verificationDB merkledb.MerkleDB) (*merkledb.ChangeProof, error)
+	GetChangeProof(ctx context.Context, request *syncpb.ChangeProofRequest, verificationDB SyncableDB) (*merkledb.ChangeProof, error)
 }
 
 type client struct {
@@ -79,7 +79,7 @@ func NewClient(config *ClientConfig) Client {
 // GetChangeProof synchronously retrieves the change proof given by [req].
 // Upon failure, retries until the context is expired.
 // The returned change proof is verified.
-func (c *client) GetChangeProof(ctx context.Context, req *syncpb.ChangeProofRequest, db merkledb.MerkleDB) (*merkledb.ChangeProof, error) {
+func (c *client) GetChangeProof(ctx context.Context, req *syncpb.ChangeProofRequest, db SyncableDB) (*merkledb.ChangeProof, error) {
 	parseFn := func(ctx context.Context, responseBytes []byte) (*merkledb.ChangeProof, error) {
 		if len(responseBytes) > int(req.BytesLimit) {
 			return nil, fmt.Errorf("%w: (%d) > %d)", errTooManyBytes, len(responseBytes), req.BytesLimit)

diff --git a/x/sync/client_test.go b/x/sync/client_test.go
@@ -27,7 +27,7 @@ import (
 
 func sendRangeRequest(
 	t *testing.T,
-	db merkledb.MerkleDB,
+	db SyncableDB,
 	request *syncpb.RangeProofRequest,
 	maxAttempts uint32,
 	modifyResponse func(*merkledb.RangeProof),
@@ -123,7 +123,7 @@ func TestGetRangeProof(t *testing.T) {
 	require.NoError(t, err)
 
 	tests := map[string]struct {
-		db                  merkledb.MerkleDB
+		db                  SyncableDB
 		request             *syncpb.RangeProofRequest
 		modifyResponse      func(*merkledb.RangeProof)
 		expectedErr         error
@@ -209,10 +209,10 @@ func TestGetRangeProof(t *testing.T) {
 			},
 			modifyResponse: func(response *merkledb.RangeProof) {
 				start := response.KeyValues[1].Key
-				proof, err := largeTrieDB.GetRangeProof(context.Background(), start, nil, defaultRequestKeyLimit)
-				if err != nil {
-					panic(err)
-				}
+				rootID, err := largeTrieDB.GetMerkleRoot(context.Background())
+				require.NoError(t, err)
+				proof, err := largeTrieDB.GetRangeProofAtRoot(context.Background(), rootID, start, nil, defaultRequestKeyLimit)
+				require.NoError(t, err)
 				response.KeyValues = proof.KeyValues
 				response.StartProof = proof.StartProof
 				response.EndProof = proof.EndProof
@@ -279,8 +279,8 @@ func TestGetRangeProof(t *testing.T) {
 
 func sendChangeRequest(
 	t *testing.T,
-	db merkledb.MerkleDB,
-	verificationDB merkledb.MerkleDB,
+	db SyncableDB,
+	verificationDB SyncableDB,
 	request *syncpb.ChangeProofRequest,
 	maxAttempts uint32,
 	modifyResponse func(*merkledb.ChangeProof),
@@ -423,7 +423,7 @@ func TestGetChangeProof(t *testing.T) {
 	require.NoError(t, err)
 
 	tests := map[string]struct {
-		db                  merkledb.MerkleDB
+		db                  SyncableDB
 		request             *syncpb.ChangeProofRequest
 		modifyResponse      func(*merkledb.ChangeProof)
 		expectedErr         error

diff --git a/x/sync/mock_client.go b/x/sync/mock_client.go
diff --git a/x/sync/network_server.go b/x/sync/network_server.go
@@ -45,11 +45,11 @@ var ErrMinProofSizeIsTooLarge = errors.New("cannot generate any proof within the
 
 type NetworkServer struct {
 	appSender common.AppSender // Used to respond to peer requests via AppResponse.
-	db        merkledb.MerkleDB
+	db        SyncableDB
 	log       logging.Logger
 }
 
-func NewNetworkServer(appSender common.AppSender, db merkledb.MerkleDB, log logging.Logger) *NetworkServer {
+func NewNetworkServer(appSender common.AppSender, db SyncableDB, log logging.Logger) *NetworkServer {
 	return &NetworkServer{
 		appSender: appSender,
 		db:        db,

diff --git a/x/sync/sync_test.go b/x/sync/sync_test.go
@@ -34,10 +34,10 @@ func newNoopTracer() trace.Tracer {
 }
 
 type mockClient struct {
-	db merkledb.MerkleDB
+	db SyncableDB
 }
 
-func (client *mockClient) GetChangeProof(ctx context.Context, request *syncpb.ChangeProofRequest, _ merkledb.MerkleDB) (*merkledb.ChangeProof, error) {
+func (client *mockClient) GetChangeProof(ctx context.Context, request *syncpb.ChangeProofRequest, _ SyncableDB) (*merkledb.ChangeProof, error) {
 	startRoot, err := ids.ToID(request.StartRootHash)
 	if err != nil {
 		return nil, err
@@ -871,7 +871,7 @@ func Test_Sync_Error_During_Sync(t *testing.T) {
 		},
 	).AnyTimes()
 	client.EXPECT().GetChangeProof(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
-		func(ctx context.Context, request *syncpb.ChangeProofRequest, _ merkledb.MerkleDB) (*merkledb.ChangeProof, error) {
+		func(ctx context.Context, request *syncpb.ChangeProofRequest, _ SyncableDB) (*merkledb.ChangeProof, error) {
 			startRoot, err := ids.ToID(request.StartRootHash)
 			require.NoError(err)
 			endRoot, err := ids.ToID(request.EndRootHash)
@@ -961,7 +961,7 @@ func Test_Sync_Result_Correct_Root_Update_Root_During(t *testing.T) {
 			},
 		).AnyTimes()
 		client.EXPECT().GetChangeProof(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
-			func(ctx context.Context, request *syncpb.ChangeProofRequest, _ merkledb.MerkleDB) (*merkledb.ChangeProof, error) {
+			func(ctx context.Context, request *syncpb.ChangeProofRequest, _ SyncableDB) (*merkledb.ChangeProof, error) {
 				<-updatedRootChan
 				startRoot, err := ids.ToID(request.StartRootHash)
 				require.NoError(err)

diff --git a/x/sync/syncable_db.go b/x/sync/syncable_db.go
@@ -0,0 +1,13 @@
+// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package sync
+
+import "github.com/ava-labs/avalanchego/x/merkledb"
+
+type SyncableDB interface {
+	merkledb.MerkleRootGetter
+	merkledb.ProofGetter
+	merkledb.ChangeProofer
+	merkledb.RangeProofer
+}
diff --git a/x/sync/syncmanager.go b/x/sync/syncmanager.go
@@ -106,7 +106,7 @@ type StateSyncManager struct {
 }
 
 type StateSyncConfig struct {
-	SyncDB                merkledb.MerkleDB
+	SyncDB                SyncableDB
 	Client                Client
 	SimultaneousWorkLimit int
 	Log                   logging.Logger