fix: F-2026-18795 | [Dual Defense] Unbounded Active-Ballot Expiry Scan on Create + Disabled 100M Expiry

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")
+	})
+}

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 {

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
 		}
 	}