Re-init utxoledger state tree after we commit to release memory

pkg/protocol/engine/ledger/ledger/ledger.go

@@ -178,7 +178,8 @@ func (l *Ledger) CommitSlot(slot iotago.SlotIndex) (stateRoot iotago.Identifier,
 
 	// Commit the changes
 	// Update the UTXO ledger
-	if err = l.utxoLedger.ApplyDiff(slot, outputs, spenders); err != nil {
+	stateTreeRoot, err := l.utxoLedger.ApplyDiff(slot, outputs, spenders)
+	if err != nil {
 		return iotago.Identifier{}, iotago.Identifier{}, iotago.Identifier{}, nil, nil, nil, ierrors.Wrapf(err, "failed to apply diff to UTXO ledger for slot %d", slot)
 	}
 
@@ -217,7 +218,7 @@ func (l *Ledger) CommitSlot(slot iotago.SlotIndex) (stateRoot iotago.Identifier,
 	// (due to removing the attachments) and then committed, which would result in a broken state of the transaction.
 	l.memPool.Evict(slot)
 
-	return l.utxoLedger.StateTreeRoot(), stateDiff.Mutations().Root(), l.accountsLedger.AccountsTreeRoot(), outputs, spenders, mutations, nil
+	return stateTreeRoot, stateDiff.Mutations().Root(), l.accountsLedger.AccountsTreeRoot(), outputs, spenders, mutations, nil
 }
 
 func (l *Ledger) AddAccount(output *utxoledger.Output, blockIssuanceCredits iotago.BlockIssuanceCredits) error {

pkg/protocol/engine/utxoledger/iteration_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/iotaledger/hive.go/kvstore/mapdb"
+	"github.com/iotaledger/hive.go/lo"
 	"github.com/iotaledger/iota-core/pkg/protocol/engine/utxoledger"
 	"github.com/iotaledger/iota-core/pkg/protocol/engine/utxoledger/tpkg"
 	iotago "github.com/iotaledger/iota.go/v4"
@@ -34,7 +35,7 @@ func TestUTXOComputeBalance(t *testing.T) {
 		tpkg.RandLedgerStateSpentWithOutput(initialOutput, index),
 	}
 
-	require.NoError(t, manager.ApplyDiffWithoutLocking(index, outputs, spents))
+	require.NoError(t, lo.Return2(manager.ApplyDiffWithoutLocking(index, outputs, spents)))
 
 	spent, err := manager.SpentOutputs()
 	require.NoError(t, err)
@@ -80,7 +81,7 @@ func TestUTXOIteration(t *testing.T) {
 		tpkg.RandLedgerStateSpentWithOutput(outputs[9], index),
 	}
 
-	require.NoError(t, manager.ApplyDiffWithoutLocking(index, outputs, spents))
+	require.NoError(t, lo.Return2(manager.ApplyDiffWithoutLocking(index, outputs, spents)))
 
 	// Prepare values to check
 	outputByID := make(map[string]struct{})

pkg/protocol/engine/utxoledger/manager.go

@@ -19,24 +19,22 @@ type Manager struct {
 	store     kvstore.KVStore
 	storeLock syncutils.RWMutex
 
-	stateTree ads.Map[iotago.Identifier, iotago.OutputID, *stateTreeMetadata]
+	stateTreeKVStore kvstore.KVStore
+	stateTree        ads.Map[iotago.Identifier, iotago.OutputID, *stateTreeMetadata]
 
 	apiProvider iotago.APIProvider
 }
 
 func New(store kvstore.KVStore, apiProvider iotago.APIProvider) *Manager {
-	return &Manager{
-		store: store,
-		stateTree: ads.NewMap[iotago.Identifier](lo.PanicOnErr(store.WithExtendedRealm(kvstore.Realm{StoreKeyPrefixStateTree})),
-			iotago.Identifier.Bytes,
-			iotago.IdentifierFromBytes,
-			iotago.OutputID.Bytes,
-			iotago.OutputIDFromBytes,
-			(*stateTreeMetadata).Bytes,
-			stateMetadataFromBytes,
-		),
-		apiProvider: apiProvider,
+	m := &Manager{
+		store:            store,
+		stateTreeKVStore: lo.PanicOnErr(store.WithExtendedRealm(kvstore.Realm{StoreKeyPrefixStateTree})),
+		apiProvider:      apiProvider,
 	}
+
+	m.reInitStateTreeWithoutLocking()
+
+	return m
 }
 
 // KVStore returns the underlying KVStore.
@@ -58,6 +56,25 @@ func (m *Manager) ClearLedgerState() (err error) {
 	return m.store.Clear()
 }
 
+func (m *Manager) ReInitStateTreeWithLocking() {
+	m.WriteLockLedger()
+	defer m.WriteUnlockLedger()
+
+	m.reInitStateTreeWithoutLocking()
+}
+
+func (m *Manager) reInitStateTreeWithoutLocking() {
+	m.stateTree = ads.NewMap[iotago.Identifier](
+		m.stateTreeKVStore,
+		iotago.Identifier.Bytes,
+		iotago.IdentifierFromBytes,
+		iotago.OutputID.Bytes,
+		iotago.OutputIDFromBytes,
+		(*stateTreeMetadata).Bytes,
+		stateMetadataFromBytes,
+	)
+}
+
 func (m *Manager) ReadLockLedger() {
 	m.storeLock.RLock()
 }
@@ -149,30 +166,30 @@ func (m *Manager) ReadLedgerSlot() (iotago.SlotIndex, error) {
 	return m.ReadLedgerIndexWithoutLocking()
 }
 
-func (m *Manager) ApplyDiffWithoutLocking(slot iotago.SlotIndex, newOutputs Outputs, newSpents Spents) error {
+func (m *Manager) ApplyDiffWithoutLocking(slot iotago.SlotIndex, newOutputs Outputs, newSpents Spents) (newStateTreeRoot iotago.Identifier, error error) {
 	mutations, err := m.store.Batched()
 	if err != nil {
-		return err
+		return iotago.EmptyIdentifier, err
 	}
 
 	for _, output := range newOutputs {
 		if err = storeOutput(output, mutations); err != nil {
 			mutations.Cancel()
 
-			return err
+			return iotago.EmptyIdentifier, err
 		}
 		if err := markAsUnspent(output, mutations); err != nil {
 			mutations.Cancel()
 
-			return err
+			return iotago.EmptyIdentifier, err
 		}
 	}
 
 	for _, spent := range newSpents {
 		if err := storeSpentAndMarkOutputAsSpent(spent, mutations); err != nil {
 			mutations.Cancel()
 
-			return err
+			return iotago.EmptyIdentifier, err
 		}
 	}
 
@@ -185,38 +202,44 @@ func (m *Manager) ApplyDiffWithoutLocking(slot iotago.SlotIndex, newOutputs Outp
 	if err := storeDiff(slotDiff, mutations); err != nil {
 		mutations.Cancel()
 
-		return err
+		return iotago.EmptyIdentifier, err
 	}
 
 	if err := storeLedgerIndex(slot, mutations); err != nil {
 		mutations.Cancel()
 
-		return err
+		return iotago.EmptyIdentifier, err
 	}
 
 	if err := mutations.Commit(); err != nil {
-		return err
+		return iotago.EmptyIdentifier, err
 	}
 
 	for _, output := range newOutputs {
 		if err := m.stateTree.Set(output.OutputID(), newStateMetadata(output)); err != nil {
-			return ierrors.Wrapf(err, "failed to set new oputput in state tree, outputID: %s", output.OutputID().ToHex())
+			return iotago.EmptyIdentifier, ierrors.Wrapf(err, "failed to set new oputput in state tree, outputID: %s", output.OutputID().ToHex())
 		}
 	}
 	for _, spent := range newSpents {
 		if _, err := m.stateTree.Delete(spent.OutputID()); err != nil {
-			return ierrors.Wrapf(err, "failed to delete spent output from state tree, outputID: %s", spent.OutputID().ToHex())
+			return iotago.EmptyIdentifier, ierrors.Wrapf(err, "failed to delete spent output from state tree, outputID: %s", spent.OutputID().ToHex())
 		}
 	}
 
 	if err := m.stateTree.Commit(); err != nil {
-		return ierrors.Wrap(err, "failed to commit state tree")
+		return iotago.EmptyIdentifier, ierrors.Wrap(err, "failed to commit state tree")
 	}
 
-	return nil
+	root := m.StateTreeRoot()
+
+	// Re-Init the state tree to release memory from the underlying SMT of all the elements that were just added.
+	// This might increase runtime (as we need to access the disk) but it will reduce memory usage significantly for a big tree.
+	m.reInitStateTreeWithoutLocking()
+
+	return root, nil
 }
 
-func (m *Manager) ApplyDiff(slot iotago.SlotIndex, newOutputs Outputs, newSpents Spents) error {
+func (m *Manager) ApplyDiff(slot iotago.SlotIndex, newOutputs Outputs, newSpents Spents) (newStateTreeRoot iotago.Identifier, error error) {
 	m.WriteLockLedger()
 	defer m.WriteUnlockLedger()
 
@@ -289,6 +312,8 @@ func (m *Manager) RollbackDiffWithoutLocking(slot iotago.SlotIndex, newOutputs O
 		return ierrors.Wrap(err, "failed to commit state tree")
 	}
 
+	m.reInitStateTreeWithoutLocking()
+
 	return nil
 }
 
@@ -321,6 +346,8 @@ func (m *Manager) AddGenesisUnspentOutputWithoutLocking(unspentOutput *Output) e
 		return ierrors.Wrap(err, "failed to commit state tree")
 	}
 
+	m.reInitStateTreeWithoutLocking()
+
 	return nil
 }
 

pkg/protocol/engine/utxoledger/manager_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/iotaledger/hive.go/db"
 	"github.com/iotaledger/hive.go/ierrors"
 	"github.com/iotaledger/hive.go/kvstore/mapdb"
+	"github.com/iotaledger/hive.go/lo"
 	"github.com/iotaledger/iota-core/pkg/protocol/engine/utxoledger"
 	"github.com/iotaledger/iota-core/pkg/protocol/engine/utxoledger/tpkg"
 	"github.com/iotaledger/iota-core/pkg/storage/database"
@@ -41,7 +42,7 @@ func TestConfirmationApplyAndRollbackToEmptyLedger(t *testing.T) {
 		tpkg.RandLedgerStateSpentWithOutput(outputs[2], slot),
 	}
 
-	require.NoError(t, manager.ApplyDiffWithoutLocking(slot, outputs, spents))
+	require.NoError(t, lo.Return2(manager.ApplyDiffWithoutLocking(slot, outputs, spents)))
 
 	require.NotEqual(t, manager.StateTreeRoot(), iotago.Identifier{})
 	require.True(t, manager.CheckStateTree())
@@ -107,7 +108,7 @@ func TestConfirmationApplyAndRollbackToPreviousLedger(t *testing.T) {
 	previousSpents := utxoledger.Spents{
 		tpkg.RandLedgerStateSpentWithOutput(previousOutputs[1], previousMsIndex),
 	}
-	require.NoError(t, manager.ApplyDiffWithoutLocking(previousMsIndex, previousOutputs, previousSpents))
+	require.NoError(t, lo.Return2(manager.ApplyDiffWithoutLocking(previousMsIndex, previousOutputs, previousSpents)))
 
 	require.True(t, manager.CheckStateTree())
 
@@ -129,7 +130,7 @@ func TestConfirmationApplyAndRollbackToPreviousLedger(t *testing.T) {
 		tpkg.RandLedgerStateSpentWithOutput(previousOutputs[2], index),
 		tpkg.RandLedgerStateSpentWithOutput(outputs[2], index),
 	}
-	require.NoError(t, manager.ApplyDiffWithoutLocking(index, outputs, spents))
+	require.NoError(t, lo.Return2(manager.ApplyDiffWithoutLocking(index, outputs, spents)))
 
 	require.True(t, manager.CheckStateTree())
 

pkg/protocol/engine/utxoledger/output_test.go

@@ -41,7 +41,7 @@ func AssertOutputUnspentAndSpentTransitions(t *testing.T, output *utxoledger.Out
 	require.True(t, has)
 
 	// Spent it with a slot.
-	require.NoError(t, manager.ApplyDiff(spent.SlotSpent(), utxoledger.Outputs{}, utxoledger.Spents{spent}))
+	require.NoError(t, lo.Return2(manager.ApplyDiff(spent.SlotSpent(), utxoledger.Outputs{}, utxoledger.Spents{spent})))
 
 	// Read Spent from DB and compare
 	readSpent, err := manager.ReadSpentForOutputIDWithoutLocking(outputID)

pkg/protocol/engine/utxoledger/slot_diff_test.go

@@ -80,7 +80,7 @@ func TestSlotDiffSerialization(t *testing.T) {
 		tpkg.RandLedgerStateSpentWithOutput(outputs[2], slot),
 	}
 
-	require.NoError(t, manager.ApplyDiffWithoutLocking(slot, outputs, spents))
+	require.NoError(t, lo.Return2(manager.ApplyDiffWithoutLocking(slot, outputs, spents)))
 
 	readDiff, err := manager.SlotDiffWithoutLocking(slot)
 	require.NoError(t, err)

pkg/protocol/engine/utxoledger/snapshot.go

@@ -215,6 +215,8 @@ func (m *Manager) Import(reader io.ReadSeeker) error {
 		return ierrors.Wrap(err, "unable to commit state tree")
 	}
 
+	m.reInitStateTreeWithoutLocking()
+
 	return nil
 }
 
@@ -306,5 +308,7 @@ func (m *Manager) Rollback(targetSlot iotago.SlotIndex) error {
 		return ierrors.Wrap(err, "unable to commit state tree")
 	}
 
+	m.reInitStateTreeWithoutLocking()
+
 	return nil
 }

pkg/protocol/engine/utxoledger/snapshot_test.go

@@ -235,13 +235,13 @@ func TestManager_Import(t *testing.T) {
 	require.NoError(t, kvstore.Copy(mapDB, mapDBAtSlot0))
 
 	output2 := tpkg.RandLedgerStateOutput()
-	require.NoError(t, manager.ApplyDiff(1,
+	require.NoError(t, lo.Return2(manager.ApplyDiff(1,
 		utxoledger.Outputs{
 			output2,
 			tpkg.RandLedgerStateOutput(),
 		}, utxoledger.Spents{
 			tpkg.RandLedgerStateSpentWithOutput(output1, 1),
-		}))
+		})))
 
 	ledgerSlot, err = manager.ReadLedgerSlot()
 	require.NoError(t, err)
@@ -250,14 +250,14 @@ func TestManager_Import(t *testing.T) {
 	mapDBAtSlot1 := mapdb.NewMapDB()
 	require.NoError(t, kvstore.Copy(mapDB, mapDBAtSlot1))
 
-	require.NoError(t, manager.ApplyDiff(2,
+	require.NoError(t, lo.Return2(manager.ApplyDiff(2,
 		utxoledger.Outputs{
 			tpkg.RandLedgerStateOutput(),
 			tpkg.RandLedgerStateOutput(),
 			tpkg.RandLedgerStateOutput(),
 		}, utxoledger.Spents{
 			tpkg.RandLedgerStateSpentWithOutput(output2, 2),
-		}))
+		})))
 
 	ledgerSlot, err = manager.ReadLedgerSlot()
 	require.NoError(t, err)
@@ -325,26 +325,26 @@ func TestManager_Export(t *testing.T) {
 	require.NoError(t, manager.AddGenesisUnspentOutput(tpkg.RandLedgerStateOutput()))
 
 	output2 := tpkg.RandLedgerStateOutput()
-	require.NoError(t, manager.ApplyDiff(1,
+	require.NoError(t, lo.Return2(manager.ApplyDiff(1,
 		utxoledger.Outputs{
 			output2,
 			tpkg.RandLedgerStateOutput(),
 		}, utxoledger.Spents{
 			tpkg.RandLedgerStateSpentWithOutput(output1, 1),
-		}))
+		})))
 
 	ledgerSlot, err := manager.ReadLedgerSlot()
 	require.NoError(t, err)
 	require.Equal(t, iotago.SlotIndex(1), ledgerSlot)
 
-	require.NoError(t, manager.ApplyDiff(2,
+	require.NoError(t, lo.Return2(manager.ApplyDiff(2,
 		utxoledger.Outputs{
 			tpkg.RandLedgerStateOutput(),
 			tpkg.RandLedgerStateOutput(),
 			tpkg.RandLedgerStateOutput(),
 		}, utxoledger.Spents{
 			tpkg.RandLedgerStateSpentWithOutput(output2, 2),
-		}))
+		})))
 
 	ledgerSlot, err = manager.ReadLedgerSlot()
 	require.NoError(t, err)