feat(store): Parallel write in the CacheMultiStore

store/CHANGELOG.md

@@ -34,6 +34,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Improvements
 
+* [#20817](https://github.com/cosmos/cosmos-sdk/pull/20817) Parallelize the `CacheMultiStore.Write` method.
 * [#19770](https://github.com/cosmos/cosmos-sdk/pull/19770) Upgrade IAVL to IAVL v1.1.1.
 
 ## v1.0.2 (January 10, 2024)

store/cachemulti/benchmark_test.go

@@ -0,0 +1,63 @@
+package cachemulti
+
+import (
+	"fmt"
+	"testing"
+
+	"cosmossdk.io/core/log"
+	"cosmossdk.io/store/iavl"
+	"cosmossdk.io/store/types"
+	dbm "github.com/cosmos/cosmos-db"
+	"github.com/stretchr/testify/require"
+)
+
+func setupStore(b *testing.B, storeCount uint) (Store, map[string]types.StoreKey) {
+	db := dbm.NewMemDB()
+
+	storeKeys := make(map[string]types.StoreKey)
+	stores := make(map[types.StoreKey]types.CacheWrapper)
+	for i := uint(0); i < storeCount; i++ {
+		key := types.NewKVStoreKey(fmt.Sprintf("store%d", i))
+		storeKeys[key.Name()] = key
+		sdb := dbm.NewPrefixDB(db, []byte(key.Name()))
+		istore, err := iavl.LoadStore(sdb, log.NewNopLogger(), key, types.CommitID{}, 1000, false, nil)
+		require.NoError(b, err)
+		stores[key] = types.KVStore(istore)
+	}
+
+	return NewStore(db, stores, storeKeys, nil, types.TraceContext{}), storeKeys
+}
+
+func benchmarkStore(b *testing.B, storeCount, keyCount uint) {
+	store, storeKeys := setupStore(b, storeCount)
+	b.ResetTimer()
+
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		for _, key := range storeKeys {
+			cstore := store.GetKVStore(key)
+			for j := uint(0); j < keyCount; j++ {
+				dataKey := fmt.Sprintf("key%s-%d", key.Name(), j)
+				dataValue := fmt.Sprintf("value%s-%d", key.Name(), j)
+				cstore.Set([]byte(dataKey), []byte(dataValue))
+			}
+		}
+		b.StartTimer()
+		store.Write()
+	}
+}
+
+func BenchmarkCacheMultiStore(b *testing.B) {
+	storeCounts := []uint{2, 4, 8, 16, 32}
+	keyCounts := []uint{100, 1000, 10000}
+
+	for _, storeCount := range storeCounts {
+		for _, keyCount := range keyCounts {
+			b.Run(fmt.Sprintf("storeCount=%d/keyCount=%d", storeCount, keyCount), func(sub *testing.B) {
+				benchmarkStore(sub, storeCount, keyCount)
+			})
+		}
+	}
+
+}

store/cachemulti/store.go

@@ -3,6 +3,7 @@
 import (
 	"fmt"
 	"io"
+	"sync"
 
 	dbm "github.com/cosmos/cosmos-db"
 
@@ -122,9 +123,15 @@
 // Write calls Write on each underlying store.
 func (cms Store) Write() {
 	cms.db.Write()
+	wg := sync.WaitGroup{}
+	wg.Add(len(cms.stores))
 	for _, store := range cms.stores {
-		store.Write()
+		go func(s types.CacheWrap) {
+			defer wg.Done()
+			s.Write()
+		}(store)
 	}
+	wg.Wait()
 }
 
 // Implements CacheWrapper.

store/rootmulti/store_test.go

@@ -116,7 +116,7 @@ func TestCacheMultiStoreWithVersion(t *testing.T) {
 	// require we cannot commit (write) to a cache-versioned multi-store
 	require.Panics(t, func() {
 		kvStore.Set(k, []byte("newValue"))
-		cms.Write()
+		kvStore.(types.CacheWrap).Write()
 	})
 }
 
@@ -1032,3 +1032,40 @@ func TestCommitStores(t *testing.T) {
 		})
 	}
 }
+
+func TestCacheMultiStoreWrite(t *testing.T) {
+	// for testing the cacheMultiStore parallel write, file based db is used
+	db, err := dbm.NewDB("test", dbm.GoLevelDBBackend, t.TempDir())
+	require.NoError(t, err)
+
+	ms := newMultiStoreWithMounts(db, pruningtypes.NewPruningOptions(pruningtypes.PruningNothing))
+	require.NoError(t, ms.LoadLatestVersion())
+
+	cacheMulti := ms.CacheMultiStore()
+
+	toVersion := int64(100)
+	keyCount := 100
+	storeKeys := []types.StoreKey{testStoreKey1, testStoreKey2, testStoreKey3}
+	for i := int64(1); i <= toVersion; i++ {
+		for _, storeKey := range storeKeys {
+			store := cacheMulti.GetKVStore(storeKey)
+			for j := 0; j < keyCount; j++ {
+				store.Set([]byte(fmt.Sprintf("key-%d-%d", i, j)), []byte(fmt.Sprintf("value-%d-%d", i, j)))
+			}
+		}
+		cacheMulti.Write()
+		ms.Commit()
+	}
+
+	// check the data
+	for _, storeKey := range storeKeys {
+		store := cacheMulti.GetKVStore(storeKey)
+		for i := int64(1); i <= toVersion; i++ {
+			for j := 0; j < keyCount; j++ {
+				key := []byte(fmt.Sprintf("key-%d-%d", i, j))
+				value := store.Get(key)
+				require.Equal(t, []byte(fmt.Sprintf("value-%d-%d", i, j)), value)
+			}
+		}
+	}
+}