From d47bb04d3fa0a9ee311328f25cfdd69838d0cede Mon Sep 17 00:00:00 2001
From: Hans Moog <3293976+hmoog@users.noreply.github.com>
Date: Mon, 11 Apr 2022 21:39:41 +0200
Subject: [PATCH] Feat: added Shutdown calls

---
 packages/ledger/branchdag/storage.go |  2 +-
 packages/ledger/storage.go           | 35 ++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/packages/ledger/branchdag/storage.go b/packages/ledger/branchdag/storage.go
index ddbdcdd6da..eca75a7871 100644
--- a/packages/ledger/branchdag/storage.go
+++ b/packages/ledger/branchdag/storage.go
@@ -24,7 +24,7 @@ type Storage struct {
 	// conflictMemberStorage is an object storage used to persist ConflictMember objects.
 	conflictMemberStorage *objectstorage.ObjectStorage[*ConflictMember]
 
-	// shutdownOnce is used to ensure that the shutdown routine is executed only a single time.
+	// shutdownOnce is used to ensure that the Shutdown routine is executed only a single time.
 	shutdownOnce sync.Once
 }
 
diff --git a/packages/ledger/storage.go b/packages/ledger/storage.go
index b55f386f17..ad8763512a 100644
--- a/packages/ledger/storage.go
+++ b/packages/ledger/storage.go
@@ -1,8 +1,11 @@
 package ledger
 
 import (
+	"sync"
+
 	"github.com/cockroachdb/errors"
 	"github.com/iotaledger/hive.go/byteutils"
+	"github.com/iotaledger/hive.go/cerrors"
 	"github.com/iotaledger/hive.go/generics/dataflow"
 	"github.com/iotaledger/hive.go/generics/objectstorage"
 	"github.com/iotaledger/hive.go/marshalutil"
@@ -33,6 +36,9 @@ type Storage struct {
 
 	// ledger contains a reference to the Ledger that created the storage.
 	ledger *Ledger
+
+	// shutdownOnce is used to ensure that the Shutdown routine is executed only a single time.
+	shutdownOnce sync.Once
 }
 
 // newStorage returns a new storage instance for the given Ledger.
@@ -163,6 +169,35 @@ func (s *Storage) CachedConsumers(outputID utxo.OutputID) (cachedConsumers objec
 	return
 }
 
+// Prune resets the database and deletes all entities.
+func (s *Storage) Prune() (err error) {
+	for _, storagePrune := range []func() error{
+		s.transactionStorage.Prune,
+		s.transactionMetadataStorage.Prune,
+		s.outputStorage.Prune,
+		s.outputMetadataStorage.Prune,
+		s.consumerStorage.Prune,
+	} {
+		if err = storagePrune(); err != nil {
+			err = errors.Errorf("failed to prune the object storage (%v): %w", err, cerrors.ErrFatal)
+			return
+		}
+	}
+
+	return
+}
+
+// Shutdown shuts down the KVStores that are used to persist data.
+func (s *Storage) Shutdown() {
+	s.shutdownOnce.Do(func() {
+		s.transactionStorage.Shutdown()
+		s.transactionMetadataStorage.Shutdown()
+		s.outputStorage.Shutdown()
+		s.outputMetadataStorage.Shutdown()
+		s.consumerStorage.Shutdown()
+	})
+}
+
 // storeTransactionCommand is a ChainedCommand that stores a Transaction.
 func (s *Storage) storeTransactionCommand(params *dataFlowParams, next dataflow.Next[*dataFlowParams]) (err error) {
 	created := false