From 653bf173c6245c7eb2e1c11ccb8d3fa63543f477 Mon Sep 17 00:00:00 2001
From: HuangYi <huang@crypto.com>
Date: Tue, 26 Nov 2024 23:10:17 +0800
Subject: [PATCH] Problem: prune cmd don't wait for async pruning to finish

---
 CHANGELOG.md             |  1 +
 client/pruning/main.go   |  8 ++++++--
 store/iavl/store.go      |  7 +++++++
 store/rootmulti/store.go | 10 ++++++++++
 4 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dc35a28c4c68..d87e63dc92c1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -55,6 +55,7 @@ Every module contains its own CHANGELOG.md. Please refer to the module you are i
 
 * (sims) [#21906](https://github.com/cosmos/cosmos-sdk/pull/21906) Skip sims test when running dry on validators
 * (cli) [#21919](https://github.com/cosmos/cosmos-sdk/pull/21919) Query address-by-acc-num by account_id instead of id.
+* (store) [#]() Prune cmd wait for async pruning to finish.
 
 ### API Breaking Changes
 
diff --git a/client/pruning/main.go b/client/pruning/main.go
index bdcff50a9bc5..28e7245fb620 100644
--- a/client/pruning/main.go
+++ b/client/pruning/main.go
@@ -89,8 +89,12 @@ Supported app-db-backend types include 'goleveldb', 'rocksdb', 'pebbledb'.`,
 			pruningHeight := latestHeight - int64(pruningOptions.KeepRecent)
 			cmd.Printf("pruning heights up to %v\n", pruningHeight)
 
-			err = rootMultiStore.PruneStores(pruningHeight)
-			if err != nil {
+			if err := rootMultiStore.PruneStores(pruningHeight); err != nil {
+				return err
+			}
+
+			// close will wait for async pruning process to finish
+			if err := rootMultiStore.Close(); err != nil {
 				return err
 			}
 
diff --git a/store/iavl/store.go b/store/iavl/store.go
index ab04b73c47b7..566556a230d2 100644
--- a/store/iavl/store.go
+++ b/store/iavl/store.go
@@ -404,6 +404,13 @@ func (st *Store) TraverseStateChanges(startVersion, endVersion int64, fn func(ve
 	return st.tree.TraverseStateChanges(startVersion, endVersion, fn)
 }
 
+func (st *Store) Close() error {
+	if closer, ok := st.tree.(io.Closer); ok {
+		return closer.Close()
+	}
+	return nil
+}
+
 // Takes a MutableTree, a key, and a flag for creating existence or absence proof and returns the
 // appropriate merkle.Proof. Since this must be called after querying for the value, this function should never error
 // Thus, it will panic on error rather than returning it
diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index 486fffebeba9..29929f360eab 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -1170,6 +1170,16 @@ func (rs *Store) flushMetadata(db corestore.KVStoreWithBatch, version int64, cIn
 	rs.logger.Debug("flushing metadata finished", "height", version)
 }
 
+func (rs *Store) Close() error {
+	errs := make([]error, 0, len(rs.stores))
+	for _, store := range rs.stores {
+		if closer, ok := store.(io.Closer); ok {
+			errs = append(errs, closer.Close())
+		}
+	}
+	return errors.Join(errs...)
+}
+
 type storeParams struct {
 	key            types.StoreKey
 	db             corestore.KVStoreWithBatch