Pruning data results

proto/sedachain/batching/v1/batching.proto

@@ -94,4 +94,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
+  // check for pruning per block.
+  uint64 max_data_results_to_check_for_prune = 3;
 }

proto/sedachain/batching/v1/genesis.proto

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

proto/sedachain/batching/v1/tx.proto

@@ -0,0 +1,28 @@
+syntax = "proto3";
+package sedachain.batching.v1;
+
+import "cosmos/msg/v1/msg.proto";
+import "gogoproto/gogo.proto";
+import "cosmos_proto/cosmos.proto";
+import "sedachain/batching/v1/batching.proto";
+
+option go_package = "github.com/sedaprotocol/seda-chain/x/batching/types";
+
+// Msg service defines the batching tx gRPC methods.
+service Msg {
+  // The UpdateParams method updates the module's parameters.
+  rpc UpdateParams(MsgUpdateParams) returns (MsgUpdateParamsResponse);
+}
+
+// The request message for the UpdateParams method.
+message MsgUpdateParams {
+  option (cosmos.msg.v1.signer) = "authority";
+
+  // Authority is the address that controls the module (defaults to x/gov unless
+  // overwritten).
+  string authority = 1 [ (cosmos_proto.scalar) = "cosmos.AddressString" ];
+  Params params = 2 [ (gogoproto.nullable) = false ];
+}
+
+// The response message for the UpdateParams method.
+message MsgUpdateParamsResponse {}

testutil/integration.go

