Reduce marshal/unmarshal operations on the genesis

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(nil, keeper.GetRecordSequences(ctx))
+	return types.NewGenesisState(keeper.GetAllEventRecords(ctx), 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, 0, len(actualParams.EventRecords))
+	require.Equal(t, len(eventRecords), len(actualParams.EventRecords))
 }

clerk/module.go

@@ -135,21 +135,32 @@ func (am AppModule) InitGenesis(ctx sdk.Context, data json.RawMessage) []abci.Va
 	return []abci.ValidatorUpdate{}
 }
 
-// ExportGenesis returns the exported genesis state as raw bytes for the auth
+// ExportGenesis returns the exported genesis state as raw bytes for the clerk
 // module.
 func (am AppModule) ExportGenesis(ctx sdk.Context) json.RawMessage {
 	gs := ExportGenesis(ctx, am.keeper)
 	return types.ModuleCdc.MustMarshalJSON(gs)
 }
 
+// ExportPartialGenesis returns the exported genesis state as raw bytes excluding the data
+// that will be returned via NextGenesisData.
+func (am AppModule) ExportPartialGenesis(ctx sdk.Context) (json.RawMessage, error) {
+	type partialGenesisState struct {
+		RecordSequences []string `json:"record_sequences" yaml:"record_sequences"`
+	}
+	return types.ModuleCdc.MustMarshalJSON(partialGenesisState{
+		RecordSequences: am.keeper.GetRecordSequences(ctx),
+	}), nil
+}
+
 // 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",
+		Path:    "event_records",
 		Data:    types.ModuleCdc.MustMarshalJSON(data),
 		NextKey: nextKey,
 	}, nil

cmd/heimdallcli/main.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	defaultLogger "log"
 	"os"
 	"path"
