Move approval store into MSM

msm/approvals.go

@@ -5,6 +5,7 @@ package metadata
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/ava-labs/simplex/common"
 	"go.uber.org/zap"
@@ -26,9 +27,13 @@ type ApprovalStore struct {
 	signatureVerifier SignatureVerifier
 	validators        NodeBLSMappings
 	logger            common.Logger
-	pkByNodeID        map[nodeID][]byte
-	approvalsByNodes  map[nodeID]approvalsByPChainHeightAndAuxInfoDigest
-	storedCount       int
+	nodeIDToPK        map[nodeID][]byte
+	// lock guards the mutable state below (approvalsByNodes, storedCount). The
+	// store is accessed concurrently: approvals are handled as they arrive while
+	// the block builder reads the accumulated approvals when building a block.
+	lock             sync.RWMutex
+	approvalsByNodes map[nodeID]approvalsByPChainHeightAndAuxInfoDigest
+	storedCount      int
 }
 
 func NewApprovalStore(signatureVerifier SignatureVerifier, validators NodeBLSMappings, logger common.Logger) *ApprovalStore {
@@ -45,13 +50,16 @@ func NewApprovalStore(signatureVerifier SignatureVerifier, validators NodeBLSMap
 	return &ApprovalStore{
 		signatureVerifier: signatureVerifier,
 		validators:        validators,
-		pkByNodeID:        pkByNodeID,
+		nodeIDToPK:        pkByNodeID,
 		logger:            logger,
 		approvalsByNodes:  approvalsByNodes,
 	}
 }
 
 func (as *ApprovalStore) Approvals() ValidatorSetApprovals {
+	as.lock.RLock()
+	defer as.lock.RUnlock()
+
 	approvals := make(ValidatorSetApprovals, 0, as.storedCount)
 	for _, approvalsByHeight := range as.approvalsByNodes {
 		for _, approval := range approvalsByHeight {
@@ -63,24 +71,28 @@ func (as *ApprovalStore) Approvals() ValidatorSetApprovals {
 
 func (as *ApprovalStore) HandleApproval(approval *ValidatorSetApproval, timestamp uint64) error {
 	// First thing we check is if the node that sent this approval is a validator.
-	pk, exists := as.getPKOfNode(approval.NodeID)
+	pk, exists := as.nodeIDToPK[approval.NodeID]
 	if !exists {
 		as.logger.Debug("Received an approval from a node that is not a validator", zap.String("nodeID",
 			fmt.Sprintf("%x", approval.NodeID)), zap.Uint64("pChainHeight", approval.PChainHeight))
 		return nil
 	}
 
-	// Second thing we check is if we already have an approval for this height from this node.
-	if as.approvalExistsAndUpToDate(approval, timestamp) {
-		as.logger.Debug("Already have an approval from the node", zap.String("nodeID",
+	// Second thing we check is if the signature of the approval is valid.
+	// We need it to be valid in order for nodes to be able to aggregate it later on along with other approvals.
+	// This is checked before taking the lock, as it only reads immutable state.
+	if err := as.checkApprovalSignature(approval, pk); err != nil {
+		as.logger.Debug("Received an approval with an invalid signature", zap.String("nodeID",
 			fmt.Sprintf("%x", approval.NodeID)), zap.Uint64("pChainHeight", approval.PChainHeight))
 		return nil
 	}
 
-	// Third thing we check is if the signature of the approval is valid.
-	// We need it to be valid in order for nodes to be able to aggregate it later on along with other approvals.
-	if err := as.checkApprovalSignature(approval, pk); err != nil {
-		as.logger.Debug("Received an approval with an invalid signature", zap.String("nodeID",
+	as.lock.Lock()
+	defer as.lock.Unlock()
+
+	// Third thing we check is if we already have an approval for this height from this node.
+	if as.approvalExistsAndUpToDate(approval, timestamp) {
+		as.logger.Debug("Already have an approval from the node", zap.String("nodeID",
 			fmt.Sprintf("%x", approval.NodeID)), zap.Uint64("pChainHeight", approval.PChainHeight))
 		return nil
 	}
@@ -145,11 +157,6 @@ func (as *ApprovalStore) checkApprovalSignature(approval *ValidatorSetApproval,
 	return as.signatureVerifier.VerifySignature(approval.Signature, toBeSigned, pk)
 }
 
-func (as *ApprovalStore) getPKOfNode(nodeID nodeID) ([]byte, bool) {
-	pk, exists := as.pkByNodeID[nodeID]
-	return pk, exists
-}
-
 func (as *ApprovalStore) approvalExistsAndUpToDate(approval *ValidatorSetApproval, timestamp uint64) bool {
 	if as.approvalsByNodes[approval.NodeID] == nil {
 		return false
@@ -167,3 +174,25 @@ func (as *ApprovalStore) approvalExistsAndUpToDate(approval *ValidatorSetApprova
 
 	return existingApproval.Timestamp >= timestamp
 }
+
+// PutApprovals copies all approvals from this store to the given approvalStore.
+func (as *ApprovalStore) PutApprovals(approvalStore *ApprovalStore) {
+	// Snapshot the approvals under our lock, then hand them to the destination
+	// store (which takes its own lock). Copying first avoids holding two store
+	// locks at once.
+	as.lock.RLock()
+	snapshot := make([]approvalAndTimestamp, 0, as.storedCount)
+	for _, approvalsByHeight := range as.approvalsByNodes {
+		for _, approval := range approvalsByHeight {
+			snapshot = append(snapshot, approvalAndTimestamp{
+				ValidatorSetApproval: approval.ValidatorSetApproval,
+				Timestamp:            approval.Timestamp,
+			})
+		}
+	}
+	as.lock.RUnlock()
+
+	for _, a := range snapshot {
+		approvalStore.HandleApproval(&a.ValidatorSetApproval, a.Timestamp)
+	}
+}

msm/approvals_test.go

@@ -222,6 +222,172 @@ func TestApprovalStoreHandleApproval(t *testing.T) {
 	}
 }
 
+func TestApprovalStorePutApprovals(t *testing.T) {
+	// findApproval locates the approval for a given (NodeID, PChainHeight) in a slice returned by
+	// Approvals(). It is used to assert which of two competing approvals (source vs. destination)
+	// survived a merge, since Approvals() does not carry the timestamp but does carry the Signature,
+	// and signApproval produces a distinct signature per call.
+	findApproval := func(got ValidatorSetApprovals, node nodeID, height uint64) (ValidatorSetApproval, bool) {
+		for _, a := range got {
+			if a.NodeID == node && a.PChainHeight == height {
+				return a, true
+			}
+		}
+		return ValidatorSetApproval{}, false
+	}
+
+	for _, tc := range []struct {
+		name string
+		// srcValidators/dstValidators size the two stores via makeValidators; NodeIDs are makeNodeID(i+1).
+		srcValidators int
+		dstValidators int
+		// srcApprovals are loaded into the source store; dstApprovals are pre-loaded into the destination
+		// store before PutApprovals is called.
+		srcApprovals []approvalAndTimestamp
+		dstApprovals []approvalAndTimestamp
+		// verify asserts on the destination (and source) store after src.PutApprovals(dst).
+		verify func(t *testing.T, dst, src *ApprovalStore, srcSent, dstSent []approvalAndTimestamp)
+	}{
+		{
+			// Copies every source approval into an empty destination that shares the same validator set,
+			// and leaves the source store untouched (PutApprovals copies, it does not move).
+			name:          "copies all approvals into empty destination",
+			srcValidators: 3,
+			dstValidators: 3,
+			srcApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 100},
+				{ValidatorSetApproval{NodeID: makeNodeID(2), PChainHeight: 8, Signature: signApproval(8, [32]byte{})}, 100},
+			},
+			verify: func(t *testing.T, dst, src *ApprovalStore, srcSent, _ []approvalAndTimestamp) {
+				got := dst.Approvals()
+				require.Len(t, got, 2)
+				require.Equal(t, 2, dst.storedCount)
+				require.ElementsMatch(t, []ValidatorSetApproval{srcSent[0].ValidatorSetApproval, srcSent[1].ValidatorSetApproval}, got)
+				// The source store is unchanged.
+				require.Len(t, src.Approvals(), 2)
+				require.Equal(t, 2, src.storedCount)
+			},
+		},
+		{
+			// Approvals from nodes absent from the destination's (smaller) validator set are dropped on
+			// carry-over. This is the epoch-change case: a validator that is not part of the new set does
+			// not have its approval carried into the new store.
+			name:          "drops approvals from nodes not in destination validator set",
+			srcValidators: 3,
+			dstValidators: 2,
+			srcApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 1, Signature: signApproval(1, [32]byte{})}, 10},
+				{ValidatorSetApproval{NodeID: makeNodeID(2), PChainHeight: 1, Signature: signApproval(1, [32]byte{})}, 10},
+				{ValidatorSetApproval{NodeID: makeNodeID(3), PChainHeight: 1, Signature: signApproval(1, [32]byte{})}, 10},
+			},
+			verify: func(t *testing.T, dst, _ *ApprovalStore, _, _ []approvalAndTimestamp) {
+				got := dst.Approvals()
+				require.Len(t, got, 2, "only approvals from nodes in the destination set carry over")
+				require.Equal(t, 2, dst.storedCount)
+				_, ok := findApproval(got, makeNodeID(3), 1)
+				require.False(t, ok, "node 3 is not in the destination validator set")
+			},
+		},
+		{
+			// Merges into a non-empty destination: a pre-existing destination approval with a newer
+			// timestamp is kept over an older one carried from the source, while a brand-new node's
+			// approval from the source is added.
+			name:          "keeps newer destination approval and adds new node",
+			srcValidators: 2,
+			dstValidators: 2,
+			dstApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 200}, // newer, already present
+			},
+			srcApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 100}, // older, must not overwrite
+				{ValidatorSetApproval{NodeID: makeNodeID(2), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 100}, // new node, added
+			},
+			verify: func(t *testing.T, dst, _ *ApprovalStore, _, dstSent []approvalAndTimestamp) {
+				got := dst.Approvals()
+				require.Len(t, got, 2)
+				require.Equal(t, 2, dst.storedCount)
+				kept, ok := findApproval(got, makeNodeID(1), 7)
+				require.True(t, ok)
+				require.Equal(t, dstSent[0].ValidatorSetApproval, kept, "the newer destination approval is retained")
+				_, ok = findApproval(got, makeNodeID(2), 7)
+				require.True(t, ok, "the new node's approval is carried over")
+			},
+		},
+		{
+			// A source approval with a newer timestamp replaces the stale destination approval at the
+			// same (NodeID, PChainHeight).
+			name:          "newer source approval replaces stale destination approval",
+			srcValidators: 2,
+			dstValidators: 2,
+			dstApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 100}, // older, present
+			},
+			srcApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 200}, // newer, replaces
+			},
+			verify: func(t *testing.T, dst, _ *ApprovalStore, srcSent, _ []approvalAndTimestamp) {
+				got := dst.Approvals()
+				require.Len(t, got, 1)
+				require.Equal(t, 1, dst.storedCount)
+				require.Equal(t, srcSent[0].ValidatorSetApproval, got[0], "the newer source approval replaces the stale one")
+			},
+		},
+		{
+			// An empty source store is a no-op: the destination is left exactly as it was.
+			name:          "empty source is a no-op",
+			srcValidators: 2,
+			dstValidators: 2,
+			dstApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 1, Signature: signApproval(1, [32]byte{})}, 10},
+			},
+			verify: func(t *testing.T, dst, _ *ApprovalStore, _, dstSent []approvalAndTimestamp) {
+				got := dst.Approvals()
+				require.Len(t, got, 1)
+				require.Equal(t, 1, dst.storedCount)
+				require.Equal(t, dstSent[0].ValidatorSetApproval, got[0])
+			},
+		},
+		{
+			// A source approval with the same timestamp as an existing destination approval at the same
+			// (NodeID, PChainHeight) does not overwrite it: the >= tie in approvalExistsAndUpToDate goes
+			// to the current destination.
+			name:          "equal-timestamp source does not overwrite destination",
+			srcValidators: 2,
+			dstValidators: 2,
+			dstApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 100},
+			},
+			srcApprovals: []approvalAndTimestamp{
+				{ValidatorSetApproval{NodeID: makeNodeID(1), PChainHeight: 7, Signature: signApproval(7, [32]byte{})}, 100},
+			},
+			verify: func(t *testing.T, dst, _ *ApprovalStore, _, dstSent []approvalAndTimestamp) {
+				got := dst.Approvals()
+				require.Len(t, got, 1)
+				require.Equal(t, 1, dst.storedCount)
+				require.Equal(t, dstSent[0].ValidatorSetApproval, got[0], "the destination approval is retained on a timestamp tie")
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			src := NewApprovalStore(&signatureVerifier{}, makeValidators(tc.srcValidators), testutil.MakeLogger(t))
+			dst := NewApprovalStore(&signatureVerifier{}, makeValidators(tc.dstValidators), testutil.MakeLogger(t))
+
+			for _, a := range tc.srcApprovals {
+				require.NoError(t, src.HandleApproval(&a.ValidatorSetApproval, a.Timestamp))
+			}
+			for _, a := range tc.dstApprovals {
+				require.NoError(t, dst.HandleApproval(&a.ValidatorSetApproval, a.Timestamp))
+			}
+
+			src.PutApprovals(dst)
+
+			// storedCount must always stay in sync with the number of retrievable approvals.
+			require.Len(t, dst.Approvals(), dst.storedCount)
+			tc.verify(t, dst, src, tc.srcApprovals, tc.dstApprovals)
+		})
+	}
+}
+
 func TestApprovalStoreHandleApprovalStoredCountStaysConsistent(t *testing.T) {
 	// Runs a mixed workload (insert, duplicate, replace, new height, prune)
 	// and asserts that storedCount equals len(Approvals()) after every step.

msm/build_decision.go

@@ -18,6 +18,7 @@ type blockBuildingDecision struct {
 	buildInnerBlock bool
 	transitionEpoch bool
 	pChainHeight    uint64
+	validatorSet    NodeBLSMappings
 }
 
 // PChainProgressListener listens for changes in the P-chain height.
@@ -33,7 +34,7 @@ type blockBuildingDecider struct {
 	waitForPendingBlock      func(ctx context.Context)
 	// hasValidatorSetChanged should return whether the validator set has changed since the
 	// P-chain height referenced by the last block in the chain and until the provided P-chain height.
-	hasValidatorSetChanged func(pChainHeight uint64) (bool, error)
+	hasValidatorSetChanged func(pChainHeight uint64) (bool, NodeBLSMappings, error)
 	getPChainHeight        func() uint64
 }
 
@@ -48,14 +49,14 @@ func (bbd *blockBuildingDecider) shouldBuildBlock(
 	for {
 		pChainHeight := bbd.getPChainHeight()
 
-		shouldTransitionEpoch, err := bbd.hasValidatorSetChanged(pChainHeight)
+		shouldTransitionEpoch, newValidatorSet, err := bbd.hasValidatorSetChanged(pChainHeight)
 		if err != nil {
 			return blockBuildingDecision{}, err
 		}
 
 		if shouldTransitionEpoch {
 			// If we should transition to a new epoch, maybe we can also build a block along the way.
-			return bbd.buildBlockWithEpochTransition(ctx, pChainHeight)
+			return bbd.buildBlockWithEpochTransition(ctx, pChainHeight, newValidatorSet)
 		}
 
 		// Else, we don't need to transition to a new epoch, but maybe we should build a block.
@@ -109,7 +110,7 @@ func (bbd *blockBuildingDecider) waitForPChainChangeOrPendingBlock(ctx context.C
 // It waits up to a limited amount of time (bbd.maxBlockBuildingWaitTime) for a block to be ready to be built,
 // and if no block is ready by then, it returns the decision to transition epoch without building a block.
 // Otherwise, it returns the decision to build a block and transition epoch along the way.
-func (bbd *blockBuildingDecider) buildBlockWithEpochTransition(ctx context.Context, pChainHeight uint64) (blockBuildingDecision, error) {
+func (bbd *blockBuildingDecider) buildBlockWithEpochTransition(ctx context.Context, pChainHeight uint64, validatorSet NodeBLSMappings) (blockBuildingDecision, error) {
 	impatientContext, cancel := context.WithTimeout(ctx, bbd.maxBlockBuildingWaitTime)
 	defer cancel()
 
@@ -124,9 +125,9 @@ func (bbd *blockBuildingDecider) buildBlockWithEpochTransition(ctx context.Conte
 	if impatientContext.Err() != nil {
 		// We have returned from waitForPendingBlock because impatientContext has timed out,
 		// which means we don't need to build a block.
-		return blockBuildingDecision{transitionEpoch: true, pChainHeight: pChainHeight}, nil
+		return blockBuildingDecision{transitionEpoch: true, pChainHeight: pChainHeight, validatorSet: validatorSet}, nil
 	}
 
 	// Block is ready to be built
-	return blockBuildingDecision{buildInnerBlock: true, transitionEpoch: true, pChainHeight: pChainHeight}, nil
+	return blockBuildingDecision{buildInnerBlock: true, transitionEpoch: true, pChainHeight: pChainHeight, validatorSet: validatorSet}, nil
 }