heimdallcli: Genesis export command optimization

bor/genesis.go

@@ -28,14 +28,5 @@ func InitGenesis(ctx sdk.Context, keeper Keeper, data types.GenesisState) {
 
 // ExportGenesis returns a GenesisState for a given context and keeper.
 func ExportGenesis(ctx sdk.Context, keeper Keeper) types.GenesisState {
-	params := keeper.GetParams(ctx)
-
-	allSpans := keeper.GetAllSpans(ctx)
-	hmTypes.SortSpanByID(allSpans)
-
-	return types.NewGenesisState(
-		params,
-		// TODO think better way to export all spans
-		allSpans,
-	)
+	return types.NewGenesisState(keeper.GetParams(ctx), nil)
 }

bor/keeper.go

@@ -349,3 +349,43 @@ func (k *Keeper) IterateSpansAndApplyFn(ctx sdk.Context, f func(span hmTypes.Spa
 		}
 	}
 }
+
+// IterateSpansAndCollect iterates over spans, collects up to 'max' entries,
+// and returns a slice containing the collected spans.
+// It continues from the last key processed in the previous batch.
+func (k *Keeper) IterateSpansAndCollect(ctx sdk.Context, nextKey []byte, max int) ([]*hmTypes.Span, []byte, error) {
+	store := ctx.KVStore(k.storeKey)
+
+	var startKey []byte
+	if nextKey != nil {
+		startKey = nextKey
+	} else {
+		startKey = SpanPrefixKey
+	}
+
+	endKey := sdk.PrefixEndBytes(SpanPrefixKey)
+
+	iterator := store.Iterator(startKey, endKey)
+	defer iterator.Close()
+
+	collectedSpans := make([]*hmTypes.Span, 0, max)
+	entriesCollected := 0
+
+	for ; iterator.Valid() && entriesCollected < max; iterator.Next() {
+		var span hmTypes.Span
+		if err := k.cdc.UnmarshalBinaryBare(iterator.Value(), &span); err != nil {
+			k.Logger(ctx).Error("IterateSpansAndCollect | UnmarshalBinaryBare", "error", err)
+			return nil, nil, err
+		}
+
+		collectedSpans = append(collectedSpans, &span)
+		entriesCollected++
+	}
+
+	// We want to return the key after last processed key because the iterator is inclusive for the start key
+	if iterator.Valid() {
+		return collectedSpans, iterator.Key(), nil
+	}
+
+	return collectedSpans, nil, nil
+}

bor/module.go

@@ -20,9 +20,10 @@ import (
 )
 
 var (
-	_ module.AppModule             = AppModule{}
-	_ module.AppModuleBasic        = AppModuleBasic{}
-	_ hmModule.HeimdallModuleBasic = AppModule{}
+	_ module.AppModule                 = AppModule{}
+	_ module.AppModuleBasic            = AppModuleBasic{}
+	_ hmModule.HeimdallModuleBasic     = AppModule{}
+	_ hmModule.StreamedGenesisExporter = AppModule{}
 	// _ module.AppModuleSimulation = AppModule{}
 )
 
