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,65 @@
+package cachemulti
+
+import (
+	"fmt"
+	"testing"
+
+	dbm "github.com/cosmos/cosmos-db"
+	"github.com/stretchr/testify/require"
+
+	"cosmossdk.io/log"
+	"cosmossdk.io/store/iavl"
+	"cosmossdk.io/store/types"
+)
+
+func setupStore(b *testing.B, storeCount uint) (Store, map[string]types.StoreKey) {
+	b.Helper()
+
+	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) {
+	b.Helper()
+	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

@@ -5,6 +5,7 @@ import (
 	"io"
 
 	dbm "github.com/cosmos/cosmos-db"
+	"golang.org/x/sync/errgroup"
 
 	"cosmossdk.io/store/cachekv"
 	"cosmossdk.io/store/dbadapter"
@@ -122,8 +123,21 @@ func (cms Store) GetStoreType() types.StoreType {
 // Write calls Write on each underlying store.
 func (cms Store) Write() {
 	cms.db.Write()
+	eg := new(errgroup.Group)
 	for _, store := range cms.stores {
-		store.Write()
+		s := store // https://golang.org/doc/faq#closures_and_goroutines
+		eg.Go(func() (err error) {
+			defer func() {
+				if r := recover(); r != nil {
+					err = fmt.Errorf("panic in Write: %v", r)
+				}
+			}()
+			s.Write()
+			return nil
+		})
+	}
+	if err := eg.Wait(); err != nil {
+		panic(err)
 	}
 }
 

store/go.mod

@@ -23,6 +23,7 @@ require (
 	github.com/stretchr/testify v1.9.0
 	github.com/tidwall/btree v1.7.0
 	golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842
+	golang.org/x/sync v0.7.0
 	google.golang.org/grpc v1.64.1
 	google.golang.org/protobuf v1.34.2
 	gotest.tools/v3 v3.5.1

store/rootmulti/store_test.go

@@ -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)
+			}
+		}
+	}
+}