Refactor

maticnetwork · Oct 17, 2024 · bf18af1 · bf18af1
1 parent c362f92
commit bf18af1
Show file tree

Hide file tree

Showing 8 changed files with 61 additions and 58 deletions.
diff --git a/bor/genesis.go b/bor/genesis.go
@@ -28,9 +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)
-	return types.NewGenesisState(
-		params,
-		nil,
-	)
+	return types.NewGenesisState(keeper.GetParams(ctx), nil)
 }
diff --git a/bor/keeper.go b/bor/keeper.go
@@ -353,12 +353,12 @@ 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, lastKey []byte, max int) ([]*hmTypes.Span, []byte, error) {
+func (k *Keeper) IterateSpansAndCollect(ctx sdk.Context, nextKey []byte, max int) ([]*hmTypes.Span, []byte, error) {
 	store := ctx.KVStore(k.storeKey)
 
 	var startKey []byte
-	if lastKey != nil {
-		startKey = lastKey
+	if nextKey != nil {
+		startKey = nextKey
 	} else {
 		startKey = SpanPrefixKey
 	}

diff --git a/bor/module.go b/bor/module.go
@@ -164,14 +164,14 @@ func (am AppModule) NewPostTxHandler() hmTypes.PostTxHandler {
 }
 
 // NextGenesisData returns the next chunk of genesis data.
-func (am AppModule) NextGenesisData(ctx sdk.Context, lastKey []byte) (*hmModule.ModuleGenesisData, error) {
-	data, lastKey, err := am.keeper.IterateSpansAndCollect(ctx, lastKey, 1000)
+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),
-		LastKey: lastKey,
+		NextKey: nextKey,
 	}, nil
 }
diff --git a/clerk/genesis.go b/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([]*types.EventRecord{}, keeper.GetRecordSequences(ctx))
+	return types.NewGenesisState(nil, keeper.GetRecordSequences(ctx))
 }
