sweep: introduce BudgetAggregator to cluster inputs by deadlines

This commit adds `BudgetAggregator` as a new implementation of
`UtxoAggregator`. This aggregator will group inputs by their deadline
heights and create input sets that can be used directly by the fee
bumper for fee calculations.

sweep/aggregator.go

@@ -3,7 +3,10 @@ package sweep
 import (
 	"sort"
 
+	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/fn"
+	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 )
 
@@ -461,3 +464,232 @@ func zipClusters(as, bs []inputCluster) []inputCluster {
 
 	return finalClusters
 }
+
+// BudgetAggregator is a budget-based aggregator that creates clusters based on
+// deadlines and budgets of inputs.
+type BudgetAggregator struct {
+	// estimator is used when crafting sweep transactions to estimate the
+	// necessary fee relative to the expected size of the sweep
+	// transaction.
+	estimator chainfee.Estimator
+
+	// maxInputs specifies the maximum number of inputs allowed in a single
+	// sweep tx.
+	maxInputs uint32
+}
+
+// Compile-time constraint to ensure BudgetAggregator implements UtxoAggregator.
+var _ UtxoAggregator = (*BudgetAggregator)(nil)
+
+// NewBudgetAggregator creates a new instance of a BudgetAggregator.
+func NewBudgetAggregator(estimator chainfee.Estimator,
+	maxInputs uint32) *BudgetAggregator {
+
+	return &BudgetAggregator{
+		estimator: estimator,
+		maxInputs: maxInputs,
+	}
+}
+
+// clusterGroup defines an alias for a set of inputs that are to be grouped.
+type clusterGroup map[fn.Option[int32]][]pendingInput
+
+// ClusterInputs creates a list of input sets from pending inputs.
+// 1. filter out inputs whose budget cannot cover min relay fee.
+// 2. group the inputs into clusters based on their deadline height.
+// 3. sort the inputs in each cluster by their budget.
+// 4. optionally split a cluster if it exceeds the max input limit.
+// 5. create input sets from each of the clusters.
+func (b *BudgetAggregator) ClusterInputs(inputs pendingInputs) []InputSet {
+	// Filter out inputs that have a budget below min relay fee.
+	filteredInputs := b.filterInputs(inputs)
+
+	// Create clusters to group inputs based on their deadline height.
+	clusters := make(clusterGroup, len(filteredInputs))
+
+	// Iterate all the inputs and group them based on their specified
+	// deadline heights.
+	for _, input := range filteredInputs {
+		height := input.params.DeadlineHeight
+		cluster, ok := clusters[height]
+		if !ok {
+			cluster = make([]pendingInput, 0)
+		}
+
+		cluster = append(cluster, *input)
+		clusters[height] = cluster
+	}
+
+	// Now that we have the clusters, we can create the input sets.
+	//
+	// NOTE: cannot pre-allocate the slice since we don't know the number
+	// of input sets in advance.
+	inputSets := make([]InputSet, 0)
+	for _, cluster := range clusters {
+		// Sort the inputs by their economical value.
+		sortedInputs := b.sortInputs(cluster)
+
+		// Create input sets from the cluster.
+		sets := b.createInputSets(sortedInputs)
+		inputSets = append(inputSets, sets...)
+	}
+
+	return inputSets
+}
+
+// createInputSet takes a set of inputs which share the same deadline height
+// and turns them into a list of `InputSet`, each set is then used to create a
+// sweep transaction.
+func (b *BudgetAggregator) createInputSets(inputs []pendingInput) []InputSet {
+	// sets holds the InputSets that we will return.
+	sets := make([]InputSet, 0)
+
+	// Copy the inputs to a new slice so we can modify it.
+	remainingInputs := make([]pendingInput, len(inputs))
+	copy(remainingInputs, inputs)
+
+	// If the number of inputs is greater than the max inputs allowed, we
+	// will split them into smaller clusters.
+	for uint32(len(remainingInputs)) > b.maxInputs {
+		log.Tracef("Cluster has %v inputs, max is %v, dividing...",
+			len(inputs), b.maxInputs)
+
+		// Copy the inputs to be put into the new set, and update the
+		// remaining inputs by removing currentInputs.
+		currentInputs := make([]pendingInput, b.maxInputs)
+		copy(currentInputs, remainingInputs[:b.maxInputs])
+		remainingInputs = remainingInputs[b.maxInputs:]
+
+		// Create an InputSet using the max allowed number of inputs.
+		set, err := NewBudgetInputSet(currentInputs)
+		if err != nil {
+			log.Errorf("unable to create input set: %v", err)
+
+			continue
+		}
+
+		sets = append(sets, set)
+	}
+
+	// Create an InputSet from the remaining inputs.
+	if len(remainingInputs) > 0 {
+		set, err := NewBudgetInputSet(remainingInputs)
+		if err != nil {
+			log.Errorf("unable to create input set: %v", err)
+			return nil
+		}
+
+		sets = append(sets, set)
+	}
+
+	return sets
+}
+
+// filterInputs filters out inputs that have a budget below the min relay fee
+// or have a required output that's below the dust.
+func (b *BudgetAggregator) filterInputs(inputs pendingInputs) pendingInputs {
+	// Get the current min relay fee for this round.
+	minFeeRate := b.estimator.RelayFeePerKW()
+
+	// filterInputs stores a map of inputs that has a budget that at least
+	// can pay the minimal fee.
+	filteredInputs := make(pendingInputs, len(inputs))
+
+	// Iterate all the inputs and filter out the ones whose budget cannot
+	// cover the min fee.
+	for _, pi := range inputs {
+		op := pi.OutPoint()
+
+		// Get the size and skip if there's an error.
+		size, _, err := pi.WitnessType().SizeUpperBound()
+		if err != nil {
+			log.Warnf("Skipped input=%v: cannot get its size: %v",
+				op, err)
+
+			continue
+		}
+
+		// Skip inputs that has too little budget.
+		minFee := minFeeRate.FeeForWeight(int64(size))
+		if pi.params.Budget < minFee {
+			log.Warnf("Skipped input=%v: has budget=%v, but the "+
+				"min fee requires %v", op, pi.params.Budget,
+				minFee)
+
+			continue
+		}
+
+		// If the input comes with a required tx out that is below
+		// dust, we won't add it.
+		//
+		// NOTE: only HtlcSecondLevelAnchorInput returns non-nil
+		// RequiredTxOut.
+		reqOut := pi.RequiredTxOut()
+		if reqOut != nil {
+			if isDustOutput(reqOut) {
+				log.Errorf("Rejected input=%v due to dust "+
+					"required output=%v", op, reqOut.Value)
+
+				continue
+			}
+		}
+
+		filteredInputs[*op] = pi
+	}
+
+	return filteredInputs
+}
+
+// sortInputs sorts the inputs based on their economical value.
+//
+// NOTE: besides the forced inputs, the sorting won't make any difference
+// because all the inputs are added to the same set. The exception is when the
+// number of inputs exceeds the maxInputs limit, it requires us to split them
+// into smaller clusters. In that case, the sorting will make a difference as
+// the budgets of the clusters will be different.
+func (b *BudgetAggregator) sortInputs(inputs []pendingInput) []pendingInput {
+	// sortedInputs is the final list of inputs sorted by their economical
+	// value.
+	sortedInputs := make([]pendingInput, 0, len(inputs))
+
+	// Copy the inputs.
+	sortedInputs = append(sortedInputs, inputs...)
+
+	// Sort the inputs based on their budgets.
+	//
+	// NOTE: We can implement more sophisticated algorithm as the budget
+	// left is a function f(minFeeRate, size) = b1 - s1 * r > b2 - s2 * r,
+	// where b1 and b2 are budgets, s1 and s2 are sizes of the inputs.
+	sort.Slice(sortedInputs, func(i, j int) bool {
+		left := sortedInputs[i].params.Budget
+		right := sortedInputs[j].params.Budget
+
+		// Make sure forced inputs are always put in the front.
+		leftForce := sortedInputs[i].params.Force
+		rightForce := sortedInputs[j].params.Force
+
+		// If both are forced inputs, we return the one with the higher
+		// budget. If neither are forced inputs, we also return the one
+		// with the higher budget.
+		if leftForce == rightForce {
+			return left > right
+		}
+
+		// Otherwise, it's either the left or the right is forced. We
+		// can simply return `leftForce` here as, if it's true, the
+		// left is forced and should be put in the front. Otherwise,
+		// the right is forced and should be put in the front.
+		return leftForce
+	})
+
+	return sortedInputs
+}
+
+// isDustOutput checks if the given output is considered as dust.
+func isDustOutput(output *wire.TxOut) bool {
+	// Fetch the dust limit for this output.
+	dustLimit := lnwallet.DustLimitForSize(len(output.PkScript))
+
+	// If the output is below the dust limit, we consider it dust.
+	return btcutil.Amount(output.Value) < dustLimit
+}