add Discard method to CacheWrap

crypto-org-chain · Apr 1, 2024 · 01ca692 · 01ca692
1 parent f69cf4b
commit 01ca692
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 3 deletions.
diff --git a/store/CHANGELOG.md b/store/CHANGELOG.md
@@ -33,6 +33,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#241](https://github.com/crypto-org-chain/cosmos-sdk/pull/241) Refactor the cache store to be btree backed, prepare to support copy-on-write atomic branching.
 * [#242](https://github.com/crypto-org-chain/cosmos-sdk/pull/242) Init cache on cache lazily, save memory allocations.
 * [#243](https://github.com/crypto-org-chain/cosmos-sdk/pull/243) Support `RunAtomic` API to use new CoW cache store.
+* [#244](https://github.com/crypto-org-chain/cosmos-sdk/pull/244) Add `Discard` method to CacheWrap to discard the write buffer.
 
 ## v1.1.0 (March 20, 2024)
 

diff --git a/store/cachekv/store.go b/store/cachekv/store.go
@@ -119,6 +119,10 @@ func (store *GStore[V]) Write() {
 	store.writeSet.Clear()
 }
 
+func (store *GStore[V]) Discard() {
+	store.writeSet.Clear()
+}
+
 // CacheWrap implements CacheWrapper.
 func (store *GStore[V]) CacheWrap() types.CacheWrap {
 	return NewGStore(store, store.isZero, store.valueLen)

diff --git a/store/cachemulti/store.go b/store/cachemulti/store.go
@@ -126,6 +126,12 @@ func (cms Store) Write() {
 	}
 }
 
+func (cms Store) Discard() {
+	for _, store := range cms.stores {
+		store.Discard()
+	}
+}
+
 // Implements CacheWrapper.
 func (cms Store) CacheWrap() types.CacheWrap {
 	return cms.CacheMultiStore().(types.CacheWrap)
@@ -202,10 +208,10 @@ func (cms Store) Restore(other Store) {
 		cms.stores[k].(types.BranchStore).Restore(v.(types.BranchStore))
 	}
 
-	for k := range cms.stores {
+	for k, v := range cms.stores {
 		if _, ok := other.stores[k]; !ok {
-			// delete the store if it's not in the other
-			delete(cms.stores, k)
+			// clear the cache store if it's not in the other
+			v.Discard()
 		}
 	}
 }

diff --git a/store/types/store.go b/store/types/store.go
@@ -339,6 +339,9 @@ type CacheWrap interface {
 
 	// Write syncs with the underlying store.
 	Write()
+
+	// Discard the write set
+	Discard()
 }
 
 type CacheWrapper interface {