Remove tempIDs while removing a wallet

iotaledger · Apr 17, 2024 · 3052423 · 3052423
1 parent e541c25
commit 3052423
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/pkg/evilwallet/evilwallet.go b/pkg/evilwallet/evilwallet.go
@@ -400,7 +400,10 @@ func (e *EvilWallet) useFreshIfInputWalletNotProvided(buildOptions *Options) (*W
 		}
 	}
 
-	w, err := e.wallets.freshWallet()
+	deleteTempIDFunc := func(tempID models.TempOutputID) {
+		e.outputManager.ClearTempIDs(tempID)
+	}
+	w, err := e.wallets.freshWallet(deleteTempIDFunc)
 	if err != nil {
 		return nil, ierrors.Wrap(err, "no Fresh wallet is available")
 	}

diff --git a/pkg/evilwallet/wallets.go b/pkg/evilwallet/wallets.go
@@ -3,6 +3,7 @@ package evilwallet
 import (
 	"go.uber.org/atomic"
 
+	"github.com/iotaledger/evil-tools/pkg/models"
 	"github.com/iotaledger/hive.go/ds/shrinkingmap"
 	"github.com/iotaledger/hive.go/lo"
 	"github.com/iotaledger/hive.go/runtime/syncutils"
@@ -117,7 +118,7 @@ func (w *Wallets) isFaucetWalletAvailable() bool {
 
 // freshWallet returns the first non-empty wallet from the faucetWallets queue. If current wallet is empty,
 // it is removed and the next enqueued one is returned.
-func (w *Wallets) freshWallet() (*Wallet, error) {
+func (w *Wallets) freshWallet(deleteTempIDFunc func(tempID models.TempOutputID)) (*Wallet, error) {
 	w.mu.Lock()
 	defer w.mu.Unlock()
 
@@ -127,7 +128,7 @@ func (w *Wallets) freshWallet() (*Wallet, error) {
 
 	wallet := lo.Return1(w.wallets.Get(w.faucetWallets[0]))
 	if wallet.IsEmpty() {
-		w.removeFreshWallet()
+		w.removeFreshWallet(deleteTempIDFunc)
 		if !w.isFaucetWalletAvailable() {
 			return nil, NoFreshOutputsAvailable
 		}
@@ -161,8 +162,16 @@ func (w *Wallets) reuseWallet(outputsNeeded int) *Wallet {
 }
 
 // removeWallet removes wallet, for Fresh wallet it will be the first wallet in a queue.
-func (w *Wallets) removeFreshWallet() {
+func (w *Wallets) removeFreshWallet(deleteTempIDFunc func(tempID models.TempOutputID)) {
 	removedID := w.faucetWallets[0]
+	wallet, exists := w.wallets.Get(removedID)
+	if exists {
+		for _, out := range wallet.UnspentOutputs() {
+			deleteTempIDFunc(out.TempID)
+		}
+		wallet.Clear()
+	}
+
 	w.faucetWallets = w.faucetWallets[1:]
 	w.wallets.Delete(removedID)
 }