Add generics to x/sync (#4275)

Author: alarso16

x/sync/client_test.go renamed to x/merkledb/client_test.go

@@ -1,7 +1,7 @@
 // Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-package sync
+package merkledb
 
 import (
 	"context"
@@ -17,50 +17,57 @@ import (
 	"github.com/ava-labs/avalanchego/network/p2p"
 	"github.com/ava-labs/avalanchego/snow/engine/common"
 	"github.com/ava-labs/avalanchego/trace"
-	"github.com/ava-labs/avalanchego/x/merkledb"
 
 	pb "github.com/ava-labs/avalanchego/proto/pb/sync"
+	xsync "github.com/ava-labs/avalanchego/x/sync"
 )
 
-var _ p2p.Handler = (*flakyHandler)(nil)
+var (
+	_ p2p.Handler = (*xsync.GetChangeProofHandler[*RangeProof, *ChangeProof])(nil)
+	_ p2p.Handler = (*xsync.GetRangeProofHandler[*RangeProof, *ChangeProof])(nil)
+	_ p2p.Handler = (*flakyHandler)(nil)
+)
 
-func newDefaultDBConfig() merkledb.Config {
-	return merkledb.Config{
+func newDefaultDBConfig() Config {
+	return Config{
 		IntermediateWriteBatchSize:  100,
-		HistoryLength:               defaultRequestKeyLimit,
-		ValueNodeCacheSize:          defaultRequestKeyLimit,
-		IntermediateWriteBufferSize: defaultRequestKeyLimit,
-		IntermediateNodeCacheSize:   defaultRequestKeyLimit,
+		HistoryLength:               xsync.DefaultRequestKeyLimit,
+		ValueNodeCacheSize:          xsync.DefaultRequestKeyLimit,
+		IntermediateWriteBufferSize: xsync.DefaultRequestKeyLimit,
+		IntermediateNodeCacheSize:   xsync.DefaultRequestKeyLimit,
 		Reg:                         prometheus.NewRegistry(),
 		Tracer:                      trace.Noop,
-		BranchFactor:                merkledb.BranchFactor16,
+		BranchFactor:                BranchFactor16,
 	}
 }
 
 func newFlakyRangeProofHandler(
 	t *testing.T,
-	db merkledb.MerkleDB,
-	modifyResponse func(response *merkledb.RangeProof),
+	db MerkleDB,
+	modifyResponse func(response *RangeProof),
 ) p2p.Handler {
-	handler := NewGetRangeProofHandler(db)
+	var (
+		c                   = counter{m: 2}
+		rangeProofMarshaler = rangeProofMarshaler
+		handler             = xsync.NewGetRangeProofHandler(db, rangeProofMarshaler)
+	)
 
-	c := counter{m: 2}
 	return &p2p.TestHandler{
 		AppRequestF: func(ctx context.Context, nodeID ids.NodeID, deadline time.Time, requestBytes []byte) ([]byte, *common.AppError) {
 			responseBytes, appErr := handler.AppRequest(ctx, nodeID, deadline, requestBytes)
 			if appErr != nil {
 				return nil, appErr
 			}
 
-			proof := &merkledb.RangeProof{}
-			require.NoError(t, proof.UnmarshalBinary(responseBytes))
+			proof, err := rangeProofMarshaler.Unmarshal(responseBytes)
+			require.NoError(t, err)
 
 			// Half of requests are modified
 			if c.Inc() == 0 {
 				modifyResponse(proof)
 			}
 
-			responseBytes, err := proof.MarshalBinary()
+			responseBytes, err = rangeProofMarshaler.Marshal(proof)
 			if err != nil {
 				return nil, &common.AppError{Code: 123, Message: err.Error()}
 			}
@@ -72,15 +79,18 @@ func newFlakyRangeProofHandler(
 
 func newFlakyChangeProofHandler(
 	t *testing.T,
-	db merkledb.MerkleDB,
-	modifyResponse func(response *merkledb.ChangeProof),
+	db MerkleDB,
+	modifyResponse func(response *ChangeProof),
 ) p2p.Handler {
-	handler := NewGetChangeProofHandler(db)
+	var (
+		c                    = counter{m: 2}
+		rangeProofMarshaler  = rangeProofMarshaler
+		changeProofMarshaler = changeProofMarshaler
+		handler              = xsync.NewGetChangeProofHandler(db, rangeProofMarshaler, changeProofMarshaler)
+	)
 
-	c := counter{m: 2}
 	return &p2p.TestHandler{
 		AppRequestF: func(ctx context.Context, nodeID ids.NodeID, deadline time.Time, requestBytes []byte) ([]byte, *common.AppError) {
-			var err error
 			responseBytes, appErr := handler.AppRequest(ctx, nodeID, deadline, requestBytes)
 			if appErr != nil {
 				return nil, appErr
@@ -89,15 +99,15 @@ func newFlakyChangeProofHandler(
 			response := &pb.GetChangeProofResponse{}
 			require.NoError(t, proto.Unmarshal(responseBytes, response))
 
-			proof := &merkledb.ChangeProof{}
-			require.NoError(t, proof.UnmarshalBinary(response.GetChangeProof()))
+			proof, err := changeProofMarshaler.Unmarshal(response.GetChangeProof())
+			require.NoError(t, err)
 
 			// Half of requests are modified
 			if c.Inc() == 0 {
 				modifyResponse(proof)
 			}
 
-			proofBytes, err := proof.MarshalBinary()
+			proofBytes, err := changeProofMarshaler.Marshal(proof)
 			require.NoError(t, err)
 			responseBytes, err = proto.Marshal(&pb.GetChangeProofResponse{
 				Response: &pb.GetChangeProofResponse_ChangeProof{
@@ -142,14 +152,3 @@ func (c *counter) Inc() int {
 	c.i++
 	return result
 }
-
-type waitingHandler struct {
-	p2p.NoOpHandler
-	handler         p2p.Handler
-	updatedRootChan chan struct{}
-}
-
-func (w *waitingHandler) AppRequest(ctx context.Context, nodeID ids.NodeID, deadline time.Time, requestBytes []byte) ([]byte, *common.AppError) {
-	<-w.updatedRootChan
-	return w.handler.AppRequest(ctx, nodeID, deadline, requestBytes)
-}

x/merkledb/db.go

@@ -24,6 +24,7 @@ import (
 	"github.com/ava-labs/avalanchego/utils/set"
 	"github.com/ava-labs/avalanchego/utils/units"
 
+	xsync "github.com/ava-labs/avalanchego/x/sync"
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
@@ -37,7 +38,8 @@ const (
 )
 
 var (
-	_ MerkleDB = (*merkleDB)(nil)
+	_ MerkleDB                            = (*merkleDB)(nil)
+	_ xsync.DB[*RangeProof, *ChangeProof] = (MerkleDB)(nil)
 
 	metadataPrefix         = []byte{0}
 	valueNodePrefix        = []byte{1}
@@ -69,13 +71,12 @@ 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.
-	// Returns [ErrInsufficientHistory] if this node has insufficient history
+	// Returns [xsync.ErrInsufficientHistory] if this node has insufficient history
 	// to generate the proof.
 	// Returns ErrEmptyProof if [endRootID] is ids.Empty.
 	// Note that [endRootID] == ids.Empty means the trie is empty
 	// (i.e. we don't need a change proof.)
-	// Returns [ErrNoEndRoot], which wraps [ErrInsufficientHistory], if the
-	// history doesn't contain the [endRootID].
+	// Returns [xsync.ErrNoEndRoot], if the history doesn't contain the [endRootID].
 	GetChangeProof(
 		ctx context.Context,
 		startRootID ids.ID,
@@ -123,6 +124,7 @@ type RangeProofer interface {
 	// Returns ErrEmptyProof if [rootID] is ids.Empty.
 	// Note that [rootID] == ids.Empty means the trie is empty
 	// (i.e. we don't need a range proof.)
+	// Returns [xsync.ErrNoEndRoot], if the history doesn't contain the [rootID].
 	GetRangeProofAtRoot(
 		ctx context.Context,
 		rootID ids.ID,