feat(x/batching): implementation of batching module pruning

We add a new map from batch number to its data results to x/batching state in order to support simple pruning. In simple pruning we remove a batch and all its associated data if necessary to keep the number of batches at the module parameter NumBatchesToKeep.

Author: hacheigriega

app/upgrades/mainnet/v1.0.7/constants.go

@@ -0,0 +1,55 @@
+package v1
+
+import (
+	"context"
+	_ "embed"
+
+	storetypes "cosmossdk.io/store/types"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	"github.com/sedaprotocol/seda-chain/app/keepers"
+	"github.com/sedaprotocol/seda-chain/app/upgrades"
+)
+
+const (
+	UpgradeName = "v" // TODO Update name and register this handler.
+)
+
+var Upgrade = upgrades.Upgrade{
+	UpgradeName:          UpgradeName,
+	CreateUpgradeHandler: CreateUpgradeHandler,
+	StoreUpgrades: storetypes.StoreUpgrades{
+		Added:   []string{},
+		Deleted: []string{},
+	},
+}
+
+func CreateUpgradeHandler(
+	mm upgrades.ModuleManager,
+	configurator module.Configurator,
+	keepers *keepers.AppKeepers,
+) upgradetypes.UpgradeHandler {
+	return func(context context.Context, _ upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		ctx := sdk.UnwrapSDKContext(context)
+
+		// Run module migrations.
+		migrations, err := mm.RunMigrations(ctx, configurator, fromVM)
+		if err != nil {
+			return nil, err
+		}
+
+		err = keepers.BatchingKeeper.SetBatchNumberAtUpgrade(ctx)
+		if err != nil {
+			return nil, err
+		}
+		err = keepers.BatchingKeeper.SetHasPruningCaughtUp(ctx, false)
+		if err != nil {
+			return nil, err
+		}
+
+		return migrations, nil
+	}
+}

proto/sedachain/batching/v1/batching.proto

@@ -86,6 +86,19 @@ message DataResult {
   bool consensus = 12 [ (gogoproto.jsontag) = "consensus" ];
 }
 
