From 63dd0354bcfa07d7cee3d2f92e9e6cafb18828d7 Mon Sep 17 00:00:00 2001 From: jait91 Date: Fri, 29 May 2026 13:44:53 +0300 Subject: [PATCH 1/3] fix: prevent SMT corruption during HA leadership churn Storage pressure could cause leadership churn while old child-round goroutines were still waiting for parent proofs. A stale finalizer could then write the current round's pending SMT data under an older block, causing blocks, block_records, aggregator_records, and smt_nodes to diverge. Fix by cancelling and waiting for in-flight round processing on Deactivate, refusing to overwrite a same-or-newer active round on reactivation, and validating round number, snapshot root, and duplicate-block contents before finalization continues. --- internal/round/batch_processor.go | 130 ++++++++++++++++-- internal/round/finalize_duplicate_test.go | 158 ++++++++++++++++++++++ internal/round/precollection_test.go | 128 ++++++++++++++---- internal/round/round_manager.go | 57 +++++++- 4 files changed, 431 insertions(+), 42 deletions(-) diff --git a/internal/round/batch_processor.go b/internal/round/batch_processor.go index e3598cbf..d049ff59 100644 --- a/internal/round/batch_processor.go +++ b/internal/round/batch_processor.go @@ -318,6 +318,14 @@ func (rm *RoundManager) pollForParentProof(ctx context.Context, rootHash string) // ErrParentProofPollTimeout marks a single poll window timeout while waiting for a parent proof. var ErrParentProofPollTimeout = errors.New("parent shard inclusion proof poll timeout") +const ( + finalizeRejectNoActiveRound = "no_active_round" + finalizeRejectRoundMismatch = "round_mismatch" + finalizeRejectSnapshotMissing = "snapshot_missing" + finalizeRejectSnapshotRoot = "root_mismatch" + finalizeRejectDuplicateMismatch = "duplicate_mismatch" +) + func (rm *RoundManager) submitShardRootWithRetry(ctx context.Context, req *api.SubmitShardRootRequest) error { if rm.rootClient == nil { return fmt.Errorf("root client not configured") @@ -406,20 +414,50 @@ func (rm *RoundManager) FinalizeBlock(ctx context.Context, block *models.Block) commitmentCount := 0 rm.roundMutex.Lock() - if rm.currentRound != nil && rm.currentRound.Number.String() == block.Index.String() { - proposalTime = rm.currentRound.ProposalTime - processingTime = rm.currentRound.ProcessingTime + round := rm.currentRound + if round == nil { + rm.roundMutex.Unlock() + rm.logFinalizationRejected(ctx, finalizeRejectNoActiveRound, "Rejected block finalization: no active round", + "blockNumber", block.Index.String(), + "rootHash", block.RootHash.String()) + return fmt.Errorf("cannot finalize block %s: no active round", block.Index.String()) + } + if round.Number == nil || round.Number.String() != block.Index.String() { + current := "" + if round.Number != nil { + current = round.Number.String() + } + rm.roundMutex.Unlock() + rm.logFinalizationRejected(ctx, finalizeRejectRoundMismatch, "Rejected block finalization: active round mismatch", + "blockNumber", block.Index.String(), + "activeRound", current, + "rootHash", block.RootHash.String()) + return fmt.Errorf("cannot finalize block %s: active round is %s", block.Index.String(), current) } - rm.roundMutex.Unlock() - rm.roundMutex.Lock() var pendingLeaves []*smt.Leaf var pendingCommitments []*models.CertificationRequest var snapshot *smt.ThreadSafeSmtSnapshot - if rm.currentRound != nil { - pendingLeaves = rm.currentRound.PendingLeaves - pendingCommitments = rm.currentRound.PendingCommitments - snapshot = rm.currentRound.Snapshot + proposalTime = round.ProposalTime + processingTime = round.ProcessingTime + pendingLeaves = append(pendingLeaves, round.PendingLeaves...) + pendingCommitments = append(pendingCommitments, round.PendingCommitments...) + snapshot = round.Snapshot + if snapshot == nil { + rm.roundMutex.Unlock() + rm.logFinalizationRejected(ctx, finalizeRejectSnapshotMissing, "Rejected block finalization: active round has no SMT snapshot", + "blockNumber", block.Index.String(), + "rootHash", block.RootHash.String()) + return fmt.Errorf("cannot finalize block %s: active round has no SMT snapshot", block.Index.String()) + } + if snapshotRoot := snapshot.GetRootHash(); snapshotRoot != block.RootHash.String() { + rm.roundMutex.Unlock() + rm.logFinalizationRejected(ctx, finalizeRejectSnapshotRoot, "Rejected block finalization: snapshot root mismatch", + "blockNumber", block.Index.String(), + "snapshotRoot", snapshotRoot, + "blockRoot", block.RootHash.String()) + return fmt.Errorf("cannot finalize block %s: snapshot root %s does not match block root %s", + block.Index.String(), snapshotRoot, block.RootHash.String()) } rm.roundMutex.Unlock() @@ -451,7 +489,14 @@ func (rm *RoundManager) FinalizeBlock(ctx context.Context, block *models.Block) if !errors.Is(err, interfaces.ErrDuplicateKey) { return fmt.Errorf("failed to store block and records: %w", err) } - rm.logger.WithContext(ctx).Info("Block already exists, continuing with remaining steps", + if err := rm.validateDuplicateBlock(ctx, block, requestIds); err != nil { + rm.logFinalizationRejected(ctx, finalizeRejectDuplicateMismatch, "Rejected duplicate block finalization", + "blockNumber", block.Index.String(), + "rootHash", block.RootHash.String(), + "error", err.Error()) + return err + } + rm.logger.WithContext(ctx).Info("Block already exists with matching root and records, continuing with remaining steps", "blockNumber", block.Index.String()) } @@ -484,7 +529,7 @@ func (rm *RoundManager) FinalizeBlock(ctx context.Context, block *models.Block) } rm.roundMutex.Lock() - if rm.currentRound != nil { + if rm.currentRound == round { rm.currentRound.Block = block rm.currentRound.PendingRootHash = "" rm.currentRound.PendingLeaves = nil @@ -590,6 +635,11 @@ func (rm *RoundManager) FinalizeBlock(ctx context.Context, block *models.Block) return nil } +func (rm *RoundManager) logFinalizationRejected(ctx context.Context, reason, message string, fields ...interface{}) { + logFields := append([]interface{}{"reason", reason}, fields...) + rm.logger.WithContext(ctx).Warn(message, logFields...) +} + // convertLeavesToNodes converts SMT leaves to storage models func (rm *RoundManager) convertLeavesToNodes(leaves []*smt.Leaf) []*models.SmtNode { if len(leaves) == 0 { @@ -635,6 +685,64 @@ func (rm *RoundManager) storeBlockAndRecords(ctx context.Context, block *models. }) } +func (rm *RoundManager) validateDuplicateBlock(ctx context.Context, block *models.Block, requestIds []api.StateID) error { + existingBlock, err := rm.getExistingBlockAnyFinalization(ctx, block.Index) + if err != nil { + return fmt.Errorf("failed to load existing block %s after duplicate insert: %w", block.Index.String(), err) + } + if existingBlock == nil { + return fmt.Errorf("duplicate block %s reported but existing block was not found", block.Index.String()) + } + if existingBlock.RootHash.String() != block.RootHash.String() { + return fmt.Errorf("duplicate block %s has root %s, attempted root %s", + block.Index.String(), existingBlock.RootHash.String(), block.RootHash.String()) + } + + existingRecords, err := rm.storage.BlockRecordsStorage().GetByBlockNumber(ctx, block.Index) + if err != nil { + return fmt.Errorf("failed to load existing block records for block %s: %w", block.Index.String(), err) + } + if existingRecords == nil { + return fmt.Errorf("duplicate block %s exists without block records", block.Index.String()) + } + if !sameStateIDs(existingRecords.StateIDs, requestIds) { + return fmt.Errorf("duplicate block %s has different request IDs: existing=%d attempted=%d", + block.Index.String(), len(existingRecords.StateIDs), len(requestIds)) + } + + return nil +} + +func (rm *RoundManager) getExistingBlockAnyFinalization(ctx context.Context, blockNumber *api.BigInt) (*models.Block, error) { + existingBlock, err := rm.storage.BlockStorage().GetByNumber(ctx, blockNumber) + if err != nil || existingBlock != nil { + return existingBlock, err + } + + unfinalizedBlocks, err := rm.storage.BlockStorage().GetUnfinalized(ctx) + if err != nil { + return nil, err + } + for _, candidate := range unfinalizedBlocks { + if candidate.Index.String() == blockNumber.String() { + return candidate, nil + } + } + return nil, nil +} + +func sameStateIDs(a, b []api.StateID) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].String() != b[i].String() { + return false + } + } + return true +} + // storeDataParallel stores SMT nodes and aggregator records in parallel. // StoreBatch handles duplicates internally (ignores duplicate key errors). func (rm *RoundManager) storeDataParallel( diff --git a/internal/round/finalize_duplicate_test.go b/internal/round/finalize_duplicate_test.go index 2ff7c6ad..2b320fb8 100644 --- a/internal/round/finalize_duplicate_test.go +++ b/internal/round/finalize_duplicate_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" @@ -507,3 +508,160 @@ func (s *FinalizeDuplicateTestSuite) Test6_BlockRecordsMatchPendingCommitmentsOn require.NoError(t, err) require.Equal(t, recordCountBefore+2, recordCountAfter, "should persist only filtered aggregator records") } + +// Test7_FinalizeBlockRejectsRoundNumberMismatch verifies that a finalized block +// cannot consume pending leaves from a different currentRound. This is the +// corruption shape where a stale finalization writes newer-round commitments +// under an older block number. +func (s *FinalizeDuplicateTestSuite) Test7_FinalizeBlockRejectsRoundNumberMismatch() { + t := s.T() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + testLogger, err := logger.New("info", "text", "stdout", false) + require.NoError(t, err) + + threadSafeSMT := smt.NewThreadSafeSMT(smt.NewSparseMerkleTree(api.SHA256, 16+256)) + rm, err := NewRoundManager(ctx, s.cfg, testLogger, s.storage.CommitmentQueue(), s.storage, nil, state.NewSyncStateTracker(), nil, events.NewEventBus(testLogger), threadSafeSMT, nil) + require.NoError(t, err) + + commitments := testutil.CreateTestCertificationRequests(t, 3, "t7_req") + + rm.currentRound = &Round{ + Number: api.NewBigInt(big.NewInt(70)), + State: RoundStateProcessing, + Commitments: commitments, + Snapshot: rm.smt.CreateSnapshot(), + } + + rm.roundMutex.Lock() + err = rm.processMiniBatch(ctx, commitments) + rm.roundMutex.Unlock() + require.NoError(t, err) + + rootHash := rm.currentRound.Snapshot.GetRootHash() + rootHashBytes, err := api.NewHexBytesFromString(rootHash) + require.NoError(t, err) + + block := models.NewBlock( + api.NewBigInt(big.NewInt(69)), + "unicity", + 0, + "1.0", + "mainnet", + rootHashBytes, + api.HexBytes{}, + api.HexBytes{}, + nil, + ) + + blockCountBefore, err := s.storage.BlockStorage().Count(ctx) + require.NoError(t, err) + blockRecordsCountBefore, err := s.storage.BlockRecordsStorage().Count(ctx) + require.NoError(t, err) + smtCountBefore, err := s.storage.SmtStorage().Count(ctx) + require.NoError(t, err) + recordCountBefore, err := s.storage.AggregatorRecordStorage().Count(ctx) + require.NoError(t, err) + + err = rm.FinalizeBlock(ctx, block) + assert.Error(t, err, "FinalizeBlock should reject block 69 while currentRound is 70") + + blockCountAfter, countErr := s.storage.BlockStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, blockCountBefore, blockCountAfter, "stale finalization must not store a block") + + blockRecordsCountAfter, countErr := s.storage.BlockRecordsStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, blockRecordsCountBefore, blockRecordsCountAfter, "stale finalization must not store block_records") + + smtCountAfter, countErr := s.storage.SmtStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, smtCountBefore, smtCountAfter, "stale finalization must not store SMT nodes") + + recordCountAfter, countErr := s.storage.AggregatorRecordStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, recordCountBefore, recordCountAfter, "stale finalization must not store aggregator records") + + rm.roundMutex.RLock() + defer rm.roundMutex.RUnlock() + if assert.NotNil(t, rm.currentRound, "stale finalization must not clear the active currentRound") { + assert.Equal(t, "70", rm.currentRound.Number.String()) + assert.Len(t, rm.currentRound.PendingCommitments, 3) + assert.NotNil(t, rm.currentRound.Snapshot) + } +} + +// Test8_DuplicateFinalizedBlockMustNotCommitDifferentSnapshot verifies that a +// retry for an already-finalized block is only idempotent if it is replaying the +// same block state. A duplicate block must not commit whatever snapshot happens +// to be in currentRound. +func (s *FinalizeDuplicateTestSuite) Test8_DuplicateFinalizedBlockMustNotCommitDifferentSnapshot() { + t := s.T() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + testLogger, err := logger.New("info", "text", "stdout", false) + require.NoError(t, err) + + threadSafeSMT := smt.NewThreadSafeSMT(smt.NewSparseMerkleTree(api.SHA256, 16+256)) + rm, err := NewRoundManager(ctx, s.cfg, testLogger, s.storage.CommitmentQueue(), s.storage, nil, state.NewSyncStateTracker(), nil, events.NewEventBus(testLogger), threadSafeSMT, nil) + require.NoError(t, err) + + emptyRoot := rm.smt.GetRootHash() + emptyRootBytes, err := api.NewHexBytesFromString(emptyRoot) + require.NoError(t, err) + + block := models.NewBlock( + api.NewBigInt(big.NewInt(71)), + "unicity", + 0, + "1.0", + "mainnet", + emptyRootBytes, + api.HexBytes{}, + api.HexBytes{}, + nil, + ) + block.Finalized = true + require.NoError(t, s.storage.BlockStorage().Store(ctx, block)) + require.NoError(t, s.storage.BlockRecordsStorage().Store(ctx, models.NewBlockRecords(block.Index, nil))) + + commitments := testutil.CreateTestCertificationRequests(t, 2, "t8_req") + rm.currentRound = &Round{ + Number: api.NewBigInt(big.NewInt(71)), + State: RoundStateProcessing, + Commitments: commitments, + Snapshot: rm.smt.CreateSnapshot(), + } + + rm.roundMutex.Lock() + err = rm.processMiniBatch(ctx, commitments) + rm.roundMutex.Unlock() + require.NoError(t, err) + require.NotEqual(t, emptyRoot, rm.currentRound.Snapshot.GetRootHash(), "test setup must create a different pending root") + + smtCountBefore, err := s.storage.SmtStorage().Count(ctx) + require.NoError(t, err) + recordCountBefore, err := s.storage.AggregatorRecordStorage().Count(ctx) + require.NoError(t, err) + + block.Finalized = false + err = rm.FinalizeBlock(ctx, block) + assert.Error(t, err, "FinalizeBlock should reject duplicate block 71 when current snapshot root differs from the stored block root") + + assert.Equal(t, emptyRoot, rm.smt.GetRootHash(), "duplicate finalized block must not commit a different currentRound snapshot") + + smtCountAfter, countErr := s.storage.SmtStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, smtCountBefore, smtCountAfter, "duplicate finalized block must not store unrelated SMT nodes") + + recordCountAfter, countErr := s.storage.AggregatorRecordStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, recordCountBefore, recordCountAfter, "duplicate finalized block must not store unrelated aggregator records") + + storedBlock, err := s.storage.BlockStorage().GetByNumber(ctx, block.Index) + require.NoError(t, err) + require.NotNil(t, storedBlock) + assert.Equal(t, emptyRoot, storedBlock.RootHash.String(), "stored block root should remain unchanged") +} diff --git a/internal/round/precollection_test.go b/internal/round/precollection_test.go index 52967b38..d0faadb4 100644 --- a/internal/round/precollection_test.go +++ b/internal/round/precollection_test.go @@ -31,10 +31,12 @@ type signalingRootAggregatorClient struct { } type blockingProofRootAggregatorClient struct { - inner *testsharding.RootAggregatorClientStub - proofPollStarted chan struct{} - releaseProof chan struct{} - startOnce sync.Once + inner *testsharding.RootAggregatorClientStub + proofPollStarted chan struct{} + proofPollCanceled chan struct{} + releaseProof chan struct{} + startOnce sync.Once + cancelOnce sync.Once } type countingProofRootAggregatorClient struct { @@ -66,9 +68,10 @@ func newSignalingRootAggregatorClient() *signalingRootAggregatorClient { func newBlockingProofRootAggregatorClient() *blockingProofRootAggregatorClient { return &blockingProofRootAggregatorClient{ - inner: testsharding.NewRootAggregatorClientStub(), - proofPollStarted: make(chan struct{}), - releaseProof: make(chan struct{}), + inner: testsharding.NewRootAggregatorClientStub(), + proofPollStarted: make(chan struct{}), + proofPollCanceled: make(chan struct{}), + releaseProof: make(chan struct{}), } } @@ -120,6 +123,9 @@ func (c *blockingProofRootAggregatorClient) GetShardProof(ctx context.Context, r select { case <-c.releaseProof: case <-ctx.Done(): + c.cancelOnce.Do(func() { + close(c.proofPollCanceled) + }) return nil, ctx.Err() } @@ -607,11 +613,9 @@ func TestPreCollectionReparenting(t *testing.T) { }) } -// TestChildPrecollector_DeactivateDuringInFlightRound deactivates while the -// first child round is blocked waiting for the parent proof. The in-flight -// round should still finish block 1 once the proof is released, but it must -// not start round 2. -func TestChildPrecollector_DeactivateDuringInFlightRound(t *testing.T) { +// TestChildPrecollector_DeactivateCancelsInFlightRound verifies leadership +// loss cancels parent-proof polling before the stale round can finalize. +func TestChildPrecollector_DeactivateCancelsInFlightRound(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -662,26 +666,102 @@ func TestChildPrecollector_DeactivateDuringInFlightRound(t *testing.T) { t.Fatal("timed out waiting for round 1 to enter proof polling") } - // Deactivate while round 1 is still in-flight. require.NoError(t, rm.Deactivate(ctx)) - // Let the in-flight round finish block 1. + select { + case <-rootClient.proofPollCanceled: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for in-flight proof polling to be cancelled") + } + close(rootClient.releaseProof) - require.Eventually(t, func() bool { - block, err := store.BlockStorage().GetByNumber(ctx, api.NewBigInt(big.NewInt(1))) - return err == nil && block != nil - }, 3*time.Second, 25*time.Millisecond, "block 1 should be created after proof release") + block1, err := store.BlockStorage().GetByNumber(ctx, api.NewBigInt(big.NewInt(1))) + require.NoError(t, err) + assert.Nil(t, block1, "deactivated in-flight round must not finalize block 1") +} - // If the bug exists, round 2 would appear shortly after block 1 finalization. - time.Sleep(500 * time.Millisecond) +func TestChildRound_ReactivateStartsFreshRoundAfterCanceledInFlight(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cfg := config.Config{ + Database: config.DatabaseConfig{ + Database: "test_child_reactivate_inflight", + }, + Processing: config.ProcessingConfig{ + RoundDuration: 100 * time.Millisecond, + MaxCommitmentsPerRound: 1000, + }, + HA: config.HAConfig{ + Enabled: true, + }, + Sharding: config.ShardingConfig{ + Mode: config.ShardingModeChild, + Child: config.ChildConfig{ + ShardID: 0b11, + ParentPollTimeout: 5 * time.Second, + ParentPollInterval: 10 * time.Millisecond, + }, + }, + } + store := testutil.SetupTestStorage(t, cfg) + + testLogger := newTestLogger(t) + rootClient := newBlockingProofRootAggregatorClient() + smtInstance := smt.NewThreadSafeSMT(smt.NewSparseMerkleTree(api.SHA256, 16+256)) - block2, err := store.BlockStorage().GetByNumber(ctx, api.NewBigInt(big.NewInt(2))) + rm, err := NewRoundManager( + ctx, + &cfg, + testLogger, + store.CommitmentQueue(), + store, + rootClient, + state.NewSyncStateTracker(), + nil, + events.NewEventBus(testLogger), + smtInstance, + nil, + ) require.NoError(t, err) - assert.Nil(t, block2, "no block 2 should be created after Deactivate") - // Wait for goroutines to drain - rm.wg.Wait() + releaseOnce := sync.Once{} + releaseProof := func() { + releaseOnce.Do(func() { + close(rootClient.releaseProof) + }) + } + defer func() { + releaseProof() + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 3*time.Second) + defer shutdownCancel() + _ = rm.Stop(shutdownCtx) + }() + + require.NoError(t, rm.Start(ctx)) + require.NoError(t, rm.Activate(ctx)) + + select { + case <-rootClient.proofPollStarted: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first round to enter proof polling") + } + require.Equal(t, 1, rootClient.inner.SubmissionCount(), "first activation should submit exactly one child root") + + require.NoError(t, rm.Deactivate(ctx)) + + select { + case <-rootClient.proofPollCanceled: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first proof polling to be cancelled") + } + + require.NoError(t, rm.Activate(ctx)) + + require.Eventually(t, func() bool { + return rootClient.inner.SubmissionCount() == 2 + }, 2*time.Second, 25*time.Millisecond, "reactivation should submit a fresh child round after cancelling the old one") } func TestChildRound_ParentProofTimeoutIsRetriable(t *testing.T) { diff --git a/internal/round/round_manager.go b/internal/round/round_manager.go index 5ad00a39..e715e919 100644 --- a/internal/round/round_manager.go +++ b/internal/round/round_manager.go @@ -101,6 +101,8 @@ type RoundManager struct { // Guards the window where SMT root advances before block finalization is persisted. finalizationMu sync.RWMutex wg sync.WaitGroup + roundWG sync.WaitGroup + activeCancel context.CancelFunc // Round duration (configurable, default 1 second) roundDuration time.Duration @@ -381,6 +383,16 @@ func (rm *RoundManager) StartNewRoundWithSnapshot( "previousRoundNumber", rm.currentRound.Number.String(), "previousRoundState", rm.currentRound.State.String(), "previousRoundAge", time.Since(rm.currentRound.StartTime).String()) + if rm.currentRound.Number != nil && rm.currentRound.Number.Cmp(roundNumber.Int) >= 0 { + activeRound := rm.currentRound.Number.String() + activeState := rm.currentRound.State.String() + rm.roundMutex.Unlock() + rm.logger.WithContext(ctx).Warn("Skipping round start - same or newer round is already active", + "requestedRound", roundNumber.String(), + "activeRound", activeRound, + "activeState", activeState) + return nil + } } if snapshot == nil { @@ -416,6 +428,7 @@ func (rm *RoundManager) StartNewRoundWithSnapshot( cp.Start(ctx, snapshot) } + rm.roundWG.Add(1) rm.roundMutex.Unlock() rm.logger.WithContext(ctx).Info("Started new round", @@ -423,6 +436,7 @@ func (rm *RoundManager) StartNewRoundWithSnapshot( "commitments", len(commitments)) rm.wg.Go(func() { + defer rm.roundWG.Done() if err := rm.processRound(ctx); err != nil { if errors.Is(err, context.Canceled) { rm.logger.WithContext(ctx).Info("Round processing stopped due to shutdown", @@ -731,12 +745,29 @@ func (rm *RoundManager) restoreSmtFromStorage(ctx context.Context) (*api.BigInt, func (rm *RoundManager) Activate(ctx context.Context) error { rm.logger.WithContext(ctx).Info("Activating round manager") + activeCtx, activeCancel := context.WithCancel(ctx) + activated := false + defer func() { + if !activated { + activeCancel() + rm.roundMutex.Lock() + rm.activeCancel = nil + rm.precollectorDisabled = true + rm.roundMutex.Unlock() + } + }() + rm.roundMutex.Lock() + if rm.activeCancel != nil { + rm.logger.WithContext(ctx).Warn("Activation requested while already active; cancelling previous activation") + rm.activeCancel() + } + rm.activeCancel = activeCancel rm.precollectorDisabled = false rm.roundMutex.Unlock() if rm.config.HA.Enabled { - recoveryResult, err := RecoverUnfinalizedBlock(ctx, rm.logger, rm.storage, rm.commitmentQueue) + recoveryResult, err := RecoverUnfinalizedBlock(activeCtx, rm.logger, rm.storage, rm.commitmentQueue) if err != nil { return fmt.Errorf("failed to recover unfinalized block on activation: %w", err) } @@ -745,7 +776,7 @@ func (rm *RoundManager) Activate(ctx context.Context) error { "blockNumber", recoveryResult.BlockNumber.String(), "requestCount", len(recoveryResult.RequestIDs)) - if err := LoadRecoveredNodesIntoSMT(ctx, rm.logger, rm.storage, rm.smt, recoveryResult.RequestIDs); err != nil { + if err := LoadRecoveredNodesIntoSMT(activeCtx, rm.logger, rm.storage, rm.smt, recoveryResult.RequestIDs); err != nil { return fmt.Errorf("failed to load recovered nodes into SMT: %w", err) } } @@ -753,15 +784,15 @@ func (rm *RoundManager) Activate(ctx context.Context) error { switch rm.config.Sharding.Mode { case config.ShardingModeStandalone: - if err := rm.bftClient.Start(ctx); err != nil { + if err := rm.bftClient.Start(activeCtx); err != nil { return fmt.Errorf("failed to start BFT client: %w", err) } case config.ShardingModeChild: - if err := rm.restoreLastAcceptedParentUC(ctx); err != nil { + if err := rm.restoreLastAcceptedParentUC(activeCtx); err != nil { return fmt.Errorf("failed to restore latest parent UC from child blocks: %w", err) } - latestBlockNumber, err := rm.storage.BlockStorage().GetLatestNumber(ctx) + latestBlockNumber, err := rm.storage.BlockStorage().GetLatestNumber(activeCtx) if err != nil { return fmt.Errorf("failed to get latest block number: %w", err) } @@ -777,14 +808,15 @@ func (rm *RoundManager) Activate(ctx context.Context) error { "latestBlock", latestBlockNumber, "nextRound", roundNumber.String()) - if err := rm.StartNewRound(ctx, api.NewBigInt(roundNumber)); err != nil { + if err := rm.StartNewRound(activeCtx, api.NewBigInt(roundNumber)); err != nil { return fmt.Errorf("failed to start new round: %w", err) } default: return fmt.Errorf("invalid shard mode: %s", rm.config.Sharding.Mode) } - rm.startCommitmentPrefetcher(ctx) + rm.startCommitmentPrefetcher(activeCtx) + activated = true return nil } @@ -843,13 +875,19 @@ func decodeUnicityCertificate(raw api.HexBytes) (*types.UnicityCertificate, erro func (rm *RoundManager) Deactivate(ctx context.Context) error { rm.logger.WithContext(ctx).Info("Deactivating round manager") + var activeCancel context.CancelFunc var cp *childPrecollector rm.roundMutex.Lock() rm.precollectorDisabled = true + activeCancel = rm.activeCancel + rm.activeCancel = nil cp = rm.precollector rm.precollector = nil rm.roundMutex.Unlock() + if activeCancel != nil { + activeCancel() + } if cp != nil { cp.Stop() } @@ -857,6 +895,11 @@ func (rm *RoundManager) Deactivate(ctx context.Context) error { if rm.bftClient != nil { rm.bftClient.Stop() } + rm.roundWG.Wait() + + rm.roundMutex.Lock() + rm.currentRound = nil + rm.roundMutex.Unlock() return nil } From f3b2d0fae36ffdaee24dd9522b7ca1a34c0b8bf3 Mon Sep 17 00:00:00 2001 From: jait91 Date: Fri, 5 Jun 2026 09:24:03 +0300 Subject: [PATCH 2/3] address review comments --- internal/round/batch_processor.go | 13 +++++++++++-- internal/round/round_manager.go | 22 +++++++++++++--------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/internal/round/batch_processor.go b/internal/round/batch_processor.go index d049ff59..d4f1be7d 100644 --- a/internal/round/batch_processor.go +++ b/internal/round/batch_processor.go @@ -401,6 +401,12 @@ func (rm *RoundManager) FinalizeBlockWithRetry(ctx context.Context, block *model // FinalizeBlock creates and persists a new block with the given data func (rm *RoundManager) FinalizeBlock(ctx context.Context, block *models.Block) error { + if block == nil { + return fmt.Errorf("cannot finalize nil block") + } + if block.Index == nil || block.Index.Int == nil { + return fmt.Errorf("cannot finalize block: missing block index") + } rm.logger.WithContext(ctx).Info("FinalizeBlock called", "blockNumber", block.Index.String(), "rootHash", block.RootHash.String(), @@ -422,7 +428,7 @@ func (rm *RoundManager) FinalizeBlock(ctx context.Context, block *models.Block) "rootHash", block.RootHash.String()) return fmt.Errorf("cannot finalize block %s: no active round", block.Index.String()) } - if round.Number == nil || round.Number.String() != block.Index.String() { + if round.Number == nil || round.Number.Cmp(block.Index.Int) != 0 { current := "" if round.Number != nil { current = round.Number.String() @@ -714,6 +720,9 @@ func (rm *RoundManager) validateDuplicateBlock(ctx context.Context, block *model } func (rm *RoundManager) getExistingBlockAnyFinalization(ctx context.Context, blockNumber *api.BigInt) (*models.Block, error) { + if blockNumber == nil || blockNumber.Int == nil { + return nil, fmt.Errorf("missing block number") + } existingBlock, err := rm.storage.BlockStorage().GetByNumber(ctx, blockNumber) if err != nil || existingBlock != nil { return existingBlock, err @@ -724,7 +733,7 @@ func (rm *RoundManager) getExistingBlockAnyFinalization(ctx context.Context, blo return nil, err } for _, candidate := range unfinalizedBlocks { - if candidate.Index.String() == blockNumber.String() { + if candidate.Index != nil && candidate.Index.Int != nil && candidate.Index.Cmp(blockNumber.Int) == 0 { return candidate, nil } } diff --git a/internal/round/round_manager.go b/internal/round/round_manager.go index e715e919..9851fe37 100644 --- a/internal/round/round_manager.go +++ b/internal/round/round_manager.go @@ -745,23 +745,27 @@ func (rm *RoundManager) restoreSmtFromStorage(ctx context.Context) (*api.BigInt, func (rm *RoundManager) Activate(ctx context.Context) error { rm.logger.WithContext(ctx).Info("Activating round manager") + rm.roundMutex.Lock() + alreadyActive := rm.activeCancel != nil + rm.roundMutex.Unlock() + if alreadyActive { + rm.logger.WithContext(ctx).Warn("Activation requested while already active; deactivating previous activation") + if err := rm.Deactivate(ctx); err != nil { + return fmt.Errorf("failed to deactivate previous activation: %w", err) + } + } + activeCtx, activeCancel := context.WithCancel(ctx) activated := false defer func() { if !activated { - activeCancel() - rm.roundMutex.Lock() - rm.activeCancel = nil - rm.precollectorDisabled = true - rm.roundMutex.Unlock() + if err := rm.Deactivate(ctx); err != nil { + rm.logger.WithContext(ctx).Error("Failed to deactivate after activation failure", "error", err.Error()) + } } }() rm.roundMutex.Lock() - if rm.activeCancel != nil { - rm.logger.WithContext(ctx).Warn("Activation requested while already active; cancelling previous activation") - rm.activeCancel() - } rm.activeCancel = activeCancel rm.precollectorDisabled = false rm.roundMutex.Unlock() From 5efb3f028d7ede75982007597c2ba2a3fd6e2145 Mon Sep 17 00:00:00 2001 From: jait91 Date: Wed, 17 Jun 2026 12:23:04 +0300 Subject: [PATCH 3/3] Handle canceled finalization during HA stand-down --- internal/round/batch_processor.go | 23 ++++- internal/round/finalize_duplicate_test.go | 109 +++++++++++++++++++++- 2 files changed, 126 insertions(+), 6 deletions(-) diff --git a/internal/round/batch_processor.go b/internal/round/batch_processor.go index d4f1be7d..ec5f2456 100644 --- a/internal/round/batch_processor.go +++ b/internal/round/batch_processor.go @@ -366,10 +366,20 @@ const ( // FinalizeBlockWithRetry retries finalization and uses recovery if block was partially stored. func (rm *RoundManager) FinalizeBlockWithRetry(ctx context.Context, block *models.Block) error { for attempt := 1; attempt <= maxFinalizeRetries; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + err := rm.FinalizeBlock(ctx, block) if err == nil { return nil } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } rm.logger.Error("FinalizeBlock failed", "attempt", attempt, @@ -393,7 +403,18 @@ func (rm *RoundManager) FinalizeBlockWithRetry(ctx context.Context, block *model if attempt < maxFinalizeRetries { rm.logger.Info("Retrying FinalizeBlock", "attempt", attempt) - time.Sleep(finalizeRetryDelay) + timer := time.NewTimer(finalizeRetryDelay) + select { + case <-timer.C: + case <-ctx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return ctx.Err() + } } } return fmt.Errorf("FinalizeBlock failed after %d attempts", maxFinalizeRetries) diff --git a/internal/round/finalize_duplicate_test.go b/internal/round/finalize_duplicate_test.go index 2b320fb8..69d81bb3 100644 --- a/internal/round/finalize_duplicate_test.go +++ b/internal/round/finalize_duplicate_test.go @@ -592,11 +592,10 @@ func (s *FinalizeDuplicateTestSuite) Test7_FinalizeBlockRejectsRoundNumberMismat } } -// Test8_DuplicateFinalizedBlockMustNotCommitDifferentSnapshot verifies that a -// retry for an already-finalized block is only idempotent if it is replaying the -// same block state. A duplicate block must not commit whatever snapshot happens -// to be in currentRound. -func (s *FinalizeDuplicateTestSuite) Test8_DuplicateFinalizedBlockMustNotCommitDifferentSnapshot() { +// Test8_DuplicateFinalizedBlockRejectsSnapshotRootMismatch verifies that a +// duplicate block retry cannot commit whatever snapshot happens to be in +// currentRound when that snapshot does not match the block root. +func (s *FinalizeDuplicateTestSuite) Test8_DuplicateFinalizedBlockRejectsSnapshotRootMismatch() { t := s.T() ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -665,3 +664,103 @@ func (s *FinalizeDuplicateTestSuite) Test8_DuplicateFinalizedBlockMustNotCommitD require.NotNil(t, storedBlock) assert.Equal(t, emptyRoot, storedBlock.RootHash.String(), "stored block root should remain unchanged") } + +// Test9_DuplicateFinalizedBlockRejectsDifferentBlockRecords verifies the +// duplicate-content guard after the snapshot-root guard has passed. +func (s *FinalizeDuplicateTestSuite) Test9_DuplicateFinalizedBlockRejectsDifferentBlockRecords() { + t := s.T() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + testLogger, err := logger.New("info", "text", "stdout", false) + require.NoError(t, err) + + threadSafeSMT := smt.NewThreadSafeSMT(smt.NewSparseMerkleTree(api.SHA256, 16+256)) + rm, err := NewRoundManager(ctx, s.cfg, testLogger, s.storage.CommitmentQueue(), s.storage, nil, state.NewSyncStateTracker(), nil, events.NewEventBus(testLogger), threadSafeSMT, nil) + require.NoError(t, err) + + commitments := testutil.CreateTestCertificationRequests(t, 2, "t9_req") + rm.currentRound = &Round{ + Number: api.NewBigInt(big.NewInt(72)), + State: RoundStateProcessing, + Commitments: commitments, + Snapshot: rm.smt.CreateSnapshot(), + } + + rm.roundMutex.Lock() + err = rm.processMiniBatch(ctx, commitments) + rm.roundMutex.Unlock() + require.NoError(t, err) + + rootHash := rm.currentRound.Snapshot.GetRootHash() + rootHashBytes, err := api.NewHexBytesFromString(rootHash) + require.NoError(t, err) + + block := models.NewBlock( + api.NewBigInt(big.NewInt(72)), + "unicity", + 0, + "1.0", + "mainnet", + rootHashBytes, + api.HexBytes{}, + api.HexBytes{}, + nil, + ) + block.Finalized = true + require.NoError(t, s.storage.BlockStorage().Store(ctx, block)) + + otherCommitments := testutil.CreateTestCertificationRequests(t, 1, "t9_existing") + require.NoError(t, s.storage.BlockRecordsStorage().Store(ctx, models.NewBlockRecords(block.Index, []api.StateID{otherCommitments[0].StateID}))) + + blockCountBefore, err := s.storage.BlockStorage().Count(ctx) + require.NoError(t, err) + blockRecordsCountBefore, err := s.storage.BlockRecordsStorage().Count(ctx) + require.NoError(t, err) + smtCountBefore, err := s.storage.SmtStorage().Count(ctx) + require.NoError(t, err) + recordCountBefore, err := s.storage.AggregatorRecordStorage().Count(ctx) + require.NoError(t, err) + + block.Finalized = false + err = rm.FinalizeBlock(ctx, block) + require.Error(t, err) + assert.Contains(t, err.Error(), "different request IDs") + + blockCountAfter, countErr := s.storage.BlockStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, blockCountBefore, blockCountAfter, "duplicate mismatch must not store another block") + + blockRecordsCountAfter, countErr := s.storage.BlockRecordsStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, blockRecordsCountBefore, blockRecordsCountAfter, "duplicate mismatch must not store block_records") + + smtCountAfter, countErr := s.storage.SmtStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, smtCountBefore, smtCountAfter, "duplicate mismatch must not store SMT nodes") + + recordCountAfter, countErr := s.storage.AggregatorRecordStorage().Count(ctx) + require.NoError(t, countErr) + assert.Equal(t, recordCountBefore, recordCountAfter, "duplicate mismatch must not store aggregator records") +} + +func (s *FinalizeDuplicateTestSuite) Test10_FinalizeBlockWithRetryPropagatesCancellation() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + block := models.NewBlock( + api.NewBigInt(big.NewInt(73)), + "unicity", + 0, + "1.0", + "mainnet", + api.HexBytes{}, + api.HexBytes{}, + api.HexBytes{}, + nil, + ) + rm := &RoundManager{} + + err := rm.FinalizeBlockWithRetry(ctx, block) + require.ErrorIs(s.T(), err, context.Canceled) +}