feat: per module fast nodes (#552) (#561)

* iavl v1 per module fast nodes * small logic change * initialize map * lint (cherry picked from commit 11b6186) Co-authored-by: Adam Tucker <[email protected]>
osmosis-labs · Mar 8, 2024 · c9ba975 · c9ba975
1 parent f3802a1
commit c9ba975
Show file tree

Hide file tree

Showing 8 changed files with 93 additions and 38 deletions.
diff --git a/baseapp/options.go b/baseapp/options.go
@@ -70,6 +70,11 @@ func SetIAVLDisableFastNode(disable bool) func(*BaseApp) {
 	return func(bapp *BaseApp) { bapp.cms.SetIAVLDisableFastNode(disable) }
 }
 
+// SetIAVLFastNodeModuleWhitelist returns a BaseApp option function that sets the modules to whitelist for IAVL fast node.
+func SetIAVLFastNodeModuleWhitelist(modulesToWhitelist []string) func(*BaseApp) {
+	return func(bapp *BaseApp) { bapp.cms.SetIAVLFastNodeModuleWhitelist(modulesToWhitelist) }
+}
+
 // SetInterBlockCache provides a BaseApp option function that sets the
 // inter-block cache.
 func SetInterBlockCache(cache sdk.MultiStorePersistentCache) func(*BaseApp) {

diff --git a/server/config/config.go b/server/config/config.go
@@ -95,6 +95,13 @@ type BaseConfig struct {
 	// IAVLDisableFastNode enables or disables the fast sync node.
 	IAVLDisableFastNode bool `mapstructure:"iavl-disable-fastnode"`
 
+	// IAVLFastNodeModuleWhitelist defines the whitelist of modules that will use fast nodes.
+	// If this is empty and IAVLDisableFastNode is false, all modules will use fast nodes.
+	// If this is empty and IAVLDisableFastNode is true, no modules will use fast nodes.
+	// If this is populated but IAVLDisableFastNode is true, no modules will use fast nodes.
+	// If this is populated and IAVLDisableFastNode is false, only modules in the whitelist will use fast nodes.
+	IAVLFastNodeModuleWhitelist []string `mapstructure:"iavl-fastnode-module-whitelist"`
+
 	// IAVLLazyLoading enable/disable the lazy loading of iavl store.
 	IAVLLazyLoading bool `mapstructure:"iavl-lazy-loading"`
 

diff --git a/server/config/toml.go b/server/config/toml.go
@@ -77,6 +77,16 @@ iavl-cache-size = {{ .BaseConfig.IAVLCacheSize }}
 # Default is false.
 iavl-disable-fastnode = {{ .BaseConfig.IAVLDisableFastNode }}
 
+# IAVLFastNodeModuleWhitelist defines the whitelist of modules that will use fast nodes.
+# If this is empty and IAVLDisableFastNode is false, all modules will use fast nodes.
+# If this is empty and IAVLDisableFastNode is true, no modules will use fast nodes.
+# If this is populated but IAVLDisableFastNode is true, no modules will use fast nodes.
+# If this is populated and IAVLDisableFastNode is false, only modules in the whitelist will use fast nodes.
+#
+# Example:
+# ["lockup", "superfluid"]
+iavl-fastnode-module-whitelist = [{{ range .BaseConfig.IAVLFastNodeModuleWhitelist }}{{ printf "%q, " . }}{{end}}]
+
 # IAVLLazyLoading enable/disable the lazy loading of iavl store.
 # Default is false.
 iavl-lazy-loading = {{ .BaseConfig.IAVLLazyLoading }}

diff --git a/server/mock/store.go b/server/mock/store.go
@@ -134,6 +134,10 @@ func (ms multiStore) SetIAVLDisableFastNode(disable bool) {
 	panic("not implemented")
 }
 
+func (ms multiStore) SetIAVLFastNodeModuleWhitelist(modulesToWhitelist []string) {
+	panic("not implemented")
+}
+
 func (ms multiStore) SetLazyLoading(bool) {
 	panic("not implemented")
 }

diff --git a/server/start.go b/server/start.go
@@ -67,14 +67,15 @@ const (
 	FlagTrace              = "trace"
 	FlagInvCheckPeriod     = "inv-check-period"
 
-	FlagPruning             = "pruning"
-	FlagPruningKeepRecent   = "pruning-keep-recent"
-	FlagPruningInterval     = "pruning-interval"
-	FlagIndexEvents         = "index-events"
-	FlagMinRetainBlocks     = "min-retain-blocks"
-	FlagIAVLCacheSize       = "iavl-cache-size"
-	FlagDisableIAVLFastNode = "iavl-disable-fastnode"
-	FlagIAVLLazyLoading     = "iavl-lazy-loading"
+	FlagPruning                     = "pruning"
+	FlagPruningKeepRecent           = "pruning-keep-recent"
+	FlagPruningInterval             = "pruning-interval"
+	FlagIndexEvents                 = "index-events"
+	FlagMinRetainBlocks             = "min-retain-blocks"
+	FlagIAVLCacheSize               = "iavl-cache-size"
+	FlagDisableIAVLFastNode         = "iavl-disable-fastnode"
+	FlagIAVLFastNodeModuleWhitelist = "iavl-fastnode-module-whitelist"
+	FlagIAVLLazyLoading             = "iavl-lazy-loading"
 
 	// state sync-related flags
 	FlagStateSyncSnapshotInterval   = "state-sync.snapshot-interval"

diff --git a/server/util.go b/server/util.go
@@ -495,6 +495,7 @@ func DefaultBaseappOptions(appOpts types.AppOptions) []func(*baseapp.BaseApp) {
 		baseapp.SetSnapshot(snapshotStore, snapshotOptions),
 		baseapp.SetIAVLCacheSize(cast.ToInt(appOpts.Get(FlagIAVLCacheSize))),
 		baseapp.SetIAVLDisableFastNode(cast.ToBool(appOpts.Get(FlagDisableIAVLFastNode))),
+		baseapp.SetIAVLFastNodeModuleWhitelist(cast.ToStringSlice(appOpts.Get(FlagIAVLFastNodeModuleWhitelist))),
 		baseapp.SetMempool(
 			mempool.NewSenderNonceMempool(
 				mempool.SenderNonceMaxTxOpt(cast.ToInt(appOpts.Get(FlagMempoolMaxTxs))),

diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
@@ -55,24 +55,25 @@ func keysForStoreKeyMap[V any](m map[types.StoreKey]V) []types.StoreKey {
 // cacheMultiStore which is used for branching other MultiStores. It implements
 // the CommitMultiStore interface.
 type Store struct {
-	db                  dbm.DB
-	logger              log.Logger
-	lastCommitInfo      *types.CommitInfo
-	pruningManager      *pruning.Manager
-	iavlCacheSize       int
-	iavlDisableFastNode bool
-	storesParams        map[types.StoreKey]storeParams
-	stores              map[types.StoreKey]types.CommitKVStore
-	keysByName          map[string]types.StoreKey
-	lazyLoading         bool
-	initialVersion      int64
-	removalMap          map[types.StoreKey]bool
-	traceWriter         io.Writer
-	traceContext        types.TraceContext
-	traceContextMutex   sync.Mutex
-	interBlockCache     types.MultiStorePersistentCache
-	listeners           map[types.StoreKey][]types.WriteListener
-	commitHeader        cmtproto.Header
+	db                          dbm.DB
+	logger                      log.Logger
+	lastCommitInfo              *types.CommitInfo
+	pruningManager              *pruning.Manager
+	iavlCacheSize               int
+	iavlDisableFastNode         bool
+	iavlFastNodeModuleWhitelist map[string]bool
+	storesParams                map[types.StoreKey]storeParams
+	stores                      map[types.StoreKey]types.CommitKVStore
+	keysByName                  map[string]types.StoreKey
+	lazyLoading                 bool
+	initialVersion              int64
+	removalMap                  map[types.StoreKey]bool
+	traceWriter                 io.Writer
+	traceContext                types.TraceContext
+	traceContextMutex           sync.Mutex
+	interBlockCache             types.MultiStorePersistentCache
+	listeners                   map[types.StoreKey][]types.WriteListener
+	commitHeader                cmtproto.Header
 }
 
 var (
@@ -86,16 +87,17 @@ var (
 // LoadVersion must be called.
 func NewStore(db dbm.DB, logger log.Logger) *Store {
 	return &Store{
-		db:                  db,
-		logger:              logger,
-		iavlCacheSize:       iavl.DefaultIAVLCacheSize,
-		iavlDisableFastNode: iavlDisablefastNodeDefault,
-		storesParams:        make(map[types.StoreKey]storeParams),
-		stores:              make(map[types.StoreKey]types.CommitKVStore),
-		keysByName:          make(map[string]types.StoreKey),
-		listeners:           make(map[types.StoreKey][]types.WriteListener),
-		removalMap:          make(map[types.StoreKey]bool),
-		pruningManager:      pruning.NewManager(db, logger),
+		db:                          db,
+		logger:                      logger,
+		iavlCacheSize:               iavl.DefaultIAVLCacheSize,
+		iavlDisableFastNode:         iavlDisablefastNodeDefault,
+		iavlFastNodeModuleWhitelist: make(map[string]bool),
+		storesParams:                make(map[types.StoreKey]storeParams),
+		stores:                      make(map[types.StoreKey]types.CommitKVStore),
+		keysByName:                  make(map[string]types.StoreKey),
+		listeners:                   make(map[types.StoreKey][]types.WriteListener),
+		removalMap:                  make(map[types.StoreKey]bool),
+		pruningManager:              pruning.NewManager(db, logger),
 	}
 }
 
@@ -125,6 +127,12 @@ func (rs *Store) SetIAVLDisableFastNode(disableFastNode bool) {
 	rs.iavlDisableFastNode = disableFastNode
 }
 
+func (rs *Store) SetIAVLFastNodeModuleWhitelist(modulesToWhitelist []string) {
+	for _, module := range modulesToWhitelist {
+		rs.iavlFastNodeModuleWhitelist[module] = true
+	}
+}
+
 // SetLazyLoading sets if the iavl store should be loaded lazily or not
 func (rs *Store) SetLazyLoading(lazyLoading bool) {
 	rs.lazyLoading = lazyLoading
@@ -942,6 +950,22 @@ func (rs *Store) loadCommitStoreFromParams(key types.StoreKey, id types.CommitID
 		db = dbm.NewPrefixDB(rs.db, []byte(prefix))
 	}
 
+	disabledFastNodes := rs.iavlDisableFastNode
+	// If fast nodes are enabled:
+	if !rs.iavlDisableFastNode {
+		// If the whitelist is empty, enable fast nodes for all modules.
+		if len(rs.iavlFastNodeModuleWhitelist) == 0 {
+			disabledFastNodes = false
+		} else {
+			// If the whitelist is not empty, enable fast nodes for only the modules in the whitelist.
+			disabledFastNodes = true
+			if _, ok := rs.iavlFastNodeModuleWhitelist[key.Name()]; ok {
+				rs.logger.Info("fast node enabled for module", "module", key.Name())
+				disabledFastNodes = false
+			}
+		}
+	}
+
 	switch params.typ {
 	case types.StoreTypeMulti:
 		panic("recursive MultiStores not yet supported")
@@ -951,9 +975,9 @@ func (rs *Store) loadCommitStoreFromParams(key types.StoreKey, id types.CommitID
 		var err error
 
 		if params.initialVersion == 0 {
-			store, err = iavl.LoadStore(db, rs.logger, key, id, rs.iavlCacheSize, rs.iavlDisableFastNode)
+			store, err = iavl.LoadStore(db, rs.logger, key, id, rs.iavlCacheSize, disabledFastNodes)
 		} else {
-			store, err = iavl.LoadStoreWithInitialVersion(db, rs.logger, key, id, params.initialVersion, rs.iavlCacheSize, rs.iavlDisableFastNode)
+			store, err = iavl.LoadStoreWithInitialVersion(db, rs.logger, key, id, params.initialVersion, rs.iavlCacheSize, disabledFastNodes)
 		}
 
 		if err != nil {

diff --git a/store/types/store.go b/store/types/store.go
@@ -191,6 +191,9 @@ type CommitMultiStore interface {
 	// SetIAVLDisableFastNode enables/disables fastnode feature on iavl.
 	SetIAVLDisableFastNode(disable bool)
 
+	// SetIAVLFastNodeModuleWhitelist sets the modules to whitelist for IAVL fast node.
+	SetIAVLFastNodeModuleWhitelist(modulesToWhitelist []string)
+
 	// RollbackToVersion rollback the db to specific version(height).
 	RollbackToVersion(version int64) error