From 0b039644fdc52644e963d13cb7d3700e48133e07 Mon Sep 17 00:00:00 2001
From: HuangYi <huang@crypto.com>
Date: Thu, 26 Oct 2023 10:30:07 +0800
Subject: [PATCH] Problem: versiondb multistore can't query mem store state

Solution:
- share the mem store instance with the main store
---
 CHANGELOG.md                 |  1 +
 app/app.go                   | 34 +++++++++++++++++-----------------
 app/versiondb.go             | 11 +++++++----
 app/versiondb_placeholder.go |  4 +---
 versiondb/multistore.go      | 10 +++++-----
 5 files changed, 31 insertions(+), 29 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 905bf90c07..ec3b2e78a1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@
 
 - [#1215](https://github.com/crypto-org-chain/cronos/pull/1215) Update ethermint to fix of concurrent write in fee history.
 - [#1217](https://github.com/crypto-org-chain/cronos/pull/1217) Use the default chain-id behavour in sdk.
+- [#]() Fix mem store in versiondb multistore.
 
 *October 17, 2023*
 
diff --git a/app/app.go b/app/app.go
index b764b33675..e589267c39 100644
--- a/app/app.go
+++ b/app/app.go
@@ -402,23 +402,6 @@ func New(
 
 	keys, memKeys, tkeys := StoreKeys(skipGravity)
 
-	// load state streaming if enabled
-	if _, _, err := streaming.LoadStreamingServices(bApp, appOpts, appCodec, logger, keys); err != nil {
-		fmt.Printf("failed to load state streaming: %s", err)
-		os.Exit(1)
-	}
-
-	// wire up the versiondb's `StreamingService` and `MultiStore`.
-	streamers := cast.ToStringSlice(appOpts.Get("store.streamers"))
-	var qms sdk.MultiStore
-	if slices.Contains(streamers, "versiondb") {
-		var err error
-		qms, err = setupVersionDB(homePath, bApp, keys, tkeys, memKeys)
-		if err != nil {
-			panic(err)
-		}
-	}
-
 	app := &App{
 		BaseApp:           bApp,
 		cdc:               cdc,
@@ -868,6 +851,23 @@ func New(
 	app.MountTransientStores(tkeys)
 	app.MountMemoryStores(memKeys)
 
+	// load state streaming if enabled
+	if _, _, err := streaming.LoadStreamingServices(bApp, appOpts, appCodec, logger, keys); err != nil {
+		fmt.Printf("failed to load state streaming: %s", err)
+		os.Exit(1)
+	}
+
+	// wire up the versiondb's `StreamingService` and `MultiStore`.
+	streamers := cast.ToStringSlice(appOpts.Get("store.streamers"))
+	var qms sdk.MultiStore
+	if slices.Contains(streamers, "versiondb") {
+		var err error
+		qms, err = app.setupVersionDB(homePath, keys, tkeys, memKeys)
+		if err != nil {
+			panic(err)
+		}
+	}
+
 	// initialize BaseApp
 	app.SetInitChainer(app.InitChainer)
 	app.SetPreBlocker(app.PreBlocker)
diff --git a/app/versiondb.go b/app/versiondb.go
index 752d17942c..460ab69187 100644
--- a/app/versiondb.go
+++ b/app/versiondb.go
@@ -7,16 +7,14 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/cosmos/cosmos-sdk/baseapp"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/crypto-org-chain/cronos/versiondb"
 	"github.com/crypto-org-chain/cronos/versiondb/tsrocksdb"
 )
 
-func setupVersionDB(
+func (app *App) setupVersionDB(
 	homePath string,
-	app *baseapp.BaseApp,
 	keys map[string]*storetypes.KVStoreKey,
 	tkeys map[string]*storetypes.TransientStoreKey,
 	memKeys map[string]*storetypes.MemoryStoreKey,
@@ -41,7 +39,12 @@ func setupVersionDB(
 
 	verDB := versiondb.NewMultiStore(versionDB, exposeStoreKeys)
 	verDB.MountTransientStores(tkeys)
-	verDB.MountMemoryStores(memKeys)
+
+	memStores := make(map[storetypes.StoreKey]storetypes.KVStore)
+	for _, memKey := range memKeys {
+		memStores[memKey] = app.CommitMultiStore().GetKVStore(memKey)
+	}
+	verDB.MountMemoryStores(memStores)
 
 	app.SetQueryMultiStore(verDB)
 	return verDB, nil
diff --git a/app/versiondb_placeholder.go b/app/versiondb_placeholder.go
index 1c46fec500..1225c8e864 100644
--- a/app/versiondb_placeholder.go
+++ b/app/versiondb_placeholder.go
@@ -6,14 +6,12 @@ package app
 import (
 	"errors"
 
-	"github.com/cosmos/cosmos-sdk/baseapp"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
-func setupVersionDB(
+func (app *App) setupVersionDB(
 	homePath string,
-	app *baseapp.BaseApp,
 	keys map[string]*storetypes.KVStoreKey,
 	tkeys map[string]*storetypes.TransientStoreKey,
 	memKeys map[string]*storetypes.MemoryStoreKey,
diff --git a/versiondb/multistore.go b/versiondb/multistore.go
index 423aec8496..d3c2b8dba9 100644
--- a/versiondb/multistore.go
+++ b/versiondb/multistore.go
@@ -5,7 +5,6 @@ import (
 	"sync"
 
 	"github.com/cosmos/cosmos-sdk/store/cachemulti"
-	"github.com/cosmos/cosmos-sdk/store/mem"
 	"github.com/cosmos/cosmos-sdk/store/transient"
 	"github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -89,10 +88,11 @@ func (s *MultiStore) MountTransientStores(keys map[string]*types.TransientStoreK
 	}
 }
 
-// MountMemoryStores simlates the same behavior as sdk to support grpc query service.
-func (s *MultiStore) MountMemoryStores(keys map[string]*types.MemoryStoreKey) {
-	for _, key := range keys {
-		s.transientStores[key] = mem.NewStore()
+// MountMemoryStores simlates the same behavior as sdk to support grpc query service,
+// it shares the existing mem store instance.
+func (s *MultiStore) MountMemoryStores(keys map[types.StoreKey]types.KVStore) {
+	for key, store := range keys {
+		s.transientStores[key] = store
 	}
 }