feat: iavl async pruning (#593)

* iavl async pruning * go mod tidy --------- Co-authored-by: Adam Tucker <[email protected]> (cherry picked from commit 0a12f99)
osmosis-labs · Apr 8, 2024 · 00b5bde · 00b5bde
1 parent 245b6ff
commit 00b5bde
Show file tree

Hide file tree

Showing 15 changed files with 91 additions and 20 deletions.
diff --git a/go.mod b/go.mod
@@ -25,7 +25,7 @@ require (
 	github.com/cosmos/go-bip39 v1.0.0
 	github.com/cosmos/gogogateway v1.2.0
 	github.com/cosmos/gogoproto v1.4.10
-	github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7
+	github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac
 	github.com/cosmos/ics23/go v0.10.0
 	github.com/cosmos/ledger-cosmos-go v0.12.1
 	github.com/golang/mock v1.6.0

diff --git a/go.sum b/go.sum
@@ -361,8 +361,8 @@ github.com/cosmos/gogogateway v1.2.0/go.mod h1:iQpLkGWxYcnCdz5iAdLcRBSw3h7NXeOkZ
 github.com/cosmos/gogoproto v1.4.2/go.mod h1:cLxOsn1ljAHSV527CHOtaIP91kK6cCrZETRBrkzItWU=
 github.com/cosmos/gogoproto v1.4.10 h1:QH/yT8X+c0F4ZDacDv3z+xE3WU1P1Z3wQoLMBRJoKuI=
 github.com/cosmos/gogoproto v1.4.10/go.mod h1:3aAZzeRWpAwr+SS/LLkICX2/kDFyaYVzckBDzygIxek=
-github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7 h1:eGi/RRQ5K+C6J8aCSRYXuS5NFiSe+S46F1zMOEQmykY=
-github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7/go.mod h1:jLeUvm6bGT1YutCaL2fIar/8vGUE8cPZvh/gXEWDaDM=
+github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac h1:D1OG5ugS4r1Jq8U331gB4mrYsX7JQsasfWkFvdva4KI=
+github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac/go.mod h1:jLeUvm6bGT1YutCaL2fIar/8vGUE8cPZvh/gXEWDaDM=
 github.com/cosmos/ics23/go v0.10.0 h1:iXqLLgp2Lp+EdpIuwXTYIQU+AiHj9mOC2X9ab++bZDM=
 github.com/cosmos/ics23/go v0.10.0/go.mod h1:ZfJSmng/TBNTBkFemHHHj5YY7VAU/MBU980F4VU1NG0=
 github.com/cosmos/keyring v1.2.0 h1:8C1lBP9xhImmIabyXW4c3vFjjLiBdGCmfLUfeZlV1Yo=

diff --git a/server/mock/store.go b/server/mock/store.go
@@ -62,6 +62,14 @@ func (ms multiStore) Commit() storetypes.CommitID {
 	panic("not implemented")
 }
 
+func (ms multiStore) SetCommitting() {
+	panic("not implemented")
+}
+
+func (ms multiStore) UnsetCommitting() {
+	panic("not implemented")
+}
+
 func (ms multiStore) LastCommitID() storetypes.CommitID {
 	panic("not implemented")
 }

diff --git a/simapp/go.mod b/simapp/go.mod
@@ -57,7 +57,7 @@ require (
 	github.com/cosmos/go-bip39 v1.0.0 // indirect
 	github.com/cosmos/gogogateway v1.2.0 // indirect
 	github.com/cosmos/gogoproto v1.4.10 // indirect
-	github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7 // indirect
+	github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac // indirect
 	github.com/cosmos/ics23/go v0.10.0 // indirect
 	github.com/cosmos/ledger-cosmos-go v0.12.1 // indirect
 	github.com/cosmos/rosetta-sdk-go v0.10.0 // indirect

diff --git a/simapp/go.sum b/simapp/go.sum
@@ -358,8 +358,8 @@ github.com/cosmos/gogogateway v1.2.0/go.mod h1:iQpLkGWxYcnCdz5iAdLcRBSw3h7NXeOkZ
 github.com/cosmos/gogoproto v1.4.2/go.mod h1:cLxOsn1ljAHSV527CHOtaIP91kK6cCrZETRBrkzItWU=
 github.com/cosmos/gogoproto v1.4.10 h1:QH/yT8X+c0F4ZDacDv3z+xE3WU1P1Z3wQoLMBRJoKuI=
 github.com/cosmos/gogoproto v1.4.10/go.mod h1:3aAZzeRWpAwr+SS/LLkICX2/kDFyaYVzckBDzygIxek=
-github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7 h1:eGi/RRQ5K+C6J8aCSRYXuS5NFiSe+S46F1zMOEQmykY=
-github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7/go.mod h1:jLeUvm6bGT1YutCaL2fIar/8vGUE8cPZvh/gXEWDaDM=
+github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac h1:D1OG5ugS4r1Jq8U331gB4mrYsX7JQsasfWkFvdva4KI=
+github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac/go.mod h1:jLeUvm6bGT1YutCaL2fIar/8vGUE8cPZvh/gXEWDaDM=
 github.com/cosmos/ics23/go v0.10.0 h1:iXqLLgp2Lp+EdpIuwXTYIQU+AiHj9mOC2X9ab++bZDM=
 github.com/cosmos/ics23/go v0.10.0/go.mod h1:ZfJSmng/TBNTBkFemHHHj5YY7VAU/MBU980F4VU1NG0=
 github.com/cosmos/keyring v1.2.0 h1:8C1lBP9xhImmIabyXW4c3vFjjLiBdGCmfLUfeZlV1Yo=

diff --git a/store/iavl/store.go b/store/iavl/store.go
@@ -54,7 +54,8 @@ func LoadStore(db dbm.DB, logger log.Logger, key types.StoreKey, id types.Commit
 // provided DB. An error is returned if the version fails to load, or if called with a positive
 // version on an empty tree.
 func LoadStoreWithInitialVersion(db dbm.DB, logger log.Logger, key types.StoreKey, id types.CommitID, initialVersion uint64, cacheSize int, disableFastNode bool) (types.CommitKVStore, error) {
-	tree := iavl.NewMutableTree(wrapper.NewIAVLDB(db), cacheSize, disableFastNode, clog.NewNopLogger(), iavl.InitialVersionOption(initialVersion))
+	// Create a new IAVL tree with the async pruning enabled
+	tree := iavl.NewMutableTree(wrapper.NewIAVLDB(db), cacheSize, disableFastNode, clog.NewNopLogger(), iavl.InitialVersionOption(initialVersion), iavl.AsyncPruningOption(true))
 
 	isUpgradeable, err := tree.IsUpgradeable()
 	if err != nil {
@@ -117,6 +118,17 @@ func (st *Store) GetImmutable(version int64) (*Store, error) {
 	}, nil
 }
 
+// SetCommitting marks the store as committing, which will prevent any
+// parallel writes to the store. It is referenced in the async pruning.
+func (st *Store) SetCommitting() {
+	st.tree.SetCommitting()
+}
+
+// UnsetCommitting marks the store as not committing.
+func (st *Store) UnsetCommitting() {
+	st.tree.UnsetCommitting()
+}
+
 // Commit commits the current store state and returns a CommitID with the new
 // version and hash.
 func (st *Store) Commit() types.CommitID {

diff --git a/store/iavl/tree.go b/store/iavl/tree.go
@@ -23,6 +23,8 @@ type (
 		Get(key []byte) ([]byte, error)
 		Set(key, value []byte) (bool, error)
 		Remove(key []byte) ([]byte, bool, error)
+		SetCommitting()
+		UnsetCommitting()
 		SaveVersion() ([]byte, int64, error)
 		DeleteVersionsTo(version int64) error
 		Version() int64
@@ -55,6 +57,14 @@ func (it *immutableTree) Remove(_ []byte) ([]byte, bool, error) {
 	panic("cannot call 'Remove' on an immutable IAVL tree")
 }
 
+func (it *immutableTree) SetCommitting() {
+	panic("cannot call 'SetCommitting' on an immutable IAVL tree")
+}
+
+func (it *immutableTree) UnsetCommitting() {
+	panic("cannot call 'UnsetCommitting' on an immutable IAVL tree")
+}
+
 func (it *immutableTree) SaveVersion() ([]byte, int64, error) {
 	panic("cannot call 'SaveVersion' on an immutable IAVL tree")
 }

diff --git a/store/mem/store.go b/store/mem/store.go
@@ -49,6 +49,12 @@ func (s Store) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.Cach
 // Commit performs a no-op as entries are persistent between commitments.
 func (s *Store) Commit() (id types.CommitID) { return }
 
+// Implements CommitStore.
+func (s *Store) SetCommitting() {}
+
+// Implements CommitStore.
+func (s *Store) UnsetCommitting() {}
+
 func (s *Store) SetPruning(pruning pruningtypes.PruningOptions) {}
 
 // GetPruning is a no-op as pruning options cannot be directly set on this store.

diff --git a/store/rootmulti/dbadapter.go b/store/rootmulti/dbadapter.go
@@ -32,6 +32,10 @@ func (cdsa commitDBStoreAdapter) LastCommitID() types.CommitID {
 	}
 }
 
+func (cdsa commitDBStoreAdapter) SetCommitting() {}
+
+func (cdsa commitDBStoreAdapter) UnsetCommitting() {}
+
 func (cdsa commitDBStoreAdapter) SetPruning(_ pruningtypes.PruningOptions) {}
 
 // GetPruning is a no-op as pruning options cannot be directly set on this store.

diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
@@ -451,7 +451,9 @@ func (rs *Store) Commit() types.CommitID {
 		rs.logger.Debug("commit header and version mismatch", "header_height", rs.commitHeader.Height, "version", version)
 	}
 
+	rs.SetCommitting()
 	rs.lastCommitInfo = commitStores(version, rs.stores, rs.removalMap)
+	rs.UnsetCommitting()
 	rs.lastCommitInfo.Timestamp = rs.commitHeader.Time
 	defer rs.flushMetadata(rs.db, version, rs.lastCommitInfo)
 
@@ -477,6 +479,20 @@ func (rs *Store) Commit() types.CommitID {
 	}
 }
 
+// SetCommitting implements Committer/CommitStore.
+func (rs *Store) SetCommitting() {
+	for _, store := range rs.stores {
+		store.SetCommitting()
+	}
+}
+
+// UnsetCommitting implements Committer/CommitStore.
+func (rs *Store) UnsetCommitting() {
+	for _, store := range rs.stores {
+		store.UnsetCommitting()
+	}
+}
+
 // CacheWrap implements CacheWrapper/Store/CommitStore.
 func (rs *Store) CacheWrap() types.CacheWrap {
 	return rs.CacheMultiStore().(types.CacheWrap)

diff --git a/store/rootmulti/store_test.go b/store/rootmulti/store_test.go
@@ -511,6 +511,10 @@ func TestMultiStore_Pruning(t *testing.T) {
 				ms.Commit()
 			}
 
+			// asyn pruning, simulate the consensus process
+			time.Sleep(150 * time.Millisecond)
+			ms.Commit()
+
 			for _, v := range tc.saved {
 				_, err := ms.CacheMultiStoreWithVersion(v)
 				require.NoError(t, err, "expected no error when loading height: %d", v)
@@ -548,16 +552,6 @@ func TestMultiStore_Pruning_SameHeightsTwice(t *testing.T) {
 
 	require.Equal(t, numVersions, lastCommitInfo.Version)
 
-	for v := int64(1); v < numVersions-int64(keepRecent); v++ {
-		err := ms.LoadVersion(v)
-		require.Error(t, err, "expected error when loading pruned height: %d", v)
-	}
-
-	for v := int64(numVersions - int64(keepRecent)); v < numVersions; v++ {
-		err := ms.LoadVersion(v)
-		require.NoError(t, err, "expected no error when loading height: %d", v)
-	}
-
 	// Get latest
 	err := ms.LoadVersion(numVersions - 1)
 	require.NoError(t, err)
@@ -574,8 +568,19 @@ func TestMultiStore_Pruning_SameHeightsTwice(t *testing.T) {
 	require.Equal(t, numVersions, lastCommitInfo.Version)
 
 	// Ensure that can commit one more height with no panic
+	time.Sleep(150 * time.Millisecond)
 	lastCommitInfo = ms.Commit()
 	require.Equal(t, numVersions+1, lastCommitInfo.Version)
+
+	for v := int64(1); v < numVersions-int64(keepRecent); v++ {
+		err := ms.LoadVersion(v)
+		require.Error(t, err, "expected error when loading pruned height: %d", v)
+	}
+
+	for v := int64(numVersions - int64(keepRecent)); v < numVersions; v++ {
+		err := ms.LoadVersion(v)
+		require.NoError(t, err, "expected no error when loading height: %d", v)
+	}
 }
 
 func TestMultiStore_PruningRestart(t *testing.T) {
@@ -613,6 +618,8 @@ func TestMultiStore_PruningRestart(t *testing.T) {
 
 	// commit one more block and ensure the heights have been pruned
 	ms.Commit()
+	time.Sleep(150 * time.Millisecond)
+	ms.Commit()
 
 	actualHeightsToPrune, err = ms.pruningManager.GetFlushAndResetPruningHeights()
 	require.NoError(t, err)

diff --git a/store/transient/store.go b/store/transient/store.go
@@ -43,6 +43,12 @@ func (ts *Store) LastCommitID() (id types.CommitID) {
 	return
 }
 
+// Implements CommitStore
+func (ts *Store) SetCommitting() {}
+
+// Implements CommitStore
+func (ts *Store) UnsetCommitting() {}
+
 // Implements Store.
 func (ts *Store) GetStoreType() types.StoreType {
 	return types.StoreTypeTransient

diff --git a/store/types/store.go b/store/types/store.go
@@ -21,6 +21,8 @@ type Store interface {
 type Committer interface {
 	Commit() CommitID
 	LastCommitID() CommitID
+	SetCommitting()
+	UnsetCommitting()
 
 	SetPruning(pruningtypes.PruningOptions)
 	GetPruning() pruningtypes.PruningOptions

diff --git a/tests/go.mod b/tests/go.mod
@@ -58,7 +58,7 @@ require (
 	github.com/cosmos/cosmos-proto v1.0.0-beta.2 // indirect
 	github.com/cosmos/go-bip39 v1.0.0 // indirect
 	github.com/cosmos/gogogateway v1.2.0 // indirect
-	github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7 // indirect
+	github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac // indirect
 	github.com/cosmos/ics23/go v0.10.0 // indirect
 	github.com/cosmos/ledger-cosmos-go v0.12.2 // indirect
 	github.com/cosmos/rosetta-sdk-go v0.10.0 // indirect

diff --git a/tests/go.sum b/tests/go.sum
@@ -358,8 +358,8 @@ github.com/cosmos/gogogateway v1.2.0/go.mod h1:iQpLkGWxYcnCdz5iAdLcRBSw3h7NXeOkZ
 github.com/cosmos/gogoproto v1.4.2/go.mod h1:cLxOsn1ljAHSV527CHOtaIP91kK6cCrZETRBrkzItWU=
 github.com/cosmos/gogoproto v1.4.10 h1:QH/yT8X+c0F4ZDacDv3z+xE3WU1P1Z3wQoLMBRJoKuI=
 github.com/cosmos/gogoproto v1.4.10/go.mod h1:3aAZzeRWpAwr+SS/LLkICX2/kDFyaYVzckBDzygIxek=
-github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7 h1:eGi/RRQ5K+C6J8aCSRYXuS5NFiSe+S46F1zMOEQmykY=
-github.com/cosmos/iavl v1.1.2-0.20240405173644-e52f7630d3b7/go.mod h1:jLeUvm6bGT1YutCaL2fIar/8vGUE8cPZvh/gXEWDaDM=
+github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac h1:D1OG5ugS4r1Jq8U331gB4mrYsX7JQsasfWkFvdva4KI=
+github.com/cosmos/iavl v1.1.2-0.20240405172238-7f92c6b356ac/go.mod h1:jLeUvm6bGT1YutCaL2fIar/8vGUE8cPZvh/gXEWDaDM=
 github.com/cosmos/ics23/go v0.10.0 h1:iXqLLgp2Lp+EdpIuwXTYIQU+AiHj9mOC2X9ab++bZDM=
 github.com/cosmos/ics23/go v0.10.0/go.mod h1:ZfJSmng/TBNTBkFemHHHj5YY7VAU/MBU980F4VU1NG0=
 github.com/cosmos/keyring v1.2.0 h1:8C1lBP9xhImmIabyXW4c3vFjjLiBdGCmfLUfeZlV1Yo=