+// DataRequestIDHeights is a collection of DataRequestIDHeight objects.
+message DataRequestIDHeights {
+  repeated DataRequestIDHeight data_request_id_heights = 1 [ (gogoproto.nullable) = false ];
+}
+
+// DataRequestIDHeight is a pair of data request ID and its posted height.
+message DataRequestIDHeight {
+  // DataRequestID is the hex-encoded data request ID.
+  string data_request_id = 1;
+  // DataRequestHeight is the height at which the data request was submitted.
+  uint64 data_request_height = 2;
+}
+
 // Params defines the parameters for the batching module.
 message Params {
   // NumBatchesToKeep is the number of batches to keep in the state without
@@ -94,7 +107,7 @@ message Params {
   // MaxBatchPrunePerBlock is the maximum number of batches to prune per
   // block.
   uint64 max_batch_prune_per_block = 2;
-  // MaxDataResultsToCheckForPrune is the maximum number of data results to
+  // MaxDataResultPrunePerBlock is the maximum number of data results to
   // check for pruning per block.
-  uint64 max_data_results_to_check_for_prune = 3;
+  uint64 max_data_result_prune_per_block = 3;
 }

proto/sedachain/batching/v1/genesis.proto

@@ -17,6 +17,9 @@ message GenesisState {
   repeated BatchAssignment batch_assignments = 5
       [ (gogoproto.nullable) = false ];
   Params params = 6 [ (gogoproto.nullable) = false ];
+  repeated GenesisDataResult legacy_data_results = 7 [ (gogoproto.nullable) = false ];
+  bool has_pruning_caught_up = 8;
+  uint64 batch_number_at_upgrade = 9;
 }
 
 // BatchAssignment represents a batch assignment for genesis export

x/batching/README.md

@@ -5,14 +5,19 @@ The batching module collects data rseults, current validators, and their signatu
 
 ## State
 ```
-0x00 | is_batched | data_request_id | data_request_height -> data_result
+0x00 | is_batched | data_request_id | data_request_height -> legacy_data_results
 0x01 | data_request_id | data_request_height              -> batch_number
 0x02                                                      -> current_batch_number
 0x03 | block_height                                       -> batch
 0x04 | batch_number                                       -> batch
 0x05 | batch_number | validator_address                   -> validator_tree_entries
 0x06 | batch_number                                       -> data_tree_entries
 0x07 | batch_number | validator_address                   -> batch_signature
+0x08                                                      -> parameters
+0x09 | is_batched | data_request_id | data_request_height -> data_result
+0x10 | batch_number                                       -> data_request_id | data_request_height
+0x11                                                      -> batch_number_at_upgrade
+0x12                                                      -> has_pruning_caught_up
 ```
 
 ### Batches

x/batching/keeper/batch.go

@@ -34,47 +34,48 @@ func (k Keeper) setBatch(ctx context.Context, batch types.Batch) error {
 	return k.batches.Set(ctx, batch.BlockHeight, batch)
 }
 
-// SetNewBatch increments the current batch number and stores a given
-// batch at that index. It also stores the given data result tree
-// entries, validator tree entries, and batch signature entries (at
-// the next batch index, to be populated with signatures later). It
-// returns an error if a batch already exists at the given batch's
-// block height or if the given batch's batch number does not match
-// the next batch number.
-func (k Keeper) SetNewBatch(ctx context.Context, batch types.Batch, dataEntries types.DataResultTreeEntries, valEntries []types.ValidatorTreeEntry) error {
+// SetNewBatch stores a new batch and its associated data at current batch number
+// and increments the current batch number. If successful, it returns the batch number
+// of the newly created batch.
+func (k Keeper) SetNewBatch(ctx sdk.Context, batch types.Batch, dataEntries types.DataResultTreeEntries, valEntries []types.ValidatorTreeEntry) (uint64, error) {
 	found, err := k.batches.Has(ctx, batch.BlockHeight)
 	if err != nil {
-		return err
+		return 0, err
 	}
 	if found {
-		return types.ErrBatchAlreadyExists.Wrapf("batch block height %d", batch.BlockHeight)
+		return 0, types.ErrBatchAlreadyExists.Wrapf("batch block height %d", batch.BlockHeight)
 	}
 
 	batchNum, err := k.GetCurrentBatchNum(ctx)
 	if err != nil {
-		return err
+		return 0, err
 	}
 	if batch.BatchNumber != batchNum {
-		return types.ErrInvalidBatchNumber.Wrapf("got %d; expected %d", batch.BatchNumber, batchNum)
+		return 0, types.ErrInvalidBatchNumber.Wrapf("got %d; expected %d", batch.BatchNumber, batchNum)
 	}
 
 	err = k.setDataResultTreeEntry(ctx, batchNum, dataEntries)
 	if err != nil {
-		return err
+		return 0, err
 	}
 
 	for _, valEntry := range valEntries {
 		err = k.setValidatorTreeEntry(ctx, batchNum, valEntry)
 		if err != nil {
-			return err
+			return 0, err
 		}
 	}
 
 	_, err = k.incrementCurrentBatchNum(ctx)
 	if err != nil {
-		return err
+		return 0, err
 	}
-	return k.batches.Set(ctx, batch.BlockHeight, batch)
+
+	err = k.batches.Set(ctx, batch.BlockHeight, batch)
+	if err != nil {
+		return 0, err
+	}
+	return batchNum, nil
 }
 
 func (k Keeper) GetBatchForHeight(ctx context.Context, blockHeight int64) (types.Batch, error) {

x/batching/keeper/batch_assignments.go

@@ -0,0 +1,59 @@
+package keeper
+
+import (
+	"context"
+	"errors"
+
+	"cosmossdk.io/collections"
+
+	"github.com/sedaprotocol/seda-chain/x/batching/types"
+)
+
+// SetBatchAssignment stores mapping between data request ID - posted height pair
+// and assigned batch number in both directions.
+func (k Keeper) SetBatchAssignment(ctx context.Context, dataReqID string, dataReqHeight, batchNumber uint64) error {
+	items, err := k.batchDataResults.Get(ctx, batchNumber)
+	if err != nil {
+		if !errors.Is(err, collections.ErrNotFound) {
+			return err
+		}
+		items.DataRequestIdHeights = make([]types.DataRequestIDHeight, 0)
+	}
+	items.DataRequestIdHeights = append(items.DataRequestIdHeights, types.DataRequestIDHeight{
+		DataRequestId:     dataReqID,
+		DataRequestHeight: dataReqHeight,
+	})
+	err = k.batchDataResults.Set(ctx, batchNumber, items)
+	if err != nil {
+		return err
+	}
+
+	return k.batchAssignments.Set(ctx, collections.Join(dataReqID, dataReqHeight), batchNumber)
+}
+
+func (k Keeper) GetBatchAssignment(ctx context.Context, dataReqID string, dataReqHeight uint64) (uint64, error) {
+	return k.batchAssignments.Get(ctx, collections.Join(dataReqID, dataReqHeight))
+}
+
+func (k Keeper) RemoveBatchAssignment(ctx context.Context, dataReqID string, dataReqHeight uint64) error {
+	return k.batchAssignments.Remove(ctx, collections.Join(dataReqID, dataReqHeight))
+}
+
+// getAllBatchAssignments retrieves all batch assignments from the store.
+// Used for genesis export.
+func (k Keeper) getAllBatchAssignments(ctx context.Context) ([]types.BatchAssignment, error) {
+	var batchAssignments []types.BatchAssignment
+	err := k.batchAssignments.Walk(ctx, nil, func(key collections.Pair[string, uint64], value uint64) (stop bool, err error) {
+		batchAssignments = append(batchAssignments, types.BatchAssignment{
+			BatchNumber:       value,
+			DataRequestId:     key.K1(),
+			DataRequestHeight: key.K2(),
+		})
+		return false, nil
+	})
+	return batchAssignments, err
+}
+
+func (k Keeper) RemoveBatchDataResults(ctx context.Context, batchNumber uint64) error {
+	return k.batchDataResults.Remove(ctx, batchNumber)
+}

x/batching/keeper/benchmark_endblock_test.go

@@ -21,16 +21,17 @@ func BenchmarkBatchPruning(b *testing.B) {
 		require.NoError(b, err)
 		batch, dataEntries, valEntries, err := f.batchingKeeper.ConstructBatch(f.Context())
 		require.NoError(b, err)
-		err = f.batchingKeeper.SetNewBatch(f.Context(), batch, dataEntries, valEntries)
+		_, err = f.batchingKeeper.SetNewBatch(f.Context(), batch, dataEntries, valEntries)
 		require.NoError(b, err)
 	}
 
 	for b.Loop() {
-		_, err := f.batchingKeeper.PruneBatches(f.Context(), numBatchesToKeep, maxBatchPrunePerBlock)
+		_, err := f.batchingKeeper.BatchPruneBatches(f.Context(), numBatchesToKeep, maxBatchPrunePerBlock)
 		require.NoError(b, err)
 	}
 }
 
+/*
 func BenchmarkDataResultPruning(b *testing.B) {
 	f := initFixture(b)
 
@@ -53,7 +54,8 @@ func BenchmarkDataResultPruning(b *testing.B) {
 		f.AddBlock()
 		f.SetRandomLastCommitHash()
 
-		err := f.batchingKeeper.PruneDataResults(f.Context(), maxDataResultsToCheckForPrune, 2000)
+		err := f.batchingKeeper.BatchPruneDataResults(f.Context(), maxDataResultsToCheckForPrune, 2000)
 		require.NoError(b, err)
 	}
 }
+*/