fix: prune retained attestation state (#408)

* fix: prune retained attestation state

* refactor: centralize attestation retention boundary

* fix: keep default attestation history for IBC

* fix: align attestation retention with pruning

Author: randygrok

modules/network/README.md

@@ -128,7 +128,20 @@ The module's `EndBlocker` is executed at the end of every block and performs the
 
 1.  **Quorum Evaluation**: It iterates through recent blocks that have received attestations and checks if the cumulative voting power of the attesters has reached the required quorum.
 2.  **Checkpoint Emission**: If quorum is met for a block, it emits an `EventSoftCheckpoint` with the `height`, `block_hash`, a bitmap of the participating attesters, and the total voting power.
-3.  **Epoch Processing**: It checks if the current block is the end of an epoch. If it is, it performs accounting tasks, such as updating the attester index map for the next epoch.
+3.  **Epoch Processing**: It checks if the current block is the end of an epoch. If it is, it performs accounting tasks, such as pruning old attestation state and updating the attester index map for the next epoch.
+
+## Attestation State Retention
+
+Attestation state is retained for the last `prune_after` epochs. At each epoch boundary, the module computes the first retained epoch as `current_epoch - prune_after` and prunes all attestation data below that boundary.
+
+The retention policy covers the attestation stores together:
+
+- raw per-height attestation bitmaps
+- stored attestation metadata
+- per-epoch participation bitmaps
+- per-height attester signatures
+
+Queries for pruned heights behave like queries for missing data. `MsgAttest` uses the same height boundary and rejects attestations below the retention window so new writes cannot recreate pruned state.
 
 ## Genesis and Queries
 

modules/network/keeper/abci.go

@@ -112,10 +112,9 @@ func (k Keeper) processEpochEnd(ctx sdk.Context, epoch uint64) error {
 		}
 	}
 