@@ -1,6 +1,7 @@
 package testutil
 
 import (
+	"crypto/rand"
 	"fmt"
 	"time"
 
@@ -206,10 +207,23 @@ func (app *IntegationApp) AddTime(seconds int64) {
 }
 
 // AddBlock increments the block number of the application context.
+// It also sets the last commit hash to a random value.
 func (app *IntegationApp) AddBlock() {
 	app.ctx = app.ctx.WithBlockHeight(app.ctx.BlockHeader().Height + 1)
 }
 
+func (app *IntegationApp) SetRandomLastCommitHash() {
+	randomBytes := make([]byte, 32)
+	_, err := rand.Read(randomBytes)
+	if err != nil {
+		panic(err)
+	}
+
+	newHeader := app.ctx.BlockHeader()
+	newHeader.LastCommitHash = randomBytes
+	app.ctx = app.ctx.WithBlockHeader(newHeader)
+}
+
 // QueryHelper returns the application query helper.
 // It can be used when registering query services.
 func (app *IntegationApp) QueryHelper() *baseapp.QueryServiceTestHelper {

x/batching/keeper/benchmark_endblock_test.go

@@ -1,11 +1,8 @@
 package keeper_test
 
 import (
-	"os"
-	"runtime/pprof"
 	"testing"
 
-	"github.com/sedaprotocol/seda-chain/x/batching/types"
 	"github.com/stretchr/testify/require"
 )
 
@@ -17,12 +14,6 @@ func BenchmarkBatchPruning(b *testing.B) {
 	numBatchesToKeep := uint64(1000)
 	maxBatchPrunePerBlock := uint64(100)
 
-	err := f.batchingKeeper.SetParams(f.Context(), types.Params{
-		NumBatchesToKeep:      numBatchesToKeep,
-		MaxBatchPrunePerBlock: maxBatchPrunePerBlock,
-	})
-	require.NoError(b, err)
-
 	for range numBatches {
 		f.AddBlock()
 
@@ -34,16 +25,35 @@ func BenchmarkBatchPruning(b *testing.B) {
 		require.NoError(b, err)
 	}
 
-	cpuFile, err := os.Create("cpu_refactor_new.out")
-	require.NoError(b, err)
-	defer cpuFile.Close()
+	for b.Loop() {
+		_, err := f.batchingKeeper.PruneBatches(f.Context(), numBatchesToKeep, maxBatchPrunePerBlock)
+		require.NoError(b, err)
+	}
+}
+
+func BenchmarkDataResultPruning(b *testing.B) {
+	f := initFixture(b)
 
-	err = pprof.StartCPUProfile(cpuFile)
-	require.NoError(b, err)
-	defer pprof.StopCPUProfile()
+	maxDataResultsToCheckForPrune := uint64(100)
+
+	// Create 10 data results for each of 1000 batches
+	for i := range uint64(100) {
+		f.AddBlock()
+
+		dataResults := generateDataResults(b, 10)
+		for _, dataResult := range dataResults {
+			err := f.batchingKeeper.SetDataResultForBatching(f.Context(), dataResult)
+			require.NoError(b, err)
+			err = f.batchingKeeper.MarkDataResultAsBatched(f.Context(), dataResult, i)
+			require.NoError(b, err)
+		}
+	}
 
 	for b.Loop() {
-		err = f.batchingKeeper.PruneBatches(f.Context())
+		f.AddBlock()
+		f.SetRandomLastCommitHash()
+
+		err := f.batchingKeeper.PruneDataResults(f.Context(), maxDataResultsToCheckForPrune, 2000)
 		require.NoError(b, err)
 	}
 }

x/batching/keeper/data_result.go

@@ -15,6 +15,11 @@ func (k Keeper) SetDataResultForBatching(ctx context.Context, result types.DataR
 	return k.dataResults.Set(ctx, collections.Join3(false, result.DrId, result.DrBlockHeight), result)
 }
 
+// RemoveDataResult removes a data result from the store.
+func (k Keeper) RemoveDataResult(ctx context.Context, batched bool, dataReqID string, dataReqHeight uint64) error {
+	return k.dataResults.Remove(ctx, collections.Join3(batched, dataReqID, dataReqHeight))
+}
+
 // MarkDataResultAsBatched removes the "unbatched" variant of the given
 // data result and stores a "batched" variant.
 func (k Keeper) MarkDataResultAsBatched(ctx context.Context, result types.DataResult, batchNum uint64) error {
@@ -147,6 +152,10 @@ func (k Keeper) GetBatchAssignment(ctx context.Context, dataReqID string, dataRe
 	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) {

x/batching/keeper/endblock.go

@@ -10,6 +10,7 @@ import (
 	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
 
+	"github.com/cosmos/cosmos-sdk/telemetry"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	"github.com/sedaprotocol/seda-chain/app/utils"
@@ -24,97 +25,44 @@ func (k Keeper) EndBlock(ctx sdk.Context) error {
 	if err != nil {
 		return err
 	}
-	if !isActivated {
-		k.Logger(ctx).Info("skip batching since proving scheme has not been activated", "index", sedatypes.SEDAKeyIndexSecp256k1)
-		return nil
-	}
 
-	batch, dataEntries, valEntries, err := k.ConstructBatch(ctx)
-	if err != nil {
-		if errors.Is(err, types.ErrNoBatchingUpdate) {
+	if isActivated {
+		batch, dataEntries, valEntries, err := k.ConstructBatch(ctx)
+		if err != nil {
+			if !errors.Is(err, types.ErrNoBatchingUpdate) {
+				return err
+			}
 			k.Logger(ctx).Info("skip batch creation due to no update", "height", ctx.BlockHeight())
-			return nil
+		} else {
+			err = k.SetNewBatch(ctx, batch, dataEntries, valEntries)
+			if err != nil {
+				return err
+			}
 		}
-		return err
-	}
-
-	err = k.SetNewBatch(ctx, batch, dataEntries, valEntries)
-	if err != nil {
-		return err
+	} else {
+		k.Logger(ctx).Info("skip batching since proving scheme has not been activated", "index", sedatypes.SEDAKeyIndexSecp256k1)
 	}
 
-	return k.PruneBatches(ctx)
-}
-
-// PruneBatches prunes batches and their associated data based on module
-// parameters NumBatchesToKeep and MaxBatchPrunePerBlock.
-func (k Keeper) PruneBatches(ctx sdk.Context) error {
 	params, err := k.GetParams(ctx)
 	if err != nil {
 		return err
 	}
 
-	// Note the current batch number here has not been used yet.
-	currentBatchNum, err := k.GetCurrentBatchNum(ctx)
-	if err != nil {
-		return err
-	}
-	firstBatchNum, err := k.firstBatchNumber.Get(ctx)
+	lastPrunedBatchNum, err := k.PruneBatches(ctx, params.NumBatchesToKeep, params.MaxBatchPrunePerBlock)
 	if err != nil {
-		return err
-	}
-	if currentBatchNum-firstBatchNum <= params.NumBatchesToKeep {
-		k.Logger(ctx).Info("skip batch pruning", "current_batch_num", currentBatchNum, "first_batch_num", firstBatchNum)
+		telemetry.SetGauge(1, types.TelemetryKeyBatchingPruningFail)
+		k.Logger(ctx).Error("error while pruning batches", "err", err)
 		return nil
 	}
 
-	// Prune range is [firstBatchNum, newFirstBatchNum)
-	firstBatchHeight, err := k.batches.Indexes.Number.MatchExact(ctx, firstBatchNum)
-	if err != nil {
-		return err
-	}
-
-	newFirstBatchNum := min(firstBatchNum+params.MaxBatchPrunePerBlock, currentBatchNum-params.NumBatchesToKeep)
-	newFirstBatchHeight, err := k.batchIndex.Get(ctx, newFirstBatchNum)
+	err = k.PruneDataResults(ctx, params.MaxDataResultsToCheckForPrune, lastPrunedBatchNum)
 	if err != nil {
-		return err
-	}
-
-	// Clear batches and their associated data.
-	batchNumRng := new(collections.Range[uint64]).StartInclusive(firstBatchNum).EndExclusive(newFirstBatchNum)
-	err = k.batchIndex.Clear(ctx, batchNumRng)
-	if err != nil {
-		return err
-	}
-	err = k.dataResultTreeEntries.Clear(ctx, batchNumRng)
-	if err != nil {
-		return err
-	}
-
-	batchHeightRng := new(collections.Range[int64]).StartInclusive(firstBatchHeight).EndExclusive(newFirstBatchHeight)
-	err = k.batchesMap.Clear(ctx, batchHeightRng)
-	if err != nil {
-		return err
-	}
-
-	valRng := new(collections.Range[collections.Pair[uint64, []byte]]).
-		StartInclusive(collections.PairPrefix[uint64, []byte](firstBatchNum)).
-		EndExclusive(collections.PairPrefix[uint64, []byte](newFirstBatchNum))
-	err = k.validatorTreeEntries.Clear(ctx, valRng)
-	if err != nil {
-		return err
-	}
-	err = k.batchSignatures.Clear(ctx, valRng)
-	if err != nil {
-		return err
-	}
-
-	err = k.firstBatchNumber.Set(ctx, newFirstBatchNum)
-	if err != nil {
-		return err
+		telemetry.SetGauge(1, types.TelemetryKeyBatchingPruningFail)
+		k.Logger(ctx).Error("error while pruning data results", "err", err)
+		return nil
 	}
 
-	k.Logger(ctx).Info("successfully pruned batch data")
+	telemetry.SetGauge(0, types.TelemetryKeyBatchingPruningFail)
 	return nil
 }