From c77159d7f717109f1ffde8156289e4e8f43be0f6 Mon Sep 17 00:00:00 2001
From: Yawning Angel <yawning@schwanenlied.me>
Date: Thu, 9 Apr 2020 11:54:32 +0000
Subject: [PATCH] go/oasis-node/cmd/debug/consim: Add `--consim.memdb`

While it's possible to use the pruner to try to prevent the ABCI induced
disk usage from growing out of control, having to save/prune is rather
slow, even when backed by tmpfs.

Enabling this option will keep ABCI state in memory, dramatically
speeding up cases where only the final dump is required, especially if
a sensible configuration is used for the pruner.
---
 go/consensus/tendermint/abci/mux.go         | 3 +++
 go/consensus/tendermint/abci/state.go       | 9 ++++++++-
 go/oasis-node/cmd/debug/consim/consim.go    | 3 +++
 go/oasis-node/cmd/debug/consim/mockchain.go | 2 ++
 4 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/go/consensus/tendermint/abci/mux.go b/go/consensus/tendermint/abci/mux.go
index 71a2c77d9a5..d249c7b6a48 100644
--- a/go/consensus/tendermint/abci/mux.go
+++ b/go/consensus/tendermint/abci/mux.go
@@ -67,6 +67,9 @@ type ApplicationConfig struct {
 
 	// OwnTxSigner is the transaction signer identity of the local node.
 	OwnTxSigner signature.PublicKey
+
+	// TestingMemDB forces the MemDB to be used for the state storage.
+	TestingMemDB bool
 }
 
 // TransactionAuthHandler is the interface for ABCI applications that handle
diff --git a/go/consensus/tendermint/abci/state.go b/go/consensus/tendermint/abci/state.go
index 6738714201c..858f8d23a92 100644
--- a/go/consensus/tendermint/abci/state.go
+++ b/go/consensus/tendermint/abci/state.go
@@ -429,8 +429,15 @@ func (s *applicationState) metricsWorker() {
 	}
 }
 
+func newStateDB(cfg *ApplicationConfig) (dbm.DB, error) {
+	if cfg.TestingMemDB {
+		return dbm.NewMemDB(), nil
+	}
+	return db.New(filepath.Join(cfg.DataDir, "abci-mux-state"), false)
+}
+
 func newApplicationState(ctx context.Context, cfg *ApplicationConfig) (*applicationState, error) {
-	db, err := db.New(filepath.Join(cfg.DataDir, "abci-mux-state"), false)
+	db, err := newStateDB(cfg)
 	if err != nil {
 		return nil, err
 	}
diff --git a/go/oasis-node/cmd/debug/consim/consim.go b/go/oasis-node/cmd/debug/consim/consim.go
index eee7f476520..f7c095ab51b 100644
--- a/go/oasis-node/cmd/debug/consim/consim.go
+++ b/go/oasis-node/cmd/debug/consim/consim.go
@@ -29,6 +29,7 @@ import (
 
 const (
 	cfgNumKept      = "consim.num_kept"
+	cfgMemDB        = "consim.memdb"
 	cfgWorkload     = "consim.workload"
 	cfgWorkloadSeed = "consim.workload.seed"
 )
@@ -115,6 +116,7 @@ func doRun(cmd *cobra.Command, args []string) error {
 		tmChainID:     tmChainID,
 		txAuthHandler: txAuthApp.(abci.TransactionAuthHandler),
 		numVersions:   viper.GetInt64(cfgNumKept),
+		memDB:         viper.GetBool(cfgMemDB),
 	}
 	mockChain, err := initMockChain(ctx, cfg)
 	if err != nil {
@@ -248,6 +250,7 @@ func Register(parentCmd *cobra.Command) {
 
 func init() {
 	flagsConsim.Int64(cfgNumKept, 0, "number of versions kept (0 = all)")
+	flagsConsim.Bool(cfgMemDB, false, "use memory to store state")
 	flagsConsim.String(cfgWorkload, fileWorkloadName, "workload to execute")
 	flagsConsim.String(cfgWorkloadSeed, "seeeeeeeeeeeeeeeeeeeeeeeeeeeeeed", "DRBG seed for workloads")
 	_ = viper.BindPFlags(flagsConsim)
diff --git a/go/oasis-node/cmd/debug/consim/mockchain.go b/go/oasis-node/cmd/debug/consim/mockchain.go
index cabfe3b474a..8fa86cbfbd0 100644
--- a/go/oasis-node/cmd/debug/consim/mockchain.go
+++ b/go/oasis-node/cmd/debug/consim/mockchain.go
@@ -26,6 +26,7 @@ type mockChainCfg struct {
 	tmChainID     string
 	txAuthHandler abci.TransactionAuthHandler
 	numVersions   int64
+	memDB         bool
 }
 
 type mockChain struct {
@@ -158,6 +159,7 @@ func initMockChain(ctx context.Context, cfg *mockChainCfg) (*mockChain, error) {
 		HaltEpochHeight: math.MaxUint64,
 		MinGasPrice:     0, // XXX: Should this be configurable?
 		OwnTxSigner:     localSigner.Public(),
+		TestingMemDB:    cfg.memDB,
 	}
 	if cfg.numVersions > 0 {
 		muxCfg.Pruning.Strategy = abci.PruneKeepN