@@ -33,7 +34,6 @@ import (
 	"github.com/tendermint/tendermint/crypto/secp256k1"
 	"github.com/tendermint/tendermint/libs/cli"
 	"github.com/tendermint/tendermint/libs/common"
-	tmCmn "github.com/tendermint/tendermint/libs/common"
 	"github.com/tendermint/tendermint/libs/log"
 	"github.com/tendermint/tendermint/privval"
 	tmTypes "github.com/tendermint/tendermint/types"
@@ -228,15 +228,14 @@ func exportCmd(ctx *server.Context, _ *codec.Codec) *cobra.Command {
 
 			happ := app.NewHeimdallApp(logger, db)
 
-			marshaledAppState, err := generateMarshalledAppState(happ, chainID, 1000)
+			savePath := file.Rootify("dump-genesis.json", config.RootDir)
+			file, err := os.Create(savePath)
 			if err != nil {
 				panic(err)
 			}
+			defer file.Close()
 
-			runtime.GC()
-
-			savePath := file.Rootify("dump-genesis.json", config.RootDir)
-			if err := tmCmn.WriteFile(savePath, marshaledAppState, 0644); err != nil {
+			if err := generateMarshalledAppState(happ, chainID, 1000, file); err != nil {
 				panic(err)
 			}
 
@@ -252,77 +251,128 @@ 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, chainID string, maxNextGenesisItems int) ([]byte, error) {
-	appState, err := getAppState(happ)
-	if err != nil {
-		return nil, err
+// generateMarshalledAppState writes the genesis doc with app state directly to a file to minimize memory usage.
+func generateMarshalledAppState(happ *app.HeimdallApp, chainID string, maxNextGenesisItems int, w io.Writer) error {
+	sdkCtx := happ.NewContext(true, abci.Header{Height: happ.LastBlockHeight()})
+	moduleManager := happ.GetModuleManager()
+
+	if _, err := w.Write([]byte("{")); err != nil {
+		return err
 	}
 
-	runtime.GC()
+	if _, err := w.Write([]byte(`"app_state":`)); err != nil {
+		return err
+	}
 
-	sdkCtx := happ.NewContext(true, abci.Header{Height: happ.LastBlockHeight()})
-	moduleManager := happ.GetModuleManager()
+	if _, err := w.Write([]byte(`{`)); err != nil {
+		return err
+	}
+
+	isFirst := true
 
 	for _, moduleName := range moduleManager.OrderExportGenesis {
-		module, ok := moduleManager.Modules[moduleName].(hmModule.StreamedGenesisExporter)
-		if !ok {
+		runtime.GC()
+
+		if !isFirst {
+			if _, err := w.Write([]byte(`,`)); err != nil {
+				return err
+			}
+		}
+
+		isFirst = false
+
+		if _, err := w.Write([]byte(`"` + moduleName + `":`)); err != nil {
+			return err
+		}
+
+		module, isStreamedGenesis := moduleManager.Modules[moduleName].(hmModule.StreamedGenesisExporter)
+		if isStreamedGenesis {
+			partialGenesis, err := module.ExportPartialGenesis(sdkCtx)
+			if err != nil {
+				return err
+			}
+
+			propertyName, data, err := fetchModuleStreamedData(sdkCtx, module, maxNextGenesisItems)
+			if err != nil {
+				return err
+			}
+
+			// remove the closing '}'
+			if _, err = w.Write(partialGenesis[0 : len(partialGenesis)-1]); err != nil {
+				return err
+			}
+
+			if _, err = w.Write([]byte(`,`)); err != nil {
+				return err
+			}
+
+			if _, err = w.Write([]byte(`"` + propertyName + `":`)); err != nil {
+				return err
+			}
+
+			if _, err = w.Write(data); err != nil {
+				return err
+			}
+
+			// add the closing '}'
+			if _, err = w.Write(partialGenesis[len(partialGenesis)-1:]); err != nil {
+				return err
+			}
+
 			continue
 		}
 
-		runtime.GC()
+		genesis := moduleManager.Modules[moduleName].ExportGenesis(sdkCtx)
 
-		if err := fetchModuleData(sdkCtx, appState, module, maxNextGenesisItems); err != nil {
-			return nil, err
+		if _, err := w.Write(genesis); err != nil {
+			return err
 		}
 	}
 
-	runtime.GC()
-
-	genesisDoc := map[string]interface{}{
-		"app_state":        appState,
-		"chain_id":         chainID,
-		"consensus_params": tmTypes.DefaultConsensusParams(),
-		"genesis_time":     time.Now(),
+	if _, err := w.Write([]byte(`}`)); err != nil {
+		return err
 	}
 
-	buf := new(bytes.Buffer)
-	encoder := json.NewEncoder(buf)
-	encoder.SetEscapeHTML(true)
-	if err = encoder.Encode(genesisDoc); err != nil {
-		return nil, err
+	if _, err := w.Write([]byte(`,`)); err != nil {
+		return err
 	}
 
-	return buf.Bytes(), nil
-}
+	consensusParams := tmTypes.DefaultConsensusParams()
+	genesisTime := time.Now().UTC().Format(time.RFC3339Nano)
 
-// getAppState generates and returns the app state.
-func getAppState(happ *app.HeimdallApp) (map[string]interface{}, error) {
-	appState, _, err := happ.ExportAppStateAndValidators()
+	consensusParamsData, err := tmTypes.GetCodec().MarshalJSON(consensusParams)
 	if err != nil {
-		return nil, err
+		return err
+	}
+
+	remainingFields := map[string]interface{}{
+		"chain_id":         chainID,
+		"consensus_params": json.RawMessage(consensusParamsData),
+		"genesis_time":     genesisTime,
 	}
 
-	appStateData, err := appState.MarshalJSON()
+	remainingFieldsData, err := json.Marshal(remainingFields)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	unmarshaledData := make(map[string]interface{})
-	if err := json.Unmarshal(appStateData, &unmarshaledData); err != nil {
-		return nil, err
+	if _, err := w.Write(remainingFieldsData[1 : len(remainingFieldsData)-1]); err != nil {
+		return err
+	}
+
+	if _, err := w.Write([]byte("}")); err != nil {
+		return err
 	}
 
-	return unmarshaledData, nil
+	return nil
 }
 
-// fetchModuleData fetches module genesis data in streamed fashion.
-func fetchModuleData(sdkCtx sdk.Context, appData map[string]interface{}, module hmModule.StreamedGenesisExporter, maxNextGenesisItems int) error {
+// fetchModuleStreamedData fetches module genesis data in streamed fashion.
+func fetchModuleStreamedData(sdkCtx sdk.Context, module hmModule.StreamedGenesisExporter, maxNextGenesisItems int) (string, json.RawMessage, error) {
 	var lastKey []byte
 	allData := []json.RawMessage{}
 	allDataLength := 0
-	var currAppendingPath string
+	// var currAppendingPath string
 
 	for {
 		data, err := module.NextGenesisData(sdkCtx, lastKey, maxNextGenesisItems)
@@ -342,37 +392,33 @@ func fetchModuleData(sdkCtx sdk.Context, appData map[string]interface{}, module
 
 			combinedData, err := combineJSONArrays(allData, allDataLength)
 			if err != nil {
-				return err
-			}
-
-			if err := AddProperty(appData, currAppendingPath, combinedData); err != nil {
-				return err
+				return "", nil, err
 			}
 
-			break
+			return data.Path, combinedData, nil
 		}
 
-		if currAppendingPath != "" && currAppendingPath != data.Path {
-			combinedData, err := combineJSONArrays(allData, allDataLength)
-			if err != nil {
-				return err
-			}
+		// if currAppendingPath != "" && currAppendingPath != data.Path {
+		// 	combinedData, err := combineJSONArrays(allData, allDataLength)
+		// 	if err != nil {
+		// 		return err
+		// 	}
 
-			if err := AddProperty(appData, currAppendingPath, combinedData); err != nil {
-				return err
-			}
+		// 	if err := AddProperty(appData, currAppendingPath, combinedData); err != nil {
+		// 		return err
+		// 	}
 
-			allData = []json.RawMessage{}
-			allDataLength = 0
-		}
+		// 	allData = []json.RawMessage{}
+		// 	allDataLength = 0
+		// }
 
-		currAppendingPath = data.Path
+		// currAppendingPath = data.Path
 
 		allData = append(allData, data.Data)
 		allDataLength += len(data.Data)
 	}
 
-	return nil
+	return "", nil, errors.New("failed to iterate module genesis data")
 }
 
 // combineJSONArrays combines multiple JSON arrays into a single JSON array.

cmd/heimdallcli/main_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"os"
 	"strconv"
@@ -32,21 +33,43 @@ func TestModulesStreamedGenesisExport(t *testing.T) {
 	ctx := happ.NewContext(true, abci.Header{Height: 1})
 	happ.GetModuleManager().InitGenesis(ctx, genesisState)
 
-	marshaledAppState, err := generateMarshalledAppState(happ, "test-chain", 2)
+	// Create a buffer to capture the output
+	var buf bytes.Buffer
+
+	// Call the refactored generateMarshalledAppState with the buffer
+	err = generateMarshalledAppState(happ, "test-chain", 2, &buf)
+	require.NoError(t, err)
+
+	// Write buf to file
+	err = os.WriteFile("./testdata/dump-genesis-streamed.json", buf.Bytes(), 0644)
 	require.NoError(t, err)
 
-	unmarshaledAppState := map[string]interface{}{}
+	// Get the bytes from the buffer
+	marshaledAppState := buf.Bytes()
+
+	// Unmarshal the output JSON into a map
+	var unmarshaledAppState map[string]interface{}
 	err = json.Unmarshal(marshaledAppState, &unmarshaledAppState)
 	require.NoError(t, err)
 
-	clerk, err := traversePath(unmarshaledAppState, "app_state.clerk")
+	// Access the "app_state" field
+	appState, ok := unmarshaledAppState["app_state"].(map[string]interface{})
+	require.True(t, ok, "app_state should be a map")
+
+	// Traverse to "clerk" module data
+	clerk, err := traversePath(appState, "clerk")
 	require.NoError(t, err)
-	eventRecords := clerk["event_records"].([]interface{})
+
+	// Validate "event_records" in the "clerk" module
+	eventRecords, ok := clerk["event_records"].([]interface{})
+	require.True(t, ok, "event_records should be an array")
 	require.Len(t, eventRecords, 6)
 	for idx, record := range eventRecords {
-		eventRecord := record.(map[string]interface{})
+		eventRecord, ok := record.(map[string]interface{})
+		require.True(t, ok, "eventRecord should be a map")
 		require.NotEmpty(t, eventRecord["id"])
-		eventIDStr := eventRecord["id"].(string)
+		eventIDStr, ok := eventRecord["id"].(string)
+		require.True(t, ok, "id should be a string")
 		eventID, err := strconv.Atoi(eventIDStr)
 		require.NoError(t, err)
 		require.Equal(t, idx+1, eventID)
@@ -58,14 +81,20 @@ func TestModulesStreamedGenesisExport(t *testing.T) {
 		require.NotEmpty(t, eventRecord["record_time"])
 	}
 
-	bor, err := traversePath(unmarshaledAppState, "app_state.bor")
+	// Traverse to "bor" module data
+	bor, err := traversePath(appState, "bor")
 	require.NoError(t, err)
-	spans := bor["spans"].([]interface{})
+
+	// Validate "spans" in the "bor" module
+	spans, ok := bor["spans"].([]interface{})
+	require.True(t, ok, "spans should be an array")
 	require.Len(t, spans, 5)
 	for idx, span := range spans {
-		spanMap := span.(map[string]interface{})
+		spanMap, ok := span.(map[string]interface{})
+		require.True(t, ok, "span should be a map")
 		require.NotEmpty(t, spanMap["span_id"])
-		spanIDStr := spanMap["span_id"].(string)
+		spanIDStr, ok := spanMap["span_id"].(string)
+		require.True(t, ok, "span_id should be a string")
 		spanID, err := strconv.Atoi(spanIDStr)
 		require.NoError(t, err)
 		require.Equal(t, idx, spanID)