Archive Node Migration: Expose Get State Store + Query Iavl Before Migration Height (#541)

## Describe your changes and provide context
- Exposes `GetStateStore` from rootmulti
- For historical queries, if before migration height, points to iavl

## Testing performed to validate your change
- Verified and tested on node

---------

Co-authored-by: Yiming Zang <[email protected]>
Co-authored-by: yzang2019 <[email protected]>

Author: Kbhat1

baseapp/abci.go

@@ -13,20 +13,20 @@ import (
 	"time"
 
 	"github.com/armon/go-metrics"
-	"github.com/gogo/protobuf/proto"
-	abci "github.com/tendermint/tendermint/abci/types"
-	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
-	"google.golang.org/grpc/codes"
-	grpcstatus "google.golang.org/grpc/status"
-
 	"github.com/cosmos/cosmos-sdk/codec"
 	snapshottypes "github.com/cosmos/cosmos-sdk/snapshots/types"
+	"github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/tasks"
 	"github.com/cosmos/cosmos-sdk/telemetry"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/types/legacytm"
 	"github.com/cosmos/cosmos-sdk/utils"
+	"github.com/gogo/protobuf/proto"
+	abci "github.com/tendermint/tendermint/abci/types"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+	"google.golang.org/grpc/codes"
+	grpcstatus "google.golang.org/grpc/status"
 )
 
 // InitChain implements the ABCI interface. It runs the initialization logic
@@ -705,7 +705,8 @@ func checkNegativeHeight(height int64) error {
 // CreateQueryContext creates a new sdk.Context for a query, taking as args
 // the block height and whether the query needs a proof or not.
 func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, error) {
-	if err := checkNegativeHeight(height); err != nil {
+	err := checkNegativeHeight(height)
+	if err != nil {
 		return sdk.Context{}, err
 	}
 
@@ -731,7 +732,15 @@ func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, e
 			)
 	}
 
-	cacheMS, err := app.cms.CacheMultiStoreWithVersion(height)
+	var cacheMS types.CacheMultiStore
+	if height < app.migrationHeight && app.qms != nil {
+		cacheMS, err = app.qms.CacheMultiStoreWithVersion(height)
+		app.logger.Info("SeiDB Archive Migration: Serving Query From Iavl", "height", height)
+	} else {
+		cacheMS, err = app.cms.CacheMultiStoreWithVersion(height)
+		app.logger.Info("SeiDB Archive Migration: Serving Query From State Store", "height", height)
+	}
+
 	if err != nil {
 		return sdk.Context{},
 			sdkerrors.Wrapf(
@@ -902,12 +911,26 @@ func handleQueryApp(app *BaseApp, path []string, req abci.RequestQuery) abci.Res
 }
 
 func handleQueryStore(app *BaseApp, path []string, req abci.RequestQuery) abci.ResponseQuery {
-	// "/store" prefix for store queries
-	queryable, ok := app.cms.(sdk.Queryable)
-	if !ok {
-		return sdkerrors.QueryResultWithDebug(sdkerrors.Wrap(sdkerrors.ErrUnknownRequest, "multistore doesn't support queries"), app.trace)
+	var (
+		queryable sdk.Queryable
+		ok        bool
+	)
+	// Check if online migration is enabled for fallback read
+	if req.Height < app.migrationHeight && app.qms != nil {
+		queryable, ok = app.qms.(sdk.Queryable)
+		if !ok {
+			return sdkerrors.QueryResultWithDebug(sdkerrors.Wrap(sdkerrors.ErrUnknownRequest, "multistore doesn't support queries"), app.trace)
+		}
+		app.logger.Info("SeiDB Archive Migration: Serving Query From Iavl", "height", req.Height)
+	} else {
+		queryable, ok = app.cms.(sdk.Queryable)
+		if !ok {
+			return sdkerrors.QueryResultWithDebug(sdkerrors.Wrap(sdkerrors.ErrUnknownRequest, "multistore doesn't support queries"), app.trace)
+		}
+		app.logger.Info("SeiDB Archive Migration: Serving Query From State Store", "height", req.Height)
 	}
 
+	// "/store" prefix for store queries
 	req.Path = "/" + strings.Join(path[1:], "/")
 
 	if req.Height <= 1 && req.Prove {

baseapp/abci_test.go

@@ -11,6 +11,7 @@ import (
 	dbm "github.com/tendermint/tm-db"
 
 	"github.com/cosmos/cosmos-sdk/testutil"
+	"github.com/cosmos/cosmos-sdk/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -200,3 +201,27 @@ func (ps *paramStore) Get(_ sdk.Context, key []byte, ptr interface{}) {
 		panic(err)
 	}
 }
+func TestHandleQueryStore_NonQueryableMultistore(t *testing.T) {
+	logger := defaultLogger()
+	db := dbm.NewMemDB()
+	name := t.Name()
+	app := NewBaseApp(name, logger, db, nil, nil, &testutil.TestAppOpts{})
+
+	// Mock a non-queryable cms
+	mockCMS := &mockNonQueryableMultiStore{}
+	app.cms = mockCMS
+
+	path := []string{"store", "test"}
+	req := abci.RequestQuery{
+		Path:   "store/test",
+		Height: 1,
+	}
+
+	resp := handleQueryStore(app, path, req)
+	require.True(t, resp.IsErr())
+	require.Contains(t, resp.Log, "multistore doesn't support queries")
+}
+
+type mockNonQueryableMultiStore struct {
+	types.CommitMultiStore
+}

baseapp/baseapp.go

@@ -179,9 +179,11 @@ type BaseApp struct { //nolint: maligned
 }
 
 type appStore struct {
-	db          dbm.DB               // common DB backend
-	cms         sdk.CommitMultiStore // Main (uncached) state
-	storeLoader StoreLoader          // function to handle store loading, may be overridden with SetStoreLoader()
+	db              dbm.DB               // common DB backend
+	cms             sdk.CommitMultiStore // Main (uncached) state
+	qms             sdk.CommitMultiStore // Query multistore used for migration only
+	migrationHeight int64
+	storeLoader     StoreLoader // function to handle store loading, may be overridden with SetStoreLoader()
 
 	// an inter-block write-through cache provided to the context during deliverState
 	interBlockCache sdk.MultiStorePersistentCache
@@ -411,6 +413,9 @@ func (app *BaseApp) MountMemoryStores(keys map[string]*sdk.MemoryStoreKey) {
 // using the default DB.
 func (app *BaseApp) MountStore(key sdk.StoreKey, typ sdk.StoreType) {
 	app.cms.MountStoreWithDB(key, typ, nil)
+	if app.qms != nil {
+		app.qms.MountStoreWithDB(key, typ, nil)
+	}
 }
 
 // LoadLatestVersion loads the latest application version. It will panic if
@@ -421,6 +426,13 @@ func (app *BaseApp) LoadLatestVersion() error {
 		return fmt.Errorf("failed to load latest version: %w", err)
 	}
 
+	if app.qms != nil {
+		err = app.storeLoader(app.qms)
+		if err != nil {
+			return fmt.Errorf("failed to load latest version: %w", err)
+		}
+	}
+
 	return app.init()
 }
 

baseapp/options.go

@@ -361,3 +361,13 @@ func (app *BaseApp) SetStreamingService(s StreamingService) {
 	// BaseApp will pass BeginBlock, DeliverTx, and EndBlock requests and responses to the streaming services to update their ABCI context
 	app.abciListeners = append(app.abciListeners, s)
 }
+
+// SetQueryMultiStore set a alternative MultiStore implementation to support online migration fallback read.
+func (app *BaseApp) SetQueryMultiStore(ms sdk.CommitMultiStore) {
+	app.qms = ms
+}
+
+// SetMigrationHeight set the migration height for online migration so that query below this height will still be served from IAVL.
+func (app *BaseApp) SetMigrationHeight(height int64) {
+	app.migrationHeight = height
+}

storev2/rootmulti/store.go

@@ -63,6 +63,7 @@ func NewStore(
 	logger log.Logger,
 	scConfig config.StateCommitConfig,
 	ssConfig config.StateStoreConfig,
+	migrateIavl bool,
 ) *Store {
 	scStore := sc.NewCommitStore(homeDir, logger, scConfig)
 	store := &Store{
@@ -81,7 +82,7 @@ func NewStore(
 		// Check whether SC was enabled before but SS was not
 		ssVersion, _ := ssStore.GetLatestVersion()
 		scVersion, _ := scStore.GetLatestVersion()
-		if ssVersion <= 0 && scVersion > 0 {
+		if ssVersion <= 0 && scVersion > 0 && !migrateIavl {
 			panic("Enabling SS store without state sync could cause data corruption")
 		}
 		if err = ss.RecoverStateStore(logger, homeDir, ssStore); err != nil {
@@ -214,6 +215,11 @@ func (rs *Store) GetStoreType() types.StoreType {
 	return types.StoreTypeMulti
 }
 
+// GetStateStore returns the ssStore instance
+func (rs *Store) GetStateStore() sstypes.StateStore {
+	return rs.ssStore
+}
+
 // Implements interface CacheWrapper
 func (rs *Store) CacheWrap(_ types.StoreKey) types.CacheWrap {
 	return rs.CacheMultiStore().(types.CacheWrap)
@@ -757,7 +763,7 @@ loop:
 			scImporter.AddNode(node)
 
 			// Check if we should also import to SS store
-			if rs.ssStore != nil && node.Height == 0 && ssImporter != nil {
+			if rs.ssStore != nil && ssImporter != nil {
 				ssImporter <- sstypes.SnapshotNode{
 					StoreKey: storeKey,
 					Key:      node.Key,

storev2/rootmulti/store_test.go

@@ -10,6 +10,6 @@ import (
 )
 
 func TestLastCommitID(t *testing.T) {
-	store := NewStore(t.TempDir(), log.NewNopLogger(), config.StateCommitConfig{}, config.StateStoreConfig{})
+	store := NewStore(t.TempDir(), log.NewNopLogger(), config.StateCommitConfig{}, config.StateStoreConfig{}, false)
 	require.Equal(t, types.CommitID{}, store.LastCommitID())
 }