diff --git a/clerk/keeper.go b/clerk/keeper.go
@@ -321,12 +321,12 @@ func (k *Keeper) HasRecordSequence(ctx sdk.Context, sequence string) bool {
 // 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, lastKey []byte, max int) ([]*types.EventRecord, []byte, error) {
+func (k *Keeper) IterateRecordsAndCollect(ctx sdk.Context, nextKey []byte, max int) ([]*types.EventRecord, []byte, error) {
 	store := ctx.KVStore(k.storeKey)
 
 	var startKey []byte
-	if lastKey != nil {
-		startKey = lastKey
+	if nextKey != nil {
+		startKey = nextKey
 	} else {
 		startKey = StateRecordPrefixKey
 	}

diff --git a/clerk/module.go b/clerk/module.go
@@ -143,15 +143,15 @@ func (am AppModule) ExportGenesis(ctx sdk.Context) json.RawMessage {
 }
 
 // NextGenesisData returns the next chunk of genesis data.
-func (am AppModule) NextGenesisData(ctx sdk.Context, lastKey []byte) (*hmModule.ModuleGenesisData, error) {
-	data, lastKey, err := am.keeper.IterateRecordsAndCollect(ctx, lastKey, 1000)
+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),
-		LastKey: lastKey,
+		NextKey: nextKey,
 	}, nil
 }
 

diff --git a/cmd/heimdallcli/main.go b/cmd/heimdallcli/main.go
@@ -228,39 +228,10 @@ func exportCmd(ctx *server.Context, _ *codec.Codec) *cobra.Command {
 
 			happ := app.NewHeimdallApp(logger, db)
 
-			marshaledAppState := []byte{}
-
-			// Anonymous function just to ensure that appState is not retained in memory
-			func() {
-				appState, err := getAppState(happ)
-				if err != nil {
-					panic(err)
-				}
-				runtime.GC()
-
-				sdkCtx := happ.NewContext(true, abci.Header{Height: happ.LastBlockHeight()})
-				moduleManager := happ.GetModuleManager()
-
-				for _, moduleName := range moduleManager.OrderExportGenesis {
-					module, ok := moduleManager.Modules[moduleName].(hmModule.StreamedGenesisExporter)
-					if !ok {
-						continue
-					}
-
-					runtime.GC()
-
-					if err := fetchModuleData(appState, module, sdkCtx); err != nil {
-						panic(err)
-					}
-				}
-
-				runtime.GC()
-
-				marshaledAppState, err = jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(appState)
-				if err != nil {
-					panic(err)
-				}
-			}()
+			marshaledAppState, err := generateMarshalledAppState(happ, 1000)
+			if err != nil {
+				panic(err)
+			}
 
 			runtime.GC()
 
@@ -281,6 +252,42 @@ func exportCmd(ctx *server.Context, _ *codec.Codec) *cobra.Command {
 	return cmd
 }
 
+// generateMarshalledAppState fetches the app state, populates it with module data,
+// and marshals it into JSON.
+func generateMarshalledAppState(happ *app.HeimdallApp, maxNextGenesisItems int) ([]byte, error) {
+	appState, err := getAppState(happ)
+	if err != nil {
+		return nil, err
+	}
+
+	runtime.GC()
+
+	sdkCtx := happ.NewContext(true, abci.Header{Height: happ.LastBlockHeight()})
+	moduleManager := happ.GetModuleManager()
+
+	for _, moduleName := range moduleManager.OrderExportGenesis {
+		module, ok := moduleManager.Modules[moduleName].(hmModule.StreamedGenesisExporter)
+		if !ok {
+			continue
+		}
+
+		runtime.GC()
+
+		if err := fetchModuleData(sdkCtx, appState, module, maxNextGenesisItems); err != nil {
+			return nil, err
+		}
+	}
+
+	runtime.GC()
+
+	marshaledAppState, err := jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(appState)
+	if err != nil {
+		return nil, err
+	}
+
+	return marshaledAppState, nil
+}
+
 // getAppState generates and returns the app state.
 func getAppState(happ *app.HeimdallApp) (map[string]interface{}, error) {
 	appState, _, err := happ.ExportAppStateAndValidators()
@@ -302,19 +309,19 @@ func getAppState(happ *app.HeimdallApp) (map[string]interface{}, error) {
 }
 
 // fetchModuleData fetches module genesis data in streamed fashion.
-func fetchModuleData(appData map[string]interface{}, module hmModule.StreamedGenesisExporter, sdkCtx sdk.Context) error {
+func fetchModuleData(sdkCtx sdk.Context, appData map[string]interface{}, module hmModule.StreamedGenesisExporter, maxNextGenesisItems int) error {
 	var lastKey []byte
 	allData := []json.RawMessage{}
 	allDataLength := 0
 	var currAppendingPath string
 
 	for {
-		data, err := module.NextGenesisData(sdkCtx, lastKey)
+		data, err := module.NextGenesisData(sdkCtx, lastKey, maxNextGenesisItems)
 		if err != nil {
 			panic(err)
 		}
 
-		lastKey = data.LastKey
+		lastKey = data.NextKey
 
 		if lastKey == nil {
 			allData = append(allData, data.Data)

diff --git a/types/module/module.go b/types/module/module.go
@@ -30,13 +30,13 @@ type ModuleGenesisData struct {
 	// Data is the JSON data chunk to be appended.
 	Data json.RawMessage
 
-	// LastKey is the last key used to append data.
-	LastKey []byte
+	// NextKey is the last key used to append data.
+	NextKey []byte
 }
 
 // StreamedGenesisExporter defines an interface for modules to export their genesis data incrementally.
 type StreamedGenesisExporter interface {
 	// NextGenesisData returns the next chunk of genesis data.
-	// Returns nil when no more data is available.
-	NextGenesisData(ctx sdk.Context, lastKey []byte) (*ModuleGenesisData, error)
+	// Returns nil NextKey when no more data is available.
+	NextGenesisData(ctx sdk.Context, nextKey []byte, max int) (*ModuleGenesisData, error)
 }