Addressing PR review comments

Author: batconjurer

icm-contracts/tests/flows/services/merkle_updater.go

@@ -63,9 +63,7 @@ func MerkleUpdater(
 
 	ethClient := ethereumNetwork.EthClient
 	_, ethFundedKey := ethereumNetwork.GetFundedAccountInfo()
-	chainID := ethereumNetwork.ChainID
 	fundedAddress, fundedKey := avalancheNetwork.GetFundedAccountInfo()
-	_ = chainID
 
 	// =========================================================================
 	// Setup: Fetch P-chain validators and compute the genesis Merkle root
@@ -144,7 +142,7 @@ func MerkleUpdater(
 
 	initialCommitment, err := contract.GetValidatorSetCommitment(callOpts, l1BlockchainID)
 	Expect(err).Should(BeNil())
-	Expect(initialCommitment.TotalWeight).Should(Equal(uint64(0)),
+	Expect(initialCommitment.TotalWeight).Should(BeZero(),
 		"L1 commitment should start empty")
 
 	log.Info("Contract state verified: P-chain initialized, L1 not yet registered")
@@ -193,37 +191,22 @@ func MerkleUpdater(
 	// Wait for the relayer to register the initial L1 validator set.
 	// totalWeight == 0 always triggers an immediate registration.
 	// =========================================================================
-	pollCtx, pollCancel := context.WithTimeout(ctx, 120*time.Second)
-	defer pollCancel()
-
-	ticker := time.NewTicker(2 * time.Second)
-	defer ticker.Stop()
-
-	var firstTotalWeight uint64
-	var firstPChainHeight uint64
-
-	for done := false; !done; {
-		select {
-		case <-pollCtx.Done():
-			Expect(pollCtx.Err()).ShouldNot(HaveOccurred(),
-				"Timed out waiting for relayer to register L1 validator set")
-		case <-ticker.C:
-			cmt, err := contract.GetValidatorSetCommitment(callOpts, l1BlockchainID)
-			if err != nil {
-				log.Warn("Failed to query on-chain commitment", zap.Error(err))
-				continue
-			}
+	firstCmt := pollForCommitmentUpdate(
+		ctx, log, contract, callOpts, l1BlockchainID, 120*time.Second,
+		"Timed out waiting for relayer to register L1 validator set",
+		func(cmt merkleregistry.ValidatorSetMerkleCommitment) bool {
 			if cmt.TotalWeight == 0 {
 				log.Debug("Validator set not yet registered, waiting...")
-				continue
+				return false
 			}
-
+			return true
+		},
+		func(cmt merkleregistry.ValidatorSetMerkleCommitment) {
 			log.Info("Initial validator set registered",
 				zap.Uint64("totalWeight", cmt.TotalWeight),
 				zap.Uint64("pChainHeight", cmt.PChainHeight),
 				zap.String("root", hex.EncodeToString(cmt.Root[:])),
 			)
-
 			Expect(cmt.PChainHeight).Should(BeNumerically(">", 0),
 				"P-chain height should be positive")
 			Expect(cmt.PChainTimestamp).Should(BeNumerically(">", 0),
@@ -246,17 +229,13 @@ func MerkleUpdater(
 				"Merkle root should match the P-chain validator set at the recorded height")
 			Expect(cmt.TotalWeight).Should(Equal(expectedWeight),
 				"Total weight should match the sum of validator weights")
-
-			firstTotalWeight = cmt.TotalWeight
-			firstPChainHeight = cmt.PChainHeight
-			done = true
-		}
-	}
+		},
+	)
 
 	firstRegistrationTime := time.Now()
 	log.Info("Initial registration complete",
-		zap.Uint64("firstTotalWeight", firstTotalWeight),
-		zap.Uint64("firstPChainHeight", firstPChainHeight),
+		zap.Uint64("firstTotalWeight", firstCmt.TotalWeight),
+		zap.Uint64("firstPChainHeight", firstCmt.PChainHeight),
 	)
 
 	// =========================================================================
@@ -314,43 +293,26 @@ func MerkleUpdater(
 
 	// 90s timeout — well under the staleness cap. If the update arrives in this window
 	// it was triggered by the root change, not by staleness.
-	rootChangeCtx, rootChangeCancel := context.WithTimeout(ctx, 90*time.Second)
-	defer rootChangeCancel()
-
-	rootChangeTicker := time.NewTicker(2 * time.Second)
-	defer rootChangeTicker.Stop()
-
-	var secondPChainHeight uint64
-	var secondTotalWeight uint64
-	var secondRoot [32]byte
-
-	for {
-		select {
-		case <-rootChangeCtx.Done():
-			Expect(false).Should(BeTrue(),
-				"Phase 1: Timed out waiting for root-change-triggered update")
-		case <-rootChangeTicker.C:
-			cmt, err := contract.GetValidatorSetCommitment(callOpts, l1BlockchainID)
-			if err != nil {
-				log.Warn("Failed to query on-chain commitment", zap.Error(err))
-				continue
-			}
-
-			if cmt.PChainHeight <= firstPChainHeight {
+	secondCmt := pollForCommitmentUpdate(
+		ctx, log, contract, callOpts, l1BlockchainID, 90*time.Second,
+		"Phase 1: Timed out waiting for root-change-triggered update",
+		func(cmt merkleregistry.ValidatorSetMerkleCommitment) bool {
+			if cmt.PChainHeight <= firstCmt.PChainHeight {
 				log.Debug("Phase 1: Waiting for root-change update...",
 					zap.Uint64("currentPChainHeight", cmt.PChainHeight),
 				)
-				continue
+				return false
 			}
-
+			return true
+		},
+		func(cmt merkleregistry.ValidatorSetMerkleCommitment) {
 			log.Info("Phase 1: Root-change update detected!",
 				zap.Uint64("totalWeight", cmt.TotalWeight),
 				zap.Uint64("pChainHeight", cmt.PChainHeight),
 				zap.String("root", hex.EncodeToString(cmt.Root[:])),
 			)
-
-			Expect(cmt.PChainHeight).Should(BeNumerically(">", firstPChainHeight))
-			Expect(cmt.TotalWeight).Should(BeNumerically(">", firstTotalWeight),
+			Expect(cmt.PChainHeight).Should(BeNumerically(">", firstCmt.PChainHeight))
+			Expect(cmt.TotalWeight).Should(BeNumerically(">", firstCmt.TotalWeight),
 				"Phase 1: New commitment should include the added validator's weight")
 			Expect(cmt.Root).ShouldNot(Equal([32]byte{}))
 
@@ -361,18 +323,13 @@ func MerkleUpdater(
 			expectedRoot := validatorupdater.BuildMerkleRoot(expectedValidators)
 			Expect(cmt.Root).Should(Equal(expectedRoot),
 				"Phase 1: Merkle root should match the updated P-chain validator set")
-
-			secondPChainHeight = cmt.PChainHeight
-			secondTotalWeight = cmt.TotalWeight
-			secondRoot = cmt.Root
-		}
-		break
-	}
+		},
+	)
 
 	secondUpdateTime := time.Now()
 	log.Info("Phase 1 PASSED: Root-change-triggered update confirmed",
-		zap.Uint64("secondTotalWeight", secondTotalWeight),
-		zap.Uint64("secondPChainHeight", secondPChainHeight),
+		zap.Uint64("secondTotalWeight", secondCmt.TotalWeight),
+		zap.Uint64("secondPChainHeight", secondCmt.PChainHeight),
 	)
 
 	// =========================================================================
@@ -400,49 +357,73 @@ func MerkleUpdater(
 		zap.Duration("remainingWait", remainingWait),
 	)
 
-	stalenessCtx, stalenessCancel := context.WithTimeout(ctx, remainingWait)
-	defer stalenessCancel()
-
-	stalenessTicker := time.NewTicker(2 * time.Second)
-	defer stalenessTicker.Stop()
-
-	for {
-		select {
-		case <-stalenessCtx.Done():
-			Expect(false).Should(BeTrue(),
-				"Phase 2: Timed out waiting for staleness-forced update")
-		case <-stalenessTicker.C:
-			cmt, err := contract.GetValidatorSetCommitment(callOpts, l1BlockchainID)
-			if err != nil {
-				log.Warn("Failed to query on-chain commitment", zap.Error(err))
-				continue
-			}
-
-			if cmt.PChainHeight <= secondPChainHeight {
+	pollForCommitmentUpdate(
+		ctx, log, contract, callOpts, l1BlockchainID, remainingWait,
+		"Phase 2: Timed out waiting for staleness-forced update",
+		func(cmt merkleregistry.ValidatorSetMerkleCommitment) bool {
+			if cmt.PChainHeight <= secondCmt.PChainHeight {
 				log.Debug("Phase 2: Still waiting for staleness update...",
 					zap.Duration("elapsed", time.Since(secondUpdateTime)),
 				)
-				continue
+				return false
 			}
-
+			return true
+		},
+		func(cmt merkleregistry.ValidatorSetMerkleCommitment) {
 			log.Info("Phase 2: Staleness-forced update detected!",
 				zap.Uint64("totalWeight", cmt.TotalWeight),
 				zap.Uint64("pChainHeight", cmt.PChainHeight),
 				zap.String("root", hex.EncodeToString(cmt.Root[:])),
 			)
-
-			Expect(cmt.PChainHeight).Should(BeNumerically(">", secondPChainHeight))
-			Expect(cmt.TotalWeight).Should(Equal(secondTotalWeight),
+			Expect(cmt.PChainHeight).Should(BeNumerically(">", secondCmt.PChainHeight))
+			Expect(cmt.TotalWeight).Should(Equal(secondCmt.TotalWeight),
 				"Phase 2: Total weight should be unchanged (no new validators)")
-			Expect(cmt.Root).Should(Equal(secondRoot),
+			Expect(cmt.Root).Should(Equal(secondCmt.Root),
 				"Phase 2: Merkle root should be unchanged (no new validators)")
 
 			log.Info("Phase 2 PASSED: Staleness-forced update confirmed")
+		},
+	)
+
+	log.Info("MerkleUpdater e2e test PASSED")
+}
+
+func pollForCommitmentUpdate(
+	ctx context.Context,
+	log logging.Logger,
+	contract *merkleregistry.MerkleValidatorSetRegistry,
+	callOpts *bind.CallOpts,
+	l1BlockchainID ids.ID,
+	timeout time.Duration,
+	timeoutMsg string,
+	isUpdated func(cmt merkleregistry.ValidatorSetMerkleCommitment) bool,
+	validate func(cmt merkleregistry.ValidatorSetMerkleCommitment),
+) merkleregistry.ValidatorSetMerkleCommitment {
+	pollCtx, pollCancel := context.WithTimeout(ctx, timeout)
+	defer pollCancel()
+
+	ticker := time.NewTicker(2 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-pollCtx.Done():
+			Expect(false).Should(BeTrue(), timeoutMsg)
+		case <-ticker.C:
+			cmt, err := contract.GetValidatorSetCommitment(callOpts, l1BlockchainID)
+			if err != nil {
+				log.Warn("Failed to query on-chain commitment", zap.Error(err))
+				continue
+			}
+			if !isUpdated(cmt) {
+				continue
+			}
+			validate(cmt)
+			return cmt
 		}
 		break
 	}
-
-	log.Info("MerkleUpdater e2e test PASSED")
+	panic("unreachable")
 }
 
 func createMerkleUpdaterRelayerConfig(

peers/clients/validator_client.go

@@ -32,6 +32,7 @@ type CanonicalValidatorState interface {
 	GetAllValidatorSets(ctx context.Context, pchainHeight uint64) (map[ids.ID]validators.WarpSet, error)
 	GetProposedValidators(ctx context.Context, subnetID ids.ID) (validators.WarpSet, error)
 	GetCurrentValidators(ctx context.Context, subnetID ids.ID) ([]platformvm.ClientPermissionlessValidator, error)
+	GetValidatorsAt(ctx context.Context, subnetID ids.ID, pchainHeight uint64) (map[ids.NodeID]*validators.GetValidatorOutput, error)
 }
 
 // CanonicalValidatorClient wraps [platformvm.Client] and implements [CanonicalValidatorState]
@@ -105,6 +106,10 @@ func (v *CanonicalValidatorClient) GetProposedValidators(
 	return validators.FlattenValidatorSet(res)
 }
 
+func (v *CanonicalValidatorClient) GetValidatorsAt(ctx context.Context, subnetID ids.ID, pchainHeight uint64) (map[ids.NodeID]*validators.GetValidatorOutput, error) {
+	return v.client.GetValidatorsAt(ctx, subnetID, pchainapi.Height(pchainHeight), v.options...)
+}
+
 // Gets the validator set of the given subnet at the given P-chain block height.
 // Uses [platform.getValidatorsAt] with supplied height
 func (v *CanonicalValidatorClient) GetAllValidatorSets(

relayer/main/main.go

@@ -849,7 +849,9 @@ func startMerkleSetUpdater(
 
 	pollInterval := time.Duration(extDest.PollIntervalSeconds) * time.Second
 	maxUpdateInterval := time.Duration(extDest.MaxUpdateIntervalSeconds) * time.Second
-
+	if pollInterval == 0 {
+		pollInterval = validatorupdater.DefaultPollInterval
+	}
 	updater := validatorupdater.NewMerkleSetUpdater(
 		logger,
 		pChainClient,

relayer/validatorupdater/diff_set_updater.go

@@ -73,7 +73,7 @@ func NewDiffSetUpdater(
 		shardSize = defaultShardSize
 	}
 	if pollInterval == 0 {
-		pollInterval = defaultPollInterval
+		pollInterval = DefaultPollInterval
 	}
 	return &DiffSetUpdater{
 		logger:                   logger,

relayer/validatorupdater/merkle_validator_updater.go

@@ -61,9 +61,6 @@ func NewMerkleSetUpdater(
 	pollInterval time.Duration,
 	maxUpdateInterval time.Duration,
 ) *MerkleSetUpdater {
-	if pollInterval == 0 {
-		pollInterval = defaultPollInterval
-	}
 	return &MerkleSetUpdater{
 		logger:              logger,
 		pChainClient:        pChainClient,
@@ -152,7 +149,7 @@ func (s *MerkleSetUpdater) checkAndUpdate(ctx context.Context) error {
 	if err := s.performUpdate(ctx, s.localPChainHeight, validatorSetUpdate, false); err != nil {
 		s.logger.Warn("Merkle root update failed, re-syncing from contract on next tick", zap.Error(err))
 		s.initialized = false
-		return err
+		return nil
 	}
 
 	s.localValidatorSet = newValidators
@@ -289,22 +286,28 @@ func (s *MerkleSetUpdater) performUpdate(
 		return fmt.Errorf("failed to sign message: %w", err)
 	}
 
-	return s.sendUpdate(ctx, signedMsg, isFirstRegistration)
+	return s.sendUpdate(ctx, signedMsg, onChainPChainHeight, isFirstRegistration)
 }
 
 func (s *MerkleSetUpdater) sendUpdate(
 	ctx context.Context,
 	signedMsg *avalancheWarp.Message,
+	onChainPChainHeight uint64,
 	isFirstRegistration bool,
 ) error {
 	var attestationValidators []*Validator
 	if isFirstRegistration {
 		// First registration is verified against the P-chain Merkle root stored in the
 		// contract constructor. Fetch the primary network validators used to build that
 		// root so the attestation proof is computed against the correct ordered set.
-		pChainWarpSet, err := s.pChainClient.GetProposedValidators(ctx, ids.Empty)
+		allValidatorSets, err := s.pChainClient.GetAllValidatorSets(ctx, onChainPChainHeight)
 		if err != nil {
-			return fmt.Errorf("failed to get P-chain validators for attestation: %w", err)
+			return fmt.Errorf("failed to get P-chain validator sets at height %d for attestation: %w",
+				onChainPChainHeight, err)
+		}
+		pChainWarpSet, ok := allValidatorSets[ids.Empty]
+		if !ok {
+			return fmt.Errorf("primary network not found in validator sets at height %d", onChainPChainHeight)
 		}
 		attestationValidators = make([]*Validator, len(pChainWarpSet.Validators))
 		for i, vdr := range pChainWarpSet.Validators {
@@ -364,22 +367,19 @@ func (s *MerkleSetUpdater) fetchSortedValidators(
 	ctx context.Context,
 	pChainHeight uint64,
 ) ([]*Validator, error) {
-	allValidatorSets, err := s.pChainClient.GetAllValidatorSets(ctx, pChainHeight)
+	subnetValidatorSet, err := s.pChainClient.GetValidatorsAt(ctx, s.subnetID, pChainHeight)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get validator sets: %w", err)
 	}
 
-	vdrSet, ok := allValidatorSets[s.subnetID]
-	if !ok {
-		return nil, fmt.Errorf("subnet %s not found in validator sets at height %d", s.subnetID, pChainHeight)
-	}
-
-	validators := make([]*Validator, len(vdrSet.Validators))
-	for i, vdr := range vdrSet.Validators {
-		validators[i] = &Validator{
+	validators := make([]*Validator, len(subnetValidatorSet))
+	ix := 0
+	for _, vdr := range subnetValidatorSet {
+		validators[ix] = &Validator{
 			UncompressedPublicKeyBytes: [96]byte(vdr.PublicKey.Serialize()),
 			Weight:                     vdr.Weight,
 		}
+		ix++
 	}
 	sort.Slice(validators, func(i, j int) bool {
 		return string(validators[i].UncompressedPublicKeyBytes[:]) < string(validators[j].UncompressedPublicKeyBytes[:])