UTXO Trimming

consensus/blake3pow/consensus.go

@@ -5,6 +5,8 @@ import (
 	"math/big"
 	"runtime"
 	"runtime/debug"
+	"sort"
+	"sync"
 	"time"
 
 	mapset "github.com/deckarep/golang-set"
@@ -16,10 +18,12 @@ import (
 	"github.com/dominant-strategies/go-quai/core/state"
 	"github.com/dominant-strategies/go-quai/core/types"
 	"github.com/dominant-strategies/go-quai/core/vm"
+	"github.com/dominant-strategies/go-quai/ethdb"
 	"github.com/dominant-strategies/go-quai/log"
 	"github.com/dominant-strategies/go-quai/multiset"
 	"github.com/dominant-strategies/go-quai/params"
 	"github.com/dominant-strategies/go-quai/trie"
+	"google.golang.org/protobuf/proto"
 	"modernc.org/mathutil"
 )
 
@@ -603,16 +607,17 @@ func (blake3pow *Blake3pow) Prepare(chain consensus.ChainHeaderReader, header *t
 
 // Finalize implements consensus.Engine, accumulating the block and uncle rewards,
 // setting the final state on the header
-func (blake3pow *Blake3pow) Finalize(chain consensus.ChainHeaderReader, header *types.WorkObject, state *state.StateDB, setRoots bool, utxosCreate, utxosDelete []common.Hash) (*multiset.MultiSet, error) {
+// Finalize returns the new MuHash of the UTXO set, the new size of the UTXO set and an error if any
+func (blake3pow *Blake3pow) Finalize(chain consensus.ChainHeaderReader, batch ethdb.Batch, header *types.WorkObject, state *state.StateDB, setRoots bool, utxoSetSize uint64, utxosCreate, utxosDelete []common.Hash) (*multiset.MultiSet, uint64, error) {
 	nodeLocation := blake3pow.config.NodeLocation
 	nodeCtx := blake3pow.config.NodeLocation.Context()
 	var multiSet *multiset.MultiSet
-	if nodeCtx == common.ZONE_CTX && chain.IsGenesisHash(header.ParentHash(nodeCtx)) {
+	if chain.IsGenesisHash(header.ParentHash(nodeCtx)) {
 		multiSet = multiset.New()
 		// Create the lockup contract account
 		lockupContract, err := vm.LockupContractAddresses[[2]byte{nodeLocation[0], nodeLocation[1]}].InternalAndQuaiAddress()
 		if err != nil {
-			return nil, err
+			return nil, 0, err
 		}
 		state.CreateAccount(lockupContract)
 		state.SetNonce(lockupContract, 1) // so it's not considered "empty"
@@ -649,31 +654,293 @@ func (blake3pow *Blake3pow) Finalize(chain consensus.ChainHeaderReader, header *
 		multiSet = rawdb.ReadMultiSet(chain.Database(), header.ParentHash(nodeCtx))
 	}
 	if multiSet == nil {
-		return nil, fmt.Errorf("Multiset is nil for block %s", header.ParentHash(nodeCtx).String())
+		return nil, 0, fmt.Errorf("Multiset is nil for block %s", header.ParentHash(nodeCtx).String())
+	}
+
+	utxoSetSize += uint64(len(utxosCreate))
+	if utxoSetSize < uint64(len(utxosDelete)) {
+		return nil, 0, fmt.Errorf("UTXO set size is less than the number of utxos to delete. This is a bug. UTXO set size: %d, UTXOs to delete: %d", utxoSetSize, len(utxosDelete))
+	}
+	utxoSetSize -= uint64(len(utxosDelete))
+
+	trimDepths := types.TrimDepths
+	if utxoSetSize > params.SoftMaxUTXOSetSize/2 {
+		var err error
+		trimDepths, err = rawdb.ReadTrimDepths(chain.Database(), header.ParentHash(nodeCtx))
+		if err != nil || trimDepths == nil {
+			blake3pow.logger.Errorf("Failed to read trim depths for block %s: %+v", header.ParentHash(nodeCtx).String(), err)
+			trimDepths = make(map[uint8]uint64, len(types.TrimDepths))
+			for denomination, depth := range types.TrimDepths { // copy the default trim depths
+				trimDepths[denomination] = depth
+			}
+		}
+		if UpdateTrimDepths(trimDepths, utxoSetSize) {
+			blake3pow.logger.Infof("Updated trim depths at height %d new depths: %+v", header.NumberU64(nodeCtx), trimDepths)
+		}
+		if !setRoots {
+			rawdb.WriteTrimDepths(batch, header.Hash(), trimDepths)
+		}
+	}
+	start := time.Now()
+	collidingKeys, err := rawdb.ReadCollidingKeys(chain.Database(), header.ParentHash(nodeCtx))
+	if err != nil {
+		blake3pow.logger.Errorf("Failed to read colliding keys for block %s: %+v", header.ParentHash(nodeCtx).String(), err)
+	}
+	newCollidingKeys := make([][]byte, 0)
+	trimmedUtxos := make([]*types.SpentUtxoEntry, 0)
+	var wg sync.WaitGroup
+	var lock sync.Mutex
+	for denomination, depth := range trimDepths {
+		if denomination <= types.MaxTrimDenomination && header.NumberU64(nodeCtx) > depth+params.MinimumTrimDepth {
+			wg.Add(1)
+			go func(denomination uint8, depth uint64) {
+				defer func() {
+					if r := recover(); r != nil {
+						blake3pow.logger.WithFields(log.Fields{
+							"error":      r,
+							"stacktrace": string(debug.Stack()),
+						}).Error("Go-Quai Panicked")
+					}
+				}()
+				nextBlockToTrim := rawdb.ReadCanonicalHash(chain.Database(), header.NumberU64(nodeCtx)-depth)
+				collisions := TrimBlock(chain, batch, denomination, true, header.NumberU64(nodeCtx)-depth, nextBlockToTrim, &utxosDelete, &trimmedUtxos, nil, &utxoSetSize, !setRoots, &lock, blake3pow.logger) // setRoots is false when we are processing the block
+				if len(collisions) > 0 {
+					lock.Lock()
+					newCollidingKeys = append(newCollidingKeys, collisions...)
+					lock.Unlock()
+				}
+				wg.Done()
+			}(denomination, depth)
+		}
+	}
+	if len(collidingKeys) > 0 {
+		wg.Add(1)
+		go func() {
+			defer func() {
+				if r := recover(); r != nil {
+					blake3pow.logger.WithFields(log.Fields{
+						"error":      r,
+						"stacktrace": string(debug.Stack()),
+					}).Error("Go-Quai Panicked")
+				}
+			}()
+			// Trim colliding/duplicate keys here - an optimization could be to do this above in parallel with the other trims
+			collisions := TrimBlock(chain, batch, 0, false, 0, common.Hash{}, &utxosDelete, &trimmedUtxos, collidingKeys, &utxoSetSize, !setRoots, &lock, blake3pow.logger)
+			if len(collisions) > 0 {
+				lock.Lock()
+				newCollidingKeys = append(newCollidingKeys, collisions...)
+				lock.Unlock()
+			}
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	if len(trimmedUtxos) > 0 {
+		blake3pow.logger.Infof("Trimmed %d UTXOs from db in %s", len(trimmedUtxos), common.PrettyDuration(time.Since(start)))
+	}
+	if !setRoots {
+		rawdb.WriteTrimmedUTXOs(batch, header.Hash(), trimmedUtxos)
+		if len(newCollidingKeys) > 0 {
+			rawdb.WriteCollidingKeys(batch, header.Hash(), newCollidingKeys)
+		}
 	}
 	for _, hash := range utxosCreate {
 		multiSet.Add(hash.Bytes())
 	}
 	for _, hash := range utxosDelete {
 		multiSet.Remove(hash.Bytes())
 	}
-	blake3pow.logger.Infof("Parent hash: %s, header hash: %s, muhash: %s, block height: %d, setroots: %t, UtxosCreated: %d, UtxosDeleted: %d", header.ParentHash(nodeCtx).String(), header.Hash().String(), multiSet.Hash().String(), header.NumberU64(nodeCtx), setRoots, len(utxosCreate), len(utxosDelete))
+	blake3pow.logger.Infof("Parent hash: %s, header hash: %s, muhash: %s, block height: %d, setroots: %t, UtxosCreated: %d, UtxosDeleted: %d, UTXOs Trimmed from DB: %d, UTXO Set Size: %d", header.ParentHash(nodeCtx).String(), header.Hash().String(), multiSet.Hash().String(), header.NumberU64(nodeCtx), setRoots, len(utxosCreate), len(utxosDelete), len(trimmedUtxos), utxoSetSize)
+
 	if setRoots {
 		header.Header().SetUTXORoot(multiSet.Hash())
 		header.Header().SetEVMRoot(state.IntermediateRoot(true))
 		header.Header().SetEtxSetRoot(state.ETXRoot())
 		header.Header().SetQuaiStateSize(state.GetQuaiTrieSize())
 	}
-	return multiSet, nil
+	return multiSet, utxoSetSize, nil
+}
+
+// TrimBlock trims all UTXOs of a given denomination that were created in a given block.
+// In the event of an attacker intentionally creating too many 9-byte keys that collide, we return the colliding keys to be trimmed in the next block.
+func TrimBlock(chain consensus.ChainHeaderReader, batch ethdb.Batch, denomination uint8, checkDenom bool, blockHeight uint64, blockHash common.Hash, utxosDelete *[]common.Hash, trimmedUtxos *[]*types.SpentUtxoEntry, collidingKeys [][]byte, utxoSetSize *uint64, deleteFromDb bool, lock *sync.Mutex, logger *log.Logger) [][]byte {
+	utxosCreated, err := rawdb.ReadPrunedUTXOKeys(chain.Database(), blockHeight)
+	if err != nil {
+		logger.Errorf("Failed to read pruned UTXOs for block %d: %+v", blockHeight, err)
+	}
+	if len(utxosCreated) == 0 {
+		// This should almost never happen, but we need to handle it
+		utxosCreated, err = rawdb.ReadCreatedUTXOKeys(chain.Database(), blockHash)
+		if err != nil {
+			logger.Errorf("Failed to read created UTXOs for block %d: %+v", blockHeight, err)
+		}
+		logger.Infof("Reading non-pruned UTXOs for block %d", blockHeight)
+		for i, key := range utxosCreated {
+			if len(key) == rawdb.UtxoKeyWithDenominationLength {
+				if key[len(key)-1] > types.MaxTrimDenomination {
+					// Don't keep it if the denomination is not trimmed
+					// The keys are sorted in order of denomination, so we can break here
+					break
+				}
+				key[rawdb.PrunedUtxoKeyWithDenominationLength+len(rawdb.UtxoPrefix)-1] = key[len(key)-1] // place the denomination at the end of the pruned key (11th byte will become 9th byte)
+			}
+			// Reduce key size to 9 bytes and cut off the prefix
+			key = key[len(rawdb.UtxoPrefix) : rawdb.PrunedUtxoKeyWithDenominationLength+len(rawdb.UtxoPrefix)]
+			utxosCreated[i] = key
+		}
+	}
+
+	logger.Infof("UTXOs created in block %d: %d", blockHeight, len(utxosCreated))
+	if len(collidingKeys) > 0 {
+		logger.Infof("Colliding keys: %d", len(collidingKeys))
+		utxosCreated = append(utxosCreated, collidingKeys...)
+		sort.Slice(utxosCreated, func(i, j int) bool {
+			return utxosCreated[i][len(utxosCreated[i])-1] < utxosCreated[j][len(utxosCreated[j])-1]
+		})
+	}
+	newCollisions := make([][]byte, 0)
+	duplicateKeys := make(map[[36]byte]bool) // cannot use rawdb.UtxoKeyLength for map as it's not const
+	// Start by grabbing all the UTXOs created in the block (that are still in the UTXO set)
+	for _, key := range utxosCreated {
+		if len(key) != rawdb.PrunedUtxoKeyWithDenominationLength {
+			continue
+		}
+		if checkDenom {
+			if key[len(key)-1] != denomination {
+				if key[len(key)-1] > denomination {
+					break // The keys are stored in order of denomination, so we can stop checking here
+				} else {
+					continue
+				}
+			} else {
+				key = append(rawdb.UtxoPrefix, key...) // prepend the db prefix
+				key = key[:len(key)-1]                 // remove the denomination byte
+			}
+		}
+		// Check key in database
+		i := 0
+		it := chain.Database().NewIterator(key, nil)
+		for it.Next() {
+			data := it.Value()
+			if len(data) == 0 {
+				logger.Infof("Empty key found, denomination: %d", denomination)
+				continue
+			}
+			// Check if the key is a duplicate
+			if len(it.Key()) == rawdb.UtxoKeyLength {
+				key36 := [36]byte(it.Key())
+				if duplicateKeys[key36] {
+					continue
+				} else {
+					duplicateKeys[key36] = true
+				}
+			} else {
+				logger.Errorf("Invalid key length: %d", len(it.Key()))
+				continue
+			}
+			utxoProto := new(types.ProtoTxOut)
+			if err := proto.Unmarshal(data, utxoProto); err != nil {
+				logger.Errorf("Failed to unmarshal ProtoTxOut: %+v data: %+v key: %+v", err, data, key)
+				continue
+			}
+
+			utxo := new(types.UtxoEntry)
+			if err := utxo.ProtoDecode(utxoProto); err != nil {
+				logger.WithFields(log.Fields{
+					"key":  key,
+					"data": data,
+					"err":  err,
+				}).Error("Invalid utxo Proto")
+				continue
+			}
+			if checkDenom && utxo.Denomination != denomination {
+				continue
+			}
+			txHash, index, err := rawdb.ReverseUtxoKey(it.Key())
+			if err != nil {
+				logger.WithField("err", err).Error("Failed to parse utxo key")
+				continue
+			}
+			lock.Lock()
+			*utxosDelete = append(*utxosDelete, types.UTXOHash(txHash, index, utxo))
+			if deleteFromDb {
+				batch.Delete(it.Key())
+				*trimmedUtxos = append(*trimmedUtxos, &types.SpentUtxoEntry{OutPoint: types.OutPoint{txHash, index}, UtxoEntry: utxo})
+			}
+			*utxoSetSize--
+			lock.Unlock()
+			i++
+			if i >= types.MaxTrimCollisionsPerKeyPerBlock {
+				// This will rarely ever happen, but if it does, we should continue trimming this key in the next block
+				logger.WithField("blockHeight", blockHeight).Error("MaxTrimCollisionsPerBlock exceeded")
+				newCollisions = append(newCollisions, key)
+				break
+			}
+		}
+		it.Release()
+	}
+	return newCollisions
+}
+
+func UpdateTrimDepths(trimDepths map[uint8]uint64, utxoSetSize uint64) bool {
+	switch {
+	case utxoSetSize > params.SoftMaxUTXOSetSize/2 && trimDepths[255] == 0: // 50% full
+		for denomination, depth := range trimDepths {
+			trimDepths[denomination] = depth - (depth / 10) // decrease lifespan of this denomination by 10%
+		}
+		trimDepths[255] = 1 // level 1
+	case utxoSetSize > params.SoftMaxUTXOSetSize-(params.SoftMaxUTXOSetSize/4) && trimDepths[255] == 1: // 75% full
+		for denomination, depth := range trimDepths {
+			trimDepths[denomination] = depth - (depth / 5) // decrease lifespan of this denomination by an additional 20%
+		}
+		trimDepths[255] = 2 // level 2
+	case utxoSetSize > params.SoftMaxUTXOSetSize-(params.SoftMaxUTXOSetSize/10) && trimDepths[255] == 2: // 90% full
+		for denomination, depth := range trimDepths {
+			trimDepths[denomination] = depth - (depth / 2) // decrease lifespan of this denomination by an additional 50%
+		}
+		trimDepths[255] = 3 // level 3
+	case utxoSetSize > params.SoftMaxUTXOSetSize && trimDepths[255] == 3:
+		for denomination, depth := range trimDepths {
+			trimDepths[denomination] = depth - (depth / 2) // decrease lifespan of this denomination by an additional 50%
+		}
+		trimDepths[255] = 4 // level 4
+
+	// Resets
+	case utxoSetSize <= params.SoftMaxUTXOSetSize/2 && trimDepths[255] == 1: // Below 50% full
+		for denomination, depth := range types.TrimDepths { // reset to the default trim depths
+			trimDepths[denomination] = depth
+		}
+		trimDepths[255] = 0 // level 0
+	case utxoSetSize <= params.SoftMaxUTXOSetSize-(params.SoftMaxUTXOSetSize/4) && trimDepths[255] == 2: // Below 75% full
+		for denomination, depth := range trimDepths {
+			trimDepths[denomination] = depth + (depth / 5) // increase lifespan of this denomination by 20%
+		}
+		trimDepths[255] = 1 // level 1
+	case utxoSetSize <= params.SoftMaxUTXOSetSize-(params.SoftMaxUTXOSetSize/10) && trimDepths[255] == 3: // Below 90% full
+		for denomination, depth := range trimDepths {
+			trimDepths[denomination] = depth + (depth / 2) // increase lifespan of this denomination by 50%
+		}
+		trimDepths[255] = 2 // level 2
+	case utxoSetSize <= params.SoftMaxUTXOSetSize && trimDepths[255] == 4: // Below 100% full
+		for denomination, depth := range trimDepths {
+			trimDepths[denomination] = depth + (depth / 2) // increase lifespan of this denomination by 50%
+		}
+		trimDepths[255] = 3 // level 3
+	default:
+		return false
+	}
+	return true
 }
 
 // FinalizeAndAssemble implements consensus.Engine, accumulating the block and
 // uncle rewards, setting the final state and assembling the block.
-func (blake3pow *Blake3pow) FinalizeAndAssemble(chain consensus.ChainHeaderReader, header *types.WorkObject, state *state.StateDB, txs []*types.Transaction, uncles []*types.WorkObjectHeader, etxs []*types.Transaction, subManifest types.BlockManifest, receipts []*types.Receipt, utxosCreate, utxosDelete []common.Hash) (*types.WorkObject, error) {
+func (blake3pow *Blake3pow) FinalizeAndAssemble(chain consensus.ChainHeaderReader, header *types.WorkObject, state *state.StateDB, txs []*types.Transaction, uncles []*types.WorkObjectHeader, etxs []*types.Transaction, subManifest types.BlockManifest, receipts []*types.Receipt, parentUtxoSetSize uint64, utxosCreate, utxosDelete []common.Hash) (*types.WorkObject, error) {
 	nodeCtx := blake3pow.config.NodeLocation.Context()
 	if nodeCtx == common.ZONE_CTX && chain.ProcessingState() {
 		// Finalize block
-		blake3pow.Finalize(chain, header, state, true, utxosCreate, utxosDelete)
+		if _, _, err := blake3pow.Finalize(chain, nil, header, state, true, parentUtxoSetSize, utxosCreate, utxosDelete); err != nil {
+			return nil, err
+		}
 	}
 
 	woBody, err := types.NewWorkObjectBody(header.Header(), txs, etxs, uncles, subManifest, receipts, trie.NewStackTrie(nil), nodeCtx)

consensus/consensus.go

@@ -189,14 +189,14 @@ type Engine interface {
 	//
 	// Note: The block header and state database might be updated to reflect any
 	// consensus rules that happen at finalization (e.g. block rewards).
-	Finalize(chain ChainHeaderReader, header *types.WorkObject, state *state.StateDB, setRoots bool, utxosCreate, utxosDelete []common.Hash) (*multiset.MultiSet, error)
+	Finalize(chain ChainHeaderReader, batch ethdb.Batch, header *types.WorkObject, state *state.StateDB, setRoots bool, parentUtxoSetSize uint64, utxosCreate, utxosDelete []common.Hash) (*multiset.MultiSet, uint64, error)
 
 	// FinalizeAndAssemble runs any post-transaction state modifications (e.g. block
 	// rewards) and assembles the final block.
 	//
 	// Note: The block header and state database might be updated to reflect any
 	// consensus rules that happen at finalization (e.g. block rewards).
-	FinalizeAndAssemble(chain ChainHeaderReader, woHeader *types.WorkObject, state *state.StateDB, txs []*types.Transaction, uncles []*types.WorkObjectHeader, etxs []*types.Transaction, subManifest types.BlockManifest, receipts []*types.Receipt, utxosCreate, utxosDelete []common.Hash) (*types.WorkObject, error)
+	FinalizeAndAssemble(chain ChainHeaderReader, woHeader *types.WorkObject, state *state.StateDB, txs []*types.Transaction, uncles []*types.WorkObjectHeader, etxs []*types.Transaction, subManifest types.BlockManifest, receipts []*types.Receipt, parentUtxoSetSize uint64, utxosCreate, utxosDelete []common.Hash) (*types.WorkObject, error)
 
 	// Seal generates a new sealing request for the given input block and pushes
 	// the result into the given channel.