From 86cc1aee0bbd4435ba2bea47a74663090f797652 Mon Sep 17 00:00:00 2001 From: Nilesh Gupta Date: Wed, 26 Aug 2026 13:57:45 +0530 Subject: [PATCH] fix(uvalidator): index pending ballots by expiry height and sweep them in EndBlock Removes the O(active-set) IAVL scan CreateBallot ran on every ballot creation. --- .../uvalidator/ballot_voting_test.go | 119 +++++++- x/uvalidator/abci.go | 22 ++ x/uvalidator/keeper/ballot.go | 110 ++++++-- x/uvalidator/keeper/ballot_test.go | 264 +++++++++++++++++- x/uvalidator/keeper/genesis_test.go | 63 +++++ x/uvalidator/keeper/keeper.go | 24 +- x/uvalidator/keeper/voting.go | 6 + x/uvalidator/module.go | 14 + x/uvalidator/types/keys.go | 7 + 9 files changed, 590 insertions(+), 39 deletions(-) diff --git a/test/integration/uvalidator/ballot_voting_test.go b/test/integration/uvalidator/ballot_voting_test.go index 64b0be3a0..ff1b55b2e 100644 --- a/test/integration/uvalidator/ballot_voting_test.go +++ b/test/integration/uvalidator/ballot_voting_test.go @@ -4,12 +4,15 @@ import ( "fmt" "testing" + "cosmossdk.io/core/appmodule" + storetypes "cosmossdk.io/store/types" sdk "github.com/cosmos/cosmos-sdk/types" stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types" "github.com/stretchr/testify/require" "github.com/pushchain/push-chain-node/app" utils "github.com/pushchain/push-chain-node/test/utils" + uvalidatorkeeper "github.com/pushchain/push-chain-node/x/uvalidator/keeper" uvalidatortypes "github.com/pushchain/push-chain-node/x/uvalidator/types" ) @@ -102,7 +105,7 @@ func TestIntegration_CreateBallot(t *testing.T) { require.Equal(t, int64(60), ballot.BlockHeightExpiry) }) - t.Run("creating a new ballot expires stale active ballots", func(t *testing.T) { + t.Run("creating a new ballot does NOT expire stale active ballots", func(t *testing.T) { chainApp, ctx, validators := setupBallotTest(t, 2) k := chainApp.UvalidatorKeeper voters := voterAddrs(t, validators) @@ -114,16 +117,24 @@ func TestIntegration_CreateBallot(t *testing.T) { voters, 1, 1) require.NoError(t, err) - // Advance height past the expiry so the next CreateBallot triggers cleanup + // Advance past the old ballot's expiry and create another one. Creation + // no longer scans the active set — that walked every active ballot and + // paid an IAVL read each time. Expiry is the EndBlocker's job now. ctx = ctx.WithBlockHeight(10) _, err = k.CreateBallot(ctx, "new-ballot", uvalidatortypes.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX, voters, 1, 100) require.NoError(t, err) - // Old ballot should now be expired old, err := k.GetBallot(ctx, "old-ballot") require.NoError(t, err) + require.Equal(t, uvalidatortypes.BallotStatus_BALLOT_STATUS_PENDING, old.Status, + "CreateBallot must not sweep; the EndBlocker owns expiry") + + // The sweep is what expires it. + require.NoError(t, k.ExpireBallotsBeforeHeight(ctx, ctx.BlockHeight())) + old, err = k.GetBallot(ctx, "old-ballot") + require.NoError(t, err) require.Equal(t, uvalidatortypes.BallotStatus_BALLOT_STATUS_EXPIRED, old.Status) }) } @@ -963,3 +974,105 @@ func TestIntegration_IsTombstonedUniversalValidator(t *testing.T) { require.Contains(t, err.Error(), "invalid signer address") }) } + +// ─── EndBlocker expiry sweep ───────────────────────────────────────────────── + +// endBlockCtx moves ctx to the given height and attaches a block gas meter. +// The test fixture's context has none, and x/feemarket's EndBlocker — which the +// module manager runs before x/uvalidator's — errors out without one. +func endBlockCtx(ctx sdk.Context, height int64) sdk.Context { + return ctx.WithBlockHeight(height).WithBlockGasMeter(storetypes.NewInfiniteGasMeter()) +} + +// TestIntegration_UvalidatorEndBlockerRuns is the guard against a +// silently-never-called EndBlock. The sweep only fires if x/uvalidator is BOTH +// listed in SetOrderEndBlockers AND satisfies appmodule.HasEndBlocker — listing +// alone is not enough, the module manager skips modules that do not implement +// the interface. This drives the app's real EndBlocker and observes the effect. +func TestIntegration_UvalidatorEndBlockerRuns(t *testing.T) { + t.Run("expires due ballots and leaves the rest alone", func(t *testing.T) { + chainApp, ctx, validators := setupBallotTest(t, 3) + k := chainApp.UvalidatorKeeper + voters := voterAddrs(t, validators) + + // ctx starts at height 1 → expiry heights 2 and 1001. + _, err := k.CreateBallot(ctx, "eb-due", + uvalidatortypes.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX, + voters, 2, 1) + require.NoError(t, err) + + _, err = k.CreateBallot(ctx, "eb-future", + uvalidatortypes.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX, + voters, 2, 1000) + require.NoError(t, err) + + // Creation must not have swept anything. + due, err := k.GetBallot(ctx, "eb-due") + require.NoError(t, err) + require.Equal(t, uvalidatortypes.BallotStatus_BALLOT_STATUS_PENDING, due.Status) + + // Advance the block and run the app's real EndBlocker chain. + ctx = endBlockCtx(ctx, 5) + _, err = chainApp.EndBlocker(ctx) + require.NoError(t, err) + + // THE assertion: if x/uvalidator's EndBlock never fires, this fails. + due, err = k.GetBallot(ctx, "eb-due") + require.NoError(t, err) + require.Equal(t, uvalidatortypes.BallotStatus_BALLOT_STATUS_EXPIRED, due.Status, + "the x/uvalidator EndBlocker did not run the ballot expiry sweep") + + future, err := k.GetBallot(ctx, "eb-future") + require.NoError(t, err) + require.Equal(t, uvalidatortypes.BallotStatus_BALLOT_STATUS_PENDING, future.Status, + "a not-yet-due ballot must survive the sweep") + }) + + t.Run("caps a large backlog at MaxExpiriesPerBlock per block", func(t *testing.T) { + chainApp, ctx, validators := setupBallotTest(t, 3) + k := chainApp.UvalidatorKeeper + voters := voterAddrs(t, validators) + + const extra = 3 + total := uvalidatorkeeper.MaxExpiriesPerBlock + extra + for i := 0; i < total; i++ { + _, err := k.CreateBallot(ctx, fmt.Sprintf("eb-cap-%03d", i), + uvalidatortypes.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX, + voters, 2, 1) + require.NoError(t, err) + } + + countExpired := func(ctx sdk.Context) int { + n := 0 + require.NoError(t, k.ExpiredBallotIDs.Walk(ctx, nil, func(string) (bool, error) { + n++ + return false, nil + })) + return n + } + + ctx = endBlockCtx(ctx, 5) + _, err := chainApp.EndBlocker(ctx) + require.NoError(t, err) + require.Equal(t, uvalidatorkeeper.MaxExpiriesPerBlock, countExpired(ctx), + "one block must expire at most MaxExpiriesPerBlock ballots") + + // The leftovers are carried to the next block, not lost. + ctx = endBlockCtx(ctx, 6) + _, err = chainApp.EndBlocker(ctx) + require.NoError(t, err) + require.Equal(t, total, countExpired(ctx), + "the backlog remainder must be swept by the following block") + }) + + t.Run("module is wired into the EndBlocker ordering", func(t *testing.T) { + chainApp, _, _ := setupBallotTest(t, 1) + + require.Contains(t, chainApp.ModuleManager.OrderEndBlockers, uvalidatortypes.ModuleName, + "x/uvalidator must be listed in SetOrderEndBlockers or its EndBlock never runs") + + _, ok := chainApp.ModuleManager.Modules[uvalidatortypes.ModuleName].(appmodule.HasEndBlocker) + require.True(t, ok, + "x/uvalidator must implement appmodule.HasEndBlocker; the module manager skips modules that do not") + }) +} diff --git a/x/uvalidator/abci.go b/x/uvalidator/abci.go index 4a0117db1..5034d5513 100644 --- a/x/uvalidator/abci.go +++ b/x/uvalidator/abci.go @@ -65,6 +65,28 @@ func BeginBlocker(ctx sdk.Context, uvalidatorKeeper keeper.Keeper) error { return nil } +// EndBlocker runs the ballot expiry sweep once per block. +// +// Expiry used to be driven off ballot *creation*, which meant it both scanned +// the whole active set on a consensus hot path and only ran when inbound +// traffic happened to arrive. Running it here decouples expiry from traffic. +// +// A failed sweep is logged and swallowed rather than returned: expiry must +// never halt the chain, and the work is idempotent — anything not expired this +// block is still in PendingByExpiry for the next one. +func EndBlocker(ctx sdk.Context, uvalidatorKeeper keeper.Keeper) error { + defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyEndBlocker) + + if err := uvalidatorKeeper.ExpireBallotsBeforeHeight(ctx, ctx.BlockHeight()); err != nil { + ctx.Logger().Error("uvalidator: ballot expiry sweep failed", + "height", ctx.BlockHeight(), + "err", err.Error(), + ) + } + + return nil +} + // AllocateTokens performs reward and fee distribution to all validators based // on the F1 fee distribution specification. func AllocateTokens(ctx context.Context, totalPreviousPower int64, bondedVotes []abci.VoteInfo, k keeper.Keeper) error { diff --git a/x/uvalidator/keeper/ballot.go b/x/uvalidator/keeper/ballot.go index bafeb3115..a1802e37a 100644 --- a/x/uvalidator/keeper/ballot.go +++ b/x/uvalidator/keeper/ballot.go @@ -2,13 +2,21 @@ package keeper import ( "context" + "errors" "fmt" + "cosmossdk.io/collections" sdk "github.com/cosmos/cosmos-sdk/types" "github.com/pushchain/push-chain-node/x/uvalidator/types" ) +// MaxExpiriesPerBlock bounds how many ballots a single expiry sweep may +// transition. It keeps the per-block cost of the sweep constant even if a +// large backlog of ballots comes due at once; anything left over stays in +// PendingByExpiry and is picked up by the next block's sweep. +const MaxExpiriesPerBlock = 50 + // CreateBallot creates a new ballot with the given parameters, stores it, and marks it as active. func (k Keeper) CreateBallot( ctx context.Context, @@ -33,10 +41,10 @@ func (k Keeper) CreateBallot( "block_height", blockHeight, ) - // First, expire any old ballots before this height - if err := k.ExpireBallotsBeforeHeight(ctx, blockHeight); err != nil { - return types.Ballot{}, err - } + // NOTE: creation deliberately does NOT sweep expired ballots. That scan used + // to run here on every create and walked the whole active set, paying an + // IAVL read + unmarshal per entry. Expiry now runs once per block in the + // module EndBlocker off the PendingByExpiry index instead. // Create ballot ballot := types.NewBallot( @@ -57,6 +65,9 @@ func (k Keeper) CreateBallot( if err := k.ActiveBallotIDs.Set(ctx, ballot.Id); err != nil { return types.Ballot{}, err } + if err := k.indexPending(ctx, ballot.Id, ballot.BlockHeightExpiry); err != nil { + return types.Ballot{}, err + } k.Logger().Debug("ballot created and marked active", "ballot_id", ballot.Id, @@ -100,9 +111,16 @@ func (k Keeper) SetBallot(ctx context.Context, ballot types.Ballot) error { return k.Ballots.Set(ctx, ballot.Id, ballot) } -// DeleteBallot removes a ballot and its ID from all collections +// DeleteBallot removes a ballot and its ID from all collections. +// +// The PendingByExpiry row is keyed by the ballot's expiry height, so the record +// must be read before it is removed. A missing record means there is nothing to +// unindex (DeleteBallot stays idempotent on absent IDs). func (k Keeper) DeleteBallot(ctx context.Context, id string) error { k.Logger().Debug("deleting ballot", "ballot_id", id) + if ballot, err := k.Ballots.Get(ctx, id); err == nil { + _ = k.unindexPending(ctx, id, ballot.BlockHeightExpiry) + } if err := k.Ballots.Remove(ctx, id); err != nil { return err } @@ -137,6 +155,9 @@ func (k Keeper) MarkBallotExpired(ctx context.Context, id string) error { if err := k.ActiveBallotIDs.Remove(ctx, id); err != nil { return err } + if err := k.unindexPending(ctx, id, ballot.BlockHeightExpiry); err != nil { + return err + } if err := k.ExpiredBallotIDs.Set(ctx, id); err != nil { return err } @@ -175,6 +196,9 @@ func (k Keeper) MarkBallotFinalized(ctx context.Context, id string, status types if err := k.ActiveBallotIDs.Remove(ctx, id); err != nil { return err } + if err := k.unindexPending(ctx, id, ballot.BlockHeightExpiry); err != nil { + return err + } if err := k.FinalizedBallotIDs.Set(ctx, id); err != nil { return err } @@ -317,45 +341,83 @@ func (k Keeper) RecomputeBallotQuorum(ctx context.Context, ballotID string) ( return oldEligibleCount, newEligibleCount, oldThreshold, newThreshold, types.BallotStatus_BALLOT_STATUS_PENDING, nil } -// ExpireBallotsBeforeHeight checks active ballots and marks expired ones. -// It uses a two-phase approach: first collect IDs to expire, then mutate, -// to avoid modifying the ActiveBallotIDs collection during iteration. +// indexPending adds the (expiryHeight, ballotID) row that mirrors an entry in +// ActiveBallotIDs. Every ActiveBallotIDs.Set must be paired with this call, or +// the ballot becomes invisible to the expiry sweep. +func (k Keeper) indexPending(ctx context.Context, id string, expiryHeight int64) error { + return k.PendingByExpiry.Set(ctx, collections.Join(expiryHeight, id)) +} + +// unindexPending drops the (expiryHeight, ballotID) row. Every +// ActiveBallotIDs.Remove must be paired with this call, or the sweep keeps +// re-visiting a ballot that is no longer active. +func (k Keeper) unindexPending(ctx context.Context, id string, expiryHeight int64) error { + return k.PendingByExpiry.Remove(ctx, collections.Join(expiryHeight, id)) +} + +// ExpireBallotsBeforeHeight marks every active ballot whose expiry height is at +// or below currentHeight as expired, up to MaxExpiriesPerBlock per call. +// +// It ranges over the PendingByExpiry index rather than ActiveBallotIDs. Because +// collections.Pair orders by its first component, iteration ends at the first +// row past currentHeight: ballots that are not yet due are never visited, and +// no Ballots.Get is needed to decide whether a ballot is due — the expiry +// height IS the key. Only ballots actually being expired are ever loaded. +// +// It keeps the original two-phase shape: IDs are collected while the iterator +// is open and mutated only after it is closed. Mutating a collection mid- +// iteration skips entries at best and panics at worst. func (k Keeper) ExpireBallotsBeforeHeight(ctx context.Context, currentHeight int64) error { - iter, err := k.ActiveBallotIDs.Iterate(ctx, nil) + // NewPrefixUntilPairRange ends at the prefix-end of currentHeight, so the + // range is inclusive of currentHeight — matching the `<=` expiry semantics. + rng := collections.NewPrefixUntilPairRange[int64, string](currentHeight) + + iter, err := k.PendingByExpiry.Iterate(ctx, rng) if err != nil { return err } - // Phase 1: collect IDs to expire - var toExpire []string + // Phase 1: collect the due index keys, bounded by MaxExpiriesPerBlock. + var due []collections.Pair[int64, string] for ; iter.Valid(); iter.Next() { - id, err := iter.Key() + key, err := iter.Key() if err != nil { iter.Close() return err } - ballot, err := k.Ballots.Get(ctx, id) - if err != nil { - iter.Close() - return err - } - - if ballot.BlockHeightExpiry <= currentHeight { - toExpire = append(toExpire, id) + due = append(due, key) + if len(due) >= MaxExpiriesPerBlock { + break } } // Close iterator explicitly before mutation phase to release the IAVL snapshot iter.Close() - if len(toExpire) > 0 { - k.Logger().Debug("expiring stale ballots", "count", len(toExpire), "current_height", currentHeight) + if len(due) == 0 { + return nil } + k.Logger().Debug("expiring stale ballots", "count", len(due), "current_height", currentHeight) + // Phase 2: expire collected ballots (safe — iterator is closed) - for _, id := range toExpire { - if err := k.MarkBallotExpired(ctx, id); err != nil { + for _, key := range due { + id := key.K2() + err := k.MarkBallotExpired(ctx, id) + switch { + case err == nil: + case errors.Is(err, collections.ErrNotFound): + // Defensive: an index row whose ballot record no longer exists must + // not wedge the sweep every block. Drop the orphan row and continue. + k.Logger().Warn("dropping orphaned ballot expiry index entry", + "ballot_id", id, + "expiry_height", key.K1(), + ) + if rErr := k.PendingByExpiry.Remove(ctx, key); rErr != nil { + return rErr + } + default: return err } } diff --git a/x/uvalidator/keeper/ballot_test.go b/x/uvalidator/keeper/ballot_test.go index f6432003d..0c525bcc9 100644 --- a/x/uvalidator/keeper/ballot_test.go +++ b/x/uvalidator/keeper/ballot_test.go @@ -4,10 +4,46 @@ import ( "fmt" "testing" + "cosmossdk.io/collections" + "github.com/pushchain/push-chain-node/x/uvalidator/keeper" "github.com/pushchain/push-chain-node/x/uvalidator/types" "github.com/stretchr/testify/require" ) +const inboundBallot = types.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX + +// pendingIndexed reports whether the (expiryHeight, ballotID) row exists in the +// PendingByExpiry index. +func pendingIndexed(t *testing.T, f *testFixture, id string, expiryHeight int64) bool { + t.Helper() + has, err := f.k.PendingByExpiry.Has(f.ctx, collections.Join(expiryHeight, id)) + require.NoError(t, err) + return has +} + +// pendingIndexIDs returns every ballot ID currently carried by the expiry index. +func pendingIndexIDs(t *testing.T, f *testFixture) []string { + t.Helper() + ids := []string{} + require.NoError(t, f.k.PendingByExpiry.Walk(f.ctx, nil, + func(key collections.Pair[int64, string]) (bool, error) { + ids = append(ids, key.K2()) + return false, nil + })) + return ids +} + +// expiredCount returns how many ballots sit in ExpiredBallotIDs. +func expiredCount(t *testing.T, f *testFixture) int { + t.Helper() + n := 0 + require.NoError(t, f.k.ExpiredBallotIDs.Walk(f.ctx, nil, func(string) (bool, error) { + n++ + return false, nil + })) + return n +} + func TestCreateAndGetBallot(t *testing.T) { f := SetupTest(t) require := require.New(t) @@ -150,32 +186,48 @@ func TestExpireBallotsBeforeHeight(t *testing.T) { require.Equal(types.BallotStatus_BALLOT_STATUS_PENDING, got2.Status) } -func TestCreateBallot_ExpiresOldOnCreate(t *testing.T) { +// Creation must no longer sweep expired ballots: that scan walked the whole +// active set on a consensus hot path. Expiry moved to the module EndBlocker. +func TestCreateBallot_DoesNotExpireOnCreate(t *testing.T) { f := SetupTest(t) require := require.New(t) - // Create a ballot that expires quickly (expiry = 1 block) + // Create a ballot that comes due quickly (expiry = 1 block) oldBallot, err := f.k.CreateBallot(f.ctx, "old", types.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX, []string{"v1"}, 1, 1) require.NoError(err) - // Manually simulate advancing block height + // Advance well past the old ballot's expiry height f.ctx = f.ctx.WithBlockHeight(oldBallot.BlockHeightCreated + 5) - // Now create a NEW ballot → should trigger expiry cleanup of the old one + // Creating a NEW ballot must NOT expire the due one — no scan on create. newBallot, err := f.k.CreateBallot(f.ctx, "new", types.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX, []string{"v2"}, 1, 10) require.NoError(err) - - // New ballot must be created fine require.Equal("new", newBallot.Id) - // Old ballot should now be expired got, err := f.k.GetBallot(f.ctx, "old") require.NoError(err) + require.Equal(types.BallotStatus_BALLOT_STATUS_PENDING, got.Status, + "CreateBallot must not expire ballots; expiry belongs to the EndBlocker sweep") + + has, err := f.k.ActiveBallotIDs.Has(f.ctx, "old") + require.NoError(err) + require.True(has, "due ballot must stay active until the sweep runs") + + // The sweep — not creation — is what expires it. + require.NoError(f.k.ExpireBallotsBeforeHeight(f.ctx, f.ctx.BlockHeight())) + + got, err = f.k.GetBallot(f.ctx, "old") + require.NoError(err) require.Equal(types.BallotStatus_BALLOT_STATUS_EXPIRED, got.Status) + + // The not-yet-due ballot survives the sweep. + gotNew, err := f.k.GetBallot(f.ctx, "new") + require.NoError(err) + require.Equal(types.BallotStatus_BALLOT_STATUS_PENDING, gotNew.Status) } func TestCreateBallot_NoExpiryTriggered(t *testing.T) { @@ -207,7 +259,7 @@ func TestCreateBallot_NoExpiryTriggered(t *testing.T) { require.Equal("newer", got2.Id) } -func TestCreateBallot_ExpiresMultipleOld(t *testing.T) { +func TestCreateBallot_DoesNotExpireMultipleOldOnCreate(t *testing.T) { f := SetupTest(t) require := require.New(t) @@ -225,19 +277,30 @@ func TestCreateBallot_ExpiresMultipleOld(t *testing.T) { // Advance height beyond both expiries f.ctx = f.ctx.WithBlockHeight(b2.BlockHeightCreated + 5) - // Create fresh ballot (triggers cleanup) + // Create fresh ballot — must NOT trigger any cleanup _, err = f.k.CreateBallot(f.ctx, "fresh", types.BallotObservationType_BALLOT_OBSERVATION_TYPE_INBOUND_TX, []string{"v3"}, 1, 5) require.NoError(err) - // Both old ballots should now be expired + // Both due ballots must still be pending — creation does not sweep got1, err := f.k.GetBallot(f.ctx, b1.Id) require.NoError(err) - require.Equal(types.BallotStatus_BALLOT_STATUS_EXPIRED, got1.Status) + require.Equal(types.BallotStatus_BALLOT_STATUS_PENDING, got1.Status) got2, err := f.k.GetBallot(f.ctx, b2.Id) require.NoError(err) + require.Equal(types.BallotStatus_BALLOT_STATUS_PENDING, got2.Status) + + // The sweep expires them. + require.NoError(f.k.ExpireBallotsBeforeHeight(f.ctx, f.ctx.BlockHeight())) + + got1, err = f.k.GetBallot(f.ctx, b1.Id) + require.NoError(err) + require.Equal(types.BallotStatus_BALLOT_STATUS_EXPIRED, got1.Status) + + got2, err = f.k.GetBallot(f.ctx, b2.Id) + require.NoError(err) require.Equal(types.BallotStatus_BALLOT_STATUS_EXPIRED, got2.Status) } @@ -401,3 +464,182 @@ func TestExpireBallotsBeforeHeight_EmptySet(t *testing.T) { err := f.k.ExpireBallotsBeforeHeight(f.ctx, 100) require.NoError(err) } + +// ─── PendingByExpiry index ─────────────────────────────────────────────────── + +// TestPendingByExpiryIndex_MirrorsEveryActiveSetWriter is the guard against a +// missed mirror site. Every writer of ActiveBallotIDs must write PendingByExpiry +// too; a miss makes the index drift from reality, which either hides a ballot +// from the sweep forever or lets the sweep act on a ballot that is no longer +// active. Each subtest asserts the index state FIRST, so the assertion that +// catches the missing mirror runs before anything else can abort the subtest. +func TestPendingByExpiryIndex_MirrorsEveryActiveSetWriter(t *testing.T) { + t.Run("CreateBallot indexes under the expiry height", func(t *testing.T) { + f := SetupTest(t) + + b, err := f.k.CreateBallot(f.ctx, "idx-create", inboundBallot, []string{"v1"}, 1, 7) + require.NoError(t, err) + + require.True(t, pendingIndexed(t, f, "idx-create", b.BlockHeightExpiry), + "CreateBallot must mirror ActiveBallotIDs into PendingByExpiry") + require.Equal(t, []string{"idx-create"}, pendingIndexIDs(t, f)) + + // The row must be keyed by the expiry height, not the creation height — + // that is the whole point of the index. + require.NotEqual(t, b.BlockHeightCreated, b.BlockHeightExpiry) + require.False(t, pendingIndexed(t, f, "idx-create", b.BlockHeightCreated), + "index must be keyed by expiry height, not creation height") + }) + + t.Run("MarkBallotExpired unindexes", func(t *testing.T) { + f := SetupTest(t) + + b, err := f.k.CreateBallot(f.ctx, "idx-expire", inboundBallot, []string{"v1"}, 1, 3) + require.NoError(t, err) + require.NoError(t, f.k.MarkBallotExpired(f.ctx, b.Id)) + + require.False(t, pendingIndexed(t, f, b.Id, b.BlockHeightExpiry), + "MarkBallotExpired must remove the PendingByExpiry row") + require.Empty(t, pendingIndexIDs(t, f)) + + // A leaked row would be re-processed by the sweep every single block, + // permanently consuming part of the per-block budget. + require.NoError(t, f.k.ExpireBallotsBeforeHeight(f.ctx, b.BlockHeightExpiry+1)) + require.Empty(t, pendingIndexIDs(t, f)) + }) + + t.Run("MarkBallotFinalized unindexes", func(t *testing.T) { + f := SetupTest(t) + + b, err := f.k.CreateBallot(f.ctx, "idx-final", inboundBallot, []string{"v1"}, 1, 3) + require.NoError(t, err) + require.NoError(t, f.k.MarkBallotFinalized(f.ctx, b.Id, types.BallotStatus_BALLOT_STATUS_PASSED)) + + require.False(t, pendingIndexed(t, f, b.Id, b.BlockHeightExpiry), + "MarkBallotFinalized must remove the PendingByExpiry row") + + // Behavioural consequence of a leaked row: the sweep would reach a + // finalized ballot and overwrite PASSED with EXPIRED. + require.NoError(t, f.k.ExpireBallotsBeforeHeight(f.ctx, b.BlockHeightExpiry+1)) + got, err := f.k.GetBallot(f.ctx, b.Id) + require.NoError(t, err) + require.Equal(t, types.BallotStatus_BALLOT_STATUS_PASSED, got.Status, + "a stale index row let the sweep overwrite a finalized ballot") + }) + + t.Run("DeleteBallot unindexes", func(t *testing.T) { + f := SetupTest(t) + + b, err := f.k.CreateBallot(f.ctx, "idx-delete", inboundBallot, []string{"v1"}, 1, 3) + require.NoError(t, err) + require.NoError(t, f.k.DeleteBallot(f.ctx, b.Id)) + + require.False(t, pendingIndexed(t, f, b.Id, b.BlockHeightExpiry), + "DeleteBallot must remove the PendingByExpiry row") + require.Empty(t, pendingIndexIDs(t, f)) + + // Deleting an absent ballot stays a no-op. + require.NoError(t, f.k.DeleteBallot(f.ctx, b.Id)) + require.Empty(t, pendingIndexIDs(t, f)) + }) + + t.Run("VoteOnBallot create path indexes the new ballot", func(t *testing.T) { + f := SetupTest(t) + + b, _, isNew, err := f.k.VoteOnBallot(f.ctx, "idx-vote", inboundBallot, + "v1", types.VoteResult_VOTE_RESULT_SUCCESS, + []string{"v1", "v2"}, 2, 9) + require.NoError(t, err) + require.True(t, isNew) + + require.True(t, pendingIndexed(t, f, "idx-vote", b.BlockHeightExpiry), + "the vote-driven create path must leave the expiry index populated") + }) +} + +// TestExpireBallotsBeforeHeight_OnlyDueRowsAreTouched pins the range semantics: +// rows past currentHeight are neither expired nor dropped from the index. +func TestExpireBallotsBeforeHeight_OnlyDueRowsAreTouched(t *testing.T) { + f := SetupTest(t) + require := require.New(t) + + // Created at height 0 → expiry heights 3, 10 and exactly 5. + due, err := f.k.CreateBallot(f.ctx, "due", inboundBallot, []string{"v1"}, 1, 3) + require.NoError(err) + atBoundary, err := f.k.CreateBallot(f.ctx, "at-boundary", inboundBallot, []string{"v1"}, 1, 5) + require.NoError(err) + future, err := f.k.CreateBallot(f.ctx, "future", inboundBallot, []string{"v1"}, 1, 10) + require.NoError(err) + + require.NoError(f.k.ExpireBallotsBeforeHeight(f.ctx, 5)) + + // The not-yet-due row must survive in the index for a later block. + require.Equal([]string{"future"}, pendingIndexIDs(t, f)) + require.True(pendingIndexed(t, f, future.Id, future.BlockHeightExpiry)) + + gotFuture, err := f.k.GetBallot(f.ctx, future.Id) + require.NoError(err) + require.Equal(types.BallotStatus_BALLOT_STATUS_PENDING, gotFuture.Status) + + // Expiry is inclusive of currentHeight (`<=` semantics). + for _, id := range []string{due.Id, atBoundary.Id} { + got, err := f.k.GetBallot(f.ctx, id) + require.NoError(err) + require.Equal(types.BallotStatus_BALLOT_STATUS_EXPIRED, got.Status, "ballot %s should be expired", id) + } +} + +// TestExpireBallotsBeforeHeight_CapsAtMaxExpiriesPerBlock proves the per-block +// bound holds and that the leftovers are carried in the index to the next block. +func TestExpireBallotsBeforeHeight_CapsAtMaxExpiriesPerBlock(t *testing.T) { + f := SetupTest(t) + require := require.New(t) + + const extra = 7 + total := keeper.MaxExpiriesPerBlock + extra + + // All due at height 1, i.e. the whole backlog comes due at once. + for i := 0; i < total; i++ { + _, err := f.k.CreateBallot(f.ctx, fmt.Sprintf("cap-%03d", i), inboundBallot, []string{"v1"}, 1, 1) + require.NoError(err) + } + require.Len(pendingIndexIDs(t, f), total) + + // First sweep: exactly MaxExpiriesPerBlock, never the whole backlog. + require.NoError(f.k.ExpireBallotsBeforeHeight(f.ctx, 100)) + require.Equal(keeper.MaxExpiriesPerBlock, expiredCount(t, f), + "a single sweep must expire at most MaxExpiriesPerBlock ballots") + require.Len(pendingIndexIDs(t, f), extra, + "the remainder must stay in the index for the next block") + + // Second sweep drains the rest. + require.NoError(f.k.ExpireBallotsBeforeHeight(f.ctx, 100)) + require.Equal(total, expiredCount(t, f)) + require.Empty(pendingIndexIDs(t, f)) +} + +// TestExpireBallotsBeforeHeight_OrphanedIndexRow covers the defensive path: an +// index row whose ballot record is gone must be dropped instead of wedging the +// sweep on every subsequent block. +func TestExpireBallotsBeforeHeight_OrphanedIndexRow(t *testing.T) { + f := SetupTest(t) + require := require.New(t) + + b, err := f.k.CreateBallot(f.ctx, "orphan", inboundBallot, []string{"v1"}, 1, 1) + require.NoError(err) + + // Drop the record behind the index row's back. + require.NoError(f.k.Ballots.Remove(f.ctx, b.Id)) + require.True(pendingIndexed(t, f, b.Id, b.BlockHeightExpiry)) + + // A healthy ballot queued behind the orphan must still get expired. + good, err := f.k.CreateBallot(f.ctx, "good", inboundBallot, []string{"v1"}, 1, 2) + require.NoError(err) + + require.NoError(f.k.ExpireBallotsBeforeHeight(f.ctx, 100)) + + require.Empty(pendingIndexIDs(t, f), "the orphaned row must be dropped, not retried forever") + gotGood, err := f.k.GetBallot(f.ctx, good.Id) + require.NoError(err) + require.Equal(types.BallotStatus_BALLOT_STATUS_EXPIRED, gotGood.Status) +} diff --git a/x/uvalidator/keeper/genesis_test.go b/x/uvalidator/keeper/genesis_test.go index b052f32ba..0f025f0c4 100755 --- a/x/uvalidator/keeper/genesis_test.go +++ b/x/uvalidator/keeper/genesis_test.go @@ -3,6 +3,7 @@ package keeper_test import ( "testing" + "cosmossdk.io/collections" sdk "github.com/cosmos/cosmos-sdk/types" "github.com/pushchain/push-chain-node/x/uvalidator/types" "github.com/stretchr/testify/require" @@ -64,3 +65,65 @@ func TestGenesisExportImportRoundTrip(t *testing.T) { require.Equal(t, len(exported.ActiveBallotIds), len(reExported.ActiveBallotIds)) require.Equal(t, exported.UniversalValidators[0].Key, reExported.UniversalValidators[0].Key) } + +// TestInitGenesisRebuildsPendingByExpiry pins the reason this change needs no +// state migration: the expiry index is derived state that InitGenesis rebuilds +// from the ballots it just restored. +func TestInitGenesisRebuildsPendingByExpiry(t *testing.T) { + f := SetupTest(t) + f.k.InitGenesis(f.ctx, &types.GenesisState{Params: types.Params{Admin: f.addrs[0].String()}}) + + genesis := &types.GenesisState{ + Params: types.Params{Admin: f.addrs[0].String()}, + Ballots: []types.Ballot{ + {Id: "g-1", Status: types.BallotStatus_BALLOT_STATUS_PENDING, BlockHeightExpiry: 42}, + {Id: "g-2", Status: types.BallotStatus_BALLOT_STATUS_PENDING, BlockHeightExpiry: 7}, + {Id: "g-3", Status: types.BallotStatus_BALLOT_STATUS_PASSED, BlockHeightExpiry: 9}, + }, + ActiveBallotIds: []string{"g-1", "g-2"}, + FinalizedBallotIds: []string{"g-3"}, + } + + f2 := SetupTest(t) + require.NoError(t, f2.k.InitGenesis(f2.ctx, genesis)) + + // Both active ballots are indexed under their own expiry heights. + for _, tc := range []struct { + id string + expiry int64 + }{{"g-1", 42}, {"g-2", 7}} { + has, err := f2.k.PendingByExpiry.Has(f2.ctx, collections.Join(tc.expiry, tc.id)) + require.NoError(t, err) + require.True(t, has, "InitGenesis must rebuild the expiry index for %s", tc.id) + } + + // The finalized ballot is not active, so it must not be indexed. + has, err := f2.k.PendingByExpiry.Has(f2.ctx, collections.Join(int64(9), "g-3")) + require.NoError(t, err) + require.False(t, has, "only active ballots belong in the expiry index") + + // The rebuilt index drives the sweep: g-2 (expiry 7) is due at height 10, + // g-1 (expiry 42) is not. + require.NoError(t, f2.k.ExpireBallotsBeforeHeight(f2.ctx, 10)) + + got, err := f2.k.GetBallot(f2.ctx, "g-2") + require.NoError(t, err) + require.Equal(t, types.BallotStatus_BALLOT_STATUS_EXPIRED, got.Status) + + got, err = f2.k.GetBallot(f2.ctx, "g-1") + require.NoError(t, err) + require.Equal(t, types.BallotStatus_BALLOT_STATUS_PENDING, got.Status) +} + +// TestInitGenesisRejectsDanglingActiveBallotID: an active ballot id with no +// ballot record cannot be indexed, so it is surfaced rather than silently +// dropped into an index that no longer matches the active set. +func TestInitGenesisRejectsDanglingActiveBallotID(t *testing.T) { + f := SetupTest(t) + err := f.k.InitGenesis(f.ctx, &types.GenesisState{ + Params: types.Params{Admin: f.addrs[0].String()}, + ActiveBallotIds: []string{"ghost"}, + }) + require.Error(t, err) + require.Contains(t, err.Error(), "ghost") +} diff --git a/x/uvalidator/keeper/keeper.go b/x/uvalidator/keeper/keeper.go index 2345d4382..c2603181c 100755 --- a/x/uvalidator/keeper/keeper.go +++ b/x/uvalidator/keeper/keeper.go @@ -2,6 +2,7 @@ package keeper import ( "context" + "fmt" "github.com/cosmos/cosmos-sdk/codec" sdk "github.com/cosmos/cosmos-sdk/types" @@ -33,6 +34,14 @@ type Keeper struct { ExpiredBallotIDs collections.KeySet[string] // set of ballot IDs that have expired (not yet pruned) FinalizedBallotIDs collections.KeySet[string] // set of ballot IDs that are PASSED or REJECTED + // PendingByExpiry is a secondary index over ActiveBallotIDs keyed by + // (expiryHeight, ballotID). collections.Pair orders by the first component, + // so the expiry sweep can range over [0, currentHeight] and stop at the + // first entry beyond it — ballots that are not due are never visited, and + // the height being part of the key means no Ballots.Get is needed to decide + // whether a ballot is due. Every ActiveBallotIDs writer must mirror here. + PendingByExpiry collections.KeySet[collections.Pair[int64, string]] + StakingKeeper types.StakingKeeper SlashingKeeper types.SlashingKeeper UtssKeeper types.UtssKeeper @@ -98,6 +107,10 @@ func NewKeeper( sb, types.FinalizedBallotIDsKey, types.FinalizedBallotIDsName, collections.StringKey, ), + PendingByExpiry: collections.NewKeySet( + sb, types.PendingByExpiryKey, types.PendingByExpiryName, + collections.PairKeyCodec(collections.Int64Key, collections.StringKey), + ), authority: authority, StakingKeeper: stakingKeeper, @@ -144,11 +157,20 @@ func (k *Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) erro } } - // Restore ActiveBallotIDs + // Restore ActiveBallotIDs, rebuilding the PendingByExpiry index from the + // ballots restored just above. This is why no state migration is needed for + // a chain that starts from (or is re-imported through) genesis. for _, id := range data.ActiveBallotIds { + ballot, err := k.Ballots.Get(ctx, id) + if err != nil { + return fmt.Errorf("active ballot id %q has no matching ballot record in genesis: %w", id, err) + } if err := k.ActiveBallotIDs.Set(ctx, id); err != nil { return err } + if err := k.PendingByExpiry.Set(ctx, collections.Join(ballot.BlockHeightExpiry, id)); err != nil { + return err + } } // Restore ExpiredBallotIDs diff --git a/x/uvalidator/keeper/voting.go b/x/uvalidator/keeper/voting.go index c0adb980b..adecb7cbc 100644 --- a/x/uvalidator/keeper/voting.go +++ b/x/uvalidator/keeper/voting.go @@ -181,6 +181,12 @@ func (k Keeper) VoteOnBallot( if err != nil { return ballot, false, false, errors.Wrap(err, "Error while voting on the ballot") } + // Mirror the active-set write into the expiry index. CreateBallot has + // already written both; both writes are idempotent, and pairing them + // here keeps the invariant local to every ActiveBallotIDs.Set site. + if err := k.indexPending(ctx, id, ballot.BlockHeightExpiry); err != nil { + return ballot, false, false, errors.Wrap(err, "Error while voting on the ballot") + } } ballot, err = k.AddVoteToBallot(ctx, ballot, voter, voteResult) diff --git a/x/uvalidator/module.go b/x/uvalidator/module.go index 408cfbd6f..155e2b18d 100755 --- a/x/uvalidator/module.go +++ b/x/uvalidator/module.go @@ -10,6 +10,7 @@ import ( abci "github.com/cometbft/cometbft/abci/types" "cosmossdk.io/client/v2/autocli" + "cosmossdk.io/core/appmodule" errorsmod "cosmossdk.io/errors" "github.com/cosmos/cosmos-sdk/client" @@ -32,6 +33,13 @@ var ( _ module.AppModuleGenesis = AppModule{} _ module.AppModule = AppModule{} + // The module manager only calls BeginBlock/EndBlock on modules that satisfy + // these interfaces — being listed in SetOrderBeginBlockers/EndBlockers is + // necessary but not sufficient. These assertions fail the build if a + // signature drifts and the hook silently stops firing. + _ appmodule.HasBeginBlocker = AppModule{} + _ appmodule.HasEndBlocker = AppModule{} + _ autocli.HasAutoCLIConfig = AppModule{} ) @@ -172,3 +180,9 @@ func (a AppModule) BeginBlock(ctx context.Context) error { return BeginBlocker(sdkCtx, a.keeper) } + +func (a AppModule) EndBlock(ctx context.Context) error { + sdkCtx := sdk.UnwrapSDKContext(ctx) + + return EndBlocker(sdkCtx, a.keeper) +} diff --git a/x/uvalidator/types/keys.go b/x/uvalidator/types/keys.go index 9f0fb92da..1b9ef7d30 100755 --- a/x/uvalidator/types/keys.go +++ b/x/uvalidator/types/keys.go @@ -49,6 +49,13 @@ var ( // FinalizedBallotIDsName is the name of the finalized ballot IDs set. FinalizedBallotIDsName = "finalized_ballot_ids" + + // PendingByExpiryKey is the key for the (expiryHeight, ballotID) index over + // the active ballot set. Next free prefix after FinalizedBallotIDsKey(6). + PendingByExpiryKey = collections.NewPrefix(7) + + // PendingByExpiryName is the name of the pending-by-expiry index. + PendingByExpiryName = "pending_ballots_by_expiry" ) const (