-	// todo (Alex): find a way to prune only bitmaps that are not used anymore
-	// if err := k.PruneOldBitmaps(ctx, epoch); err != nil {
-	// 	return fmt.Errorf("pruning old data at epoch %d: %w", epoch, err)
-	// }
+	if err := k.PruneOldBitmaps(ctx, epoch); err != nil {
+		return fmt.Errorf("pruning old data at epoch %d: %w", epoch, err)
+	}
 
 	if err := k.BuildValidatorIndexMap(ctx); err != nil {
 		return fmt.Errorf("rebuilding validator index map at epoch %d: %w", epoch, err)

modules/network/keeper/keeper.go

@@ -308,40 +308,62 @@ func (k Keeper) IsSoftConfirmed(ctx sdk.Context, height int64) (bool, error) {
 	return k.CheckQuorum(ctx, votedPower, totalPower)
 }
 
-// PruneOldBitmaps removes bitmaps older than PruneAfter epochs
-func (k Keeper) PruneOldBitmaps(ctx sdk.Context, currentEpoch uint64) error {
+type attestationRetentionBoundary struct {
+	firstRetainedEpoch  uint64
+	firstRetainedHeight int64
+}
+
+func (b attestationRetentionBoundary) prunesHeight(height int64) bool {
+	return height < b.firstRetainedHeight
+}
+
+func (k Keeper) attestationRetentionBoundary(ctx sdk.Context, currentEpoch uint64) *attestationRetentionBoundary {
 	params := k.GetParams(ctx)
-	if params.PruneAfter == 0 { // Avoid pruning if PruneAfter is zero or not set
+	if params.PruneAfter == 0 || params.EpochLength == 0 {
 		return nil
 	}
 	if currentEpoch <= params.PruneAfter {
 		return nil
 	}
 
-	pruneBeforeEpoch := currentEpoch - params.PruneAfter
-	pruneHeight := int64(pruneBeforeEpoch * params.EpochLength) // Assuming EpochLength defines blocks per epoch
+	firstRetainedEpoch := currentEpoch - params.PruneAfter
+	return &attestationRetentionBoundary{
+		firstRetainedEpoch:  firstRetainedEpoch,
+		firstRetainedHeight: int64(firstRetainedEpoch * params.EpochLength),
+	}
+}
+
+// PruneOldBitmaps removes attestation state older than PruneAfter epochs.
+func (k Keeper) PruneOldBitmaps(ctx sdk.Context, currentEpoch uint64) error {
+	boundary := k.attestationRetentionBoundary(ctx, currentEpoch)
+	if boundary == nil {
+		return nil
+	}
 
 	// Prune attestation bitmaps (raw bitmaps)
-	attestationRange := new(collections.Range[int64]).StartInclusive(0).EndExclusive(pruneHeight)
+	attestationRange := new(collections.Range[int64]).StartInclusive(0).EndExclusive(boundary.firstRetainedHeight)
 	if err := k.AttestationBitmap.Clear(ctx, attestationRange); err != nil {
-		return fmt.Errorf("clearing attestation bitmaps before height %d: %w", pruneHeight, err)
+		return fmt.Errorf("clearing attestation bitmaps before height %d: %w", boundary.firstRetainedHeight, err)
 	}
 	// Prune stored attestation info (full AttestationBitmap objects)
-	storedAttestationInfoRange := new(collections.Range[int64]).StartInclusive(0).EndExclusive(pruneHeight)
+	storedAttestationInfoRange := new(collections.Range[int64]).StartInclusive(0).EndExclusive(boundary.firstRetainedHeight)
 	if err := k.StoredAttestationInfo.Clear(ctx, storedAttestationInfoRange); err != nil {
-		return fmt.Errorf("clearing stored attestation info before height %d: %w", pruneHeight, err)
+		return fmt.Errorf("clearing stored attestation info before height %d: %w", boundary.firstRetainedHeight, err)
 	}
 
 	// Prune epoch bitmaps
-	epochRange := new(collections.Range[uint64]).StartInclusive(0).EndExclusive(pruneBeforeEpoch)
+	epochRange := new(collections.Range[uint64]).StartInclusive(0).EndExclusive(boundary.firstRetainedEpoch)
 	if err := k.EpochBitmap.Clear(ctx, epochRange); err != nil {
-		return fmt.Errorf("clearing epoch bitmaps before epoch %d: %w", pruneBeforeEpoch, err)
+		return fmt.Errorf("clearing epoch bitmaps before epoch %d: %w", boundary.firstRetainedEpoch, err)
 	}
 
-	// TODO: Consider pruning signatures associated with pruned heights.
-	// This would involve iterating k.Signatures and removing entries where height < pruneHeight.
+	signatureRange := new(collections.Range[collections.Pair[int64, string]]).
+		EndExclusive(collections.Join(boundary.firstRetainedHeight, ""))
+	if err := k.Signatures.Clear(ctx, signatureRange); err != nil {
+		return fmt.Errorf("clearing signatures before height %d: %w", boundary.firstRetainedHeight, err)
+	}
 
-	k.Logger(ctx).Info("Pruned old bitmaps and attestation info", "prunedBeforeEpoch", pruneBeforeEpoch, "prunedBeforeHeight", pruneHeight)
+	k.Logger(ctx).Info("Pruned old attestation state", "prunedBeforeEpoch", boundary.firstRetainedEpoch, "prunedBeforeHeight", boundary.firstRetainedHeight)
 	return nil
 }
 

modules/network/keeper/msg_server.go

@@ -60,21 +60,14 @@ func (k msgServer) Attest(goCtx context.Context, msg *types.MsgAttest) (*types.M
 		return nil, sdkerr.Wrapf(sdkerrors.ErrInvalidRequest, "attestation height %d exceeds max allowed height %d", msg.Height, maxFutureHeight)
 	}
 
-	// Enforce attestation height lower bound: reject heights that fall below
-	// the PruneAfter retention window. Attesting pruned/about-to-be-pruned
-	// heights wastes storage and serves no purpose. This uses the same epoch
-	// calculation as PruneOldBitmaps so the two stay aligned.
-	params := k.GetParams(ctx)
-	minHeight := int64(1)
-	if params.PruneAfter > 0 && params.EpochLength > 0 {
-		currentEpoch := uint64(currentHeight) / params.EpochLength
-		if currentEpoch > params.PruneAfter {
-			minHeight = int64((currentEpoch - params.PruneAfter) * params.EpochLength)
+	// Txs run before EndBlocker, so the current epoch's pruning has not run yet.
+	currentEpoch := k.GetCurrentEpoch(ctx)
+	if currentEpoch > 0 {
+		boundary := k.attestationRetentionBoundary(ctx, currentEpoch-1)
+		if boundary != nil && boundary.prunesHeight(msg.Height) {
+			return nil, sdkerr.Wrapf(sdkerrors.ErrInvalidRequest, "attestation height %d is below retention window (min %d)", msg.Height, boundary.firstRetainedHeight)
 		}
 	}
-	if msg.Height < minHeight {
-		return nil, sdkerr.Wrapf(sdkerrors.ErrInvalidRequest, "attestation height %d is below retention window (min %d)", msg.Height, minHeight)
-	}
 	bitmap, err := k.GetAttestationBitmap(ctx, msg.Height)
 	if err != nil && !errors.Is(err, collections.ErrNotFound) {
 		return nil, sdkerr.Wrap(err, "get attestation bitmap")

modules/network/keeper/msg_server_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/log"
 	"cosmossdk.io/math"
 	storetypes "cosmossdk.io/store/types"
@@ -359,12 +360,17 @@ func newTestServer(t *testing.T, sk *MockStakingKeeper) (msgServer, Keeper, sdk.
 func TestAttestHeightBounds(t *testing.T) {
 	myValAddr := sdk.ValAddress("validator1")
 	ownerAddr := sdk.ValAddress("attester_owner")
-	// With DefaultParams: EpochLength=1, PruneAfter=15
-	// At blockHeight=100: currentEpoch=100, minHeight=(100-7)*1=93
+	shortRetentionParams := types.DefaultParams()
+	shortRetentionParams.PruneAfter = 15
+	epochRetentionParams := types.DefaultParams()
+	epochRetentionParams.EpochLength = 10
+	epochRetentionParams.PruneAfter = 2
+
 	specs := map[string]struct {
-		blockHeight int64
-		attestH     int64
-		expErr      error
+		blockHeight    int64
+		attestH        int64
+		paramsOverride *types.Params
+		expErr         error
 	}{
 		"future height rejected": {
 			blockHeight: 100,
@@ -385,18 +391,30 @@ func TestAttestHeightBounds(t *testing.T) {
 			attestH:     101,
 		},
 		"stale height rejected": {
-			blockHeight: 100,
-			attestH:     1,
-			expErr:      sdkerrors.ErrInvalidRequest,
+			blockHeight:    100,
+			attestH:        1,
+			paramsOverride: &shortRetentionParams,
+			expErr:         sdkerrors.ErrInvalidRequest,
 		},
 		"below retention window rejected": {
-			blockHeight: 100,
-			attestH:     84, // minHeight = 85
-			expErr:      sdkerrors.ErrInvalidRequest,
+			blockHeight:    100,
+			attestH:        83, // minHeight = 84
+			paramsOverride: &shortRetentionParams,
+			expErr:         sdkerrors.ErrInvalidRequest,
 		},
 		"at retention boundary accepted": {
-			blockHeight: 100,
-			attestH:     93, // exactly minHeight
+			blockHeight:    100,
+			attestH:        84,
+			paramsOverride: &shortRetentionParams,
+		},
+		"retained checkpoint accepted before next epoch pruning runs": {
+			blockHeight:    50,
+			attestH:        20,
+			paramsOverride: &epochRetentionParams,
+		},
+		"default retention keeps IBC handshake history": {
+			blockHeight: 139,
+			attestH:     25,
 		},
 		"early chain no stale rejection": {
 			blockHeight: 16,
@@ -419,7 +437,11 @@ func TestAttestHeightBounds(t *testing.T) {
 				Height:  spec.blockHeight,
 			}, false, logger).WithContext(t.Context())
 
-			require.NoError(t, keeper.SetParams(ctx, types.DefaultParams()))
+			params := types.DefaultParams()
+			if spec.paramsOverride != nil {
+				params = *spec.paramsOverride
+			}
+			require.NoError(t, keeper.SetParams(ctx, params))
 
 			joinMsg := &types.MsgJoinAttesterSet{
 				Authority:        ownerAddr.String(),
@@ -448,6 +470,109 @@ func TestAttestHeightBounds(t *testing.T) {
 	}
 }
 
+func TestPruneOldBitmapsRemovesAllAttestationStateBelowRetentionWindow(t *testing.T) {
+	sk := NewMockStakingKeeper()
+	_, keeper, ctx := newTestServer(t, &sk)
+
+	params := types.DefaultParams()
+	params.EpochLength = 10
+	params.PruneAfter = 2
+	require.NoError(t, keeper.SetParams(ctx, params))
+
+	oldHeight := int64(19)
+	boundaryHeight := int64(20)
+	oldEpoch := uint64(1)
+	boundaryEpoch := uint64(2)
+	attester := sdk.ValAddress("validator1").String()
+
+	require.NoError(t, keeper.SetAttestationBitmap(ctx, oldHeight, []byte{0x01}))
+	require.NoError(t, keeper.StoredAttestationInfo.Set(ctx, oldHeight, types.AttestationBitmap{
+		Height: oldHeight,
+		Bitmap: []byte{0x01},
+	}))
+	require.NoError(t, keeper.SetEpochBitmap(ctx, oldEpoch, []byte{0x01}))
+	require.NoError(t, keeper.SetSignature(ctx, oldHeight, attester, []byte("old-signature")))
+
+	require.NoError(t, keeper.SetAttestationBitmap(ctx, boundaryHeight, []byte{0x02}))
+	require.NoError(t, keeper.StoredAttestationInfo.Set(ctx, boundaryHeight, types.AttestationBitmap{
+		Height: boundaryHeight,
+		Bitmap: []byte{0x02},
+	}))
+	require.NoError(t, keeper.SetEpochBitmap(ctx, boundaryEpoch, []byte{0x02}))
+	require.NoError(t, keeper.SetSignature(ctx, boundaryHeight, attester, []byte("boundary-signature")))
+
+	require.NoError(t, keeper.PruneOldBitmaps(ctx, 4))
+
+	_, err := keeper.GetAttestationBitmap(ctx, oldHeight)
+	require.ErrorIs(t, err, collections.ErrNotFound)
+	_, err = keeper.StoredAttestationInfo.Get(ctx, oldHeight)
+	require.ErrorIs(t, err, collections.ErrNotFound)
+	require.Nil(t, keeper.GetEpochBitmap(ctx, oldEpoch))
+	hasOldSignature, err := keeper.HasSignature(ctx, oldHeight, attester)
+	require.NoError(t, err)
+	require.False(t, hasOldSignature)
+
+	bitmap, err := keeper.GetAttestationBitmap(ctx, boundaryHeight)
+	require.NoError(t, err)
+	require.Equal(t, []byte{0x02}, bitmap)
+	_, err = keeper.StoredAttestationInfo.Get(ctx, boundaryHeight)
+	require.NoError(t, err)
+	require.Equal(t, []byte{0x02}, keeper.GetEpochBitmap(ctx, boundaryEpoch))
+	hasBoundarySignature, err := keeper.HasSignature(ctx, boundaryHeight, attester)
+	require.NoError(t, err)
+	require.True(t, hasBoundarySignature)
+}
+
+func TestAttestationRetentionBoundary(t *testing.T) {
+	sk := NewMockStakingKeeper()
+	_, keeper, ctx := newTestServer(t, &sk)
+
+	params := types.DefaultParams()
+	params.EpochLength = 10
+	params.PruneAfter = 2
+	require.NoError(t, keeper.SetParams(ctx, params))
+
+	require.Nil(t, keeper.attestationRetentionBoundary(ctx, 2))
+
+	boundary := keeper.attestationRetentionBoundary(ctx, 4)
+	require.NotNil(t, boundary)
+	require.Equal(t, uint64(2), boundary.firstRetainedEpoch)
+	require.Equal(t, int64(20), boundary.firstRetainedHeight)
+	require.True(t, boundary.prunesHeight(19))
+	require.False(t, boundary.prunesHeight(20))
+}
+
+func TestEndBlockerPrunesAttestationStateOnEpochBoundary(t *testing.T) {
+	sk := NewMockStakingKeeper()
+	_, keeper, ctx := newTestServer(t, &sk)
+
+	params := types.DefaultParams()
+	params.EpochLength = 10
+	params.PruneAfter = 2
+	require.NoError(t, keeper.SetParams(ctx, params))
+
+	ctx = ctx.WithBlockHeight(49)
+	attester := sdk.ValAddress("validator1").String()
+	require.NoError(t, keeper.SetAttestationBitmap(ctx, 19, []byte{0x01}))
+	require.NoError(t, keeper.SetSignature(ctx, 19, attester, []byte("old-signature")))
+	require.NoError(t, keeper.SetAttestationBitmap(ctx, 20, []byte{0x02}))
+	require.NoError(t, keeper.SetSignature(ctx, 20, attester, []byte("boundary-signature")))
+
+	require.NoError(t, keeper.EndBlocker(ctx))
+
+	_, err := keeper.GetAttestationBitmap(ctx, 19)
+	require.ErrorIs(t, err, collections.ErrNotFound)
+	hasOldSignature, err := keeper.HasSignature(ctx, 19, attester)
+	require.NoError(t, err)
+	require.False(t, hasOldSignature)
+
+	_, err = keeper.GetAttestationBitmap(ctx, 20)
+	require.NoError(t, err)
+	hasBoundarySignature, err := keeper.HasSignature(ctx, 20, attester)
+	require.NoError(t, err)
+	require.True(t, hasBoundarySignature)
+}
+
 var _ types.StakingKeeper = &MockStakingKeeper{}
 
 type MockStakingKeeper struct {

modules/network/types/params.go

@@ -11,7 +11,7 @@ var (
 	DefaultEpochLength      = uint64(1)                         // Changed from 10 to 1 to allow attestations on every block
 	DefaultQuorumFraction   = math.LegacyNewDecWithPrec(667, 3) // 2/3
 	DefaultMinParticipation = math.LegacyNewDecWithPrec(5, 1)   // 1/2
-	DefaultPruneAfter       = uint64(15)                        // also used as number of blocks for attestations to land
+	DefaultPruneAfter       = uint64(1000)
 	DefaultSignMode         = SignMode_SIGN_MODE_CHECKPOINT
 )