@@ -161,3 +162,16 @@ func (am AppModule) NewSideTxHandler() hmTypes.SideTxHandler {
 func (am AppModule) NewPostTxHandler() hmTypes.PostTxHandler {
 	return NewPostTxHandler(am.keeper, am.contractCaller)
 }
+
+// NextGenesisData returns the next chunk of genesis data.
+func (am AppModule) NextGenesisData(ctx sdk.Context, nextKey []byte, max int) (*hmModule.ModuleGenesisData, error) {
+	data, nextKey, err := am.keeper.IterateSpansAndCollect(ctx, nextKey, max)
+	if err != nil {
+		return nil, err
+	}
+	return &hmModule.ModuleGenesisData{
+		Path:    "bor.spans",
+		Data:    types.ModuleCdc.MustMarshalJSON(data),
+		NextKey: nextKey,
+	}, nil
+}

clerk/genesis.go

@@ -24,5 +24,5 @@ func InitGenesis(ctx sdk.Context, keeper Keeper, data types.GenesisState) {
 
 // ExportGenesis returns a GenesisState for a given context and keeper.
 func ExportGenesis(ctx sdk.Context, keeper Keeper) types.GenesisState {
-	return types.NewGenesisState(keeper.GetAllEventRecords(ctx), keeper.GetRecordSequences(ctx))
+	return types.NewGenesisState(nil, keeper.GetRecordSequences(ctx))
 }

clerk/genesis_test.go

@@ -66,5 +66,5 @@ func (suite *GenesisTestSuite) TestInitExportGenesis() {
 	actualParams := clerk.ExportGenesis(ctx, app.ClerkKeeper)
 
 	require.Equal(t, len(recordSequences), len(actualParams.RecordSequences))
-	require.Equal(t, len(eventRecords), len(actualParams.EventRecords))
+	require.Equal(t, 0, len(actualParams.EventRecords))
 }

clerk/keeper.go

@@ -317,3 +317,43 @@ func (k *Keeper) HasRecordSequence(ctx sdk.Context, sequence string) bool {
 	store := ctx.KVStore(k.storeKey)
 	return store.Has(GetRecordSequenceKey(sequence))
 }
+
+// IterateRecordsAndCollect iterates over EventRecords, collects up to 'max' entries,
+// and returns a slice containing the collected records.
+// It continues from the last key processed in the previous batch.
+func (k *Keeper) IterateRecordsAndCollect(ctx sdk.Context, nextKey []byte, max int) ([]*types.EventRecord, []byte, error) {
+	store := ctx.KVStore(k.storeKey)
+
+	var startKey []byte
+	if nextKey != nil {
+		startKey = nextKey
+	} else {
+		startKey = StateRecordPrefixKey
+	}
+
+	endKey := sdk.PrefixEndBytes(StateRecordPrefixKey)
+
+	iterator := store.Iterator(startKey, endKey)
+	defer iterator.Close()
+
+	collectedRecords := make([]*types.EventRecord, 0, max)
+	entriesCollected := 0
+
+	for ; iterator.Valid() && entriesCollected < max; iterator.Next() {
+		var record types.EventRecord
+		if err := k.cdc.UnmarshalBinaryBare(iterator.Value(), &record); err != nil {
+			k.Logger(ctx).Error("IterateRecordsAndCollect | UnmarshalBinaryBare", "error", err)
+			return nil, nil, err
+		}
+
+		collectedRecords = append(collectedRecords, &record)
+		entriesCollected++
+	}
+
+	// We want to return the key after last processed key because the iterator is inclusive for the start key
+	if iterator.Valid() {
+		return collectedRecords, iterator.Key(), nil
+	}
+
+	return collectedRecords, nil, nil
+}

clerk/module.go

@@ -20,9 +20,10 @@ import (
 )
 
 var (
-	_ module.AppModule             = AppModule{}
-	_ module.AppModuleBasic        = AppModuleBasic{}
-	_ hmModule.HeimdallModuleBasic = AppModule{}
+	_ module.AppModule                 = AppModule{}
+	_ module.AppModuleBasic            = AppModuleBasic{}
+	_ hmModule.HeimdallModuleBasic     = AppModule{}
+	_ hmModule.StreamedGenesisExporter = AppModule{}
 	// _ module.AppModuleSimulation = AppModule{}
 )
 
@@ -141,6 +142,19 @@ func (am AppModule) ExportGenesis(ctx sdk.Context) json.RawMessage {
 	return types.ModuleCdc.MustMarshalJSON(gs)
 }
 
+// NextGenesisData returns the next chunk of genesis data.
+func (am AppModule) NextGenesisData(ctx sdk.Context, nextKey []byte, max int) (*hmModule.ModuleGenesisData, error) {
+	data, nextKey, err := am.keeper.IterateRecordsAndCollect(ctx, nextKey, max)
+	if err != nil {
+		return nil, err
+	}
+	return &hmModule.ModuleGenesisData{
+		Path:    "clerk.event_records",
+		Data:    types.ModuleCdc.MustMarshalJSON(data),
+		NextKey: nextKey,
+	}, nil
+}
+
 // BeginBlock returns the begin blocker for the auth module.
 func (AppModule) BeginBlock(_ sdk.Context, _ abci.RequestBeginBlock) {}