Improve coin selection unit tests

tapdb/assets_store_test.go

@@ -36,6 +36,8 @@ type assetGenOptions struct {
 
 	groupAnchorGen *asset.Genesis
 
+	groupAnchorGenPoint *wire.OutPoint
+
 	noGroupKey bool
 
 	groupKeyPriv *btcec.PrivateKey
@@ -87,6 +89,12 @@ func withGroupAnchorGen(g *asset.Genesis) assetGenOpt {
 	}
 }
 
+func withGroupAnchorGenPoint(op wire.OutPoint) assetGenOpt {
+	return func(opt *assetGenOptions) {
+		opt.groupAnchorGenPoint = &op
+	}
+}
+
 func withAssetGenPoint(op wire.OutPoint) assetGenOpt {
 	return func(opt *assetGenOptions) {
 		opt.genesisPoint = op
@@ -151,6 +159,9 @@ func randAsset(t *testing.T, genOpts ...assetGenOpt) *asset.Asset {
 	if opts.groupAnchorGen != nil {
 		initialGen = *opts.groupAnchorGen
 	}
+	if opts.groupAnchorGenPoint != nil {
+		initialGen.FirstPrevOut = *opts.groupAnchorGenPoint
+	}
 
 	groupReq := asset.NewGroupKeyRequestNoErr(
 		t, groupKeyDesc, initialGen, protoAsset, nil,
@@ -433,6 +444,8 @@ type assetDesc struct {
 
 	groupAnchorGen *asset.Genesis
 
+	groupAnchorGenPoint *wire.OutPoint
+
 	anchorPoint wire.OutPoint
 
 	keyGroup *btcec.PrivateKey
@@ -548,6 +561,11 @@ func (a *assetGenerator) genAssets(t *testing.T, assetStore *AssetStore,
 				desc.groupAnchorGen,
 			))
 		}
+		if desc.groupAnchorGenPoint != nil {
+			opts = append(opts, withGroupAnchorGenPoint(
+				*desc.groupAnchorGenPoint,
+			))
+		}
 		if desc.assetVersion != nil {
 			opts = append(opts, withAssetVersionGen(
 				desc.assetVersion,
@@ -612,7 +630,7 @@ func (a *assetGenerator) bindAssetID(i int, op wire.OutPoint) *asset.ID {
 	return &id
 }
 
-func (a *assetGenerator) bindKeyGroup(i int, op wire.OutPoint) *btcec.PublicKey {
+func (a *assetGenerator) bindGroupKey(i int, op wire.OutPoint) *btcec.PublicKey {
 	gen := a.assetGens[i]
 	gen.FirstPrevOut = op
 	genTweak := gen.ID()
@@ -901,6 +919,7 @@ func TestSelectCommitment(t *testing.T) {
 		constraints tapfreighter.CommitmentConstraints
 
 		numAssets int
+		sum       int64
 
 		err error
 	}{
@@ -911,7 +930,7 @@ func TestSelectCommitment(t *testing.T) {
 			assets: []assetDesc{
 				{
 					assetGen: assetGen.assetGens[0],
-					amt:      5,
+					amt:      6,
 
 					anchorPoint: assetGen.anchorPoints[0],
 				},
@@ -923,6 +942,7 @@ func TestSelectCommitment(t *testing.T) {
 				MinAmt: 2,
 			},
 			numAssets: 1,
+			sum:       6,
 		},
 
 		// Asset matches all the params, but too small of a UTXO.  only
@@ -968,10 +988,10 @@ func TestSelectCommitment(t *testing.T) {
 			err:       tapfreighter.ErrMatchingAssetsNotFound,
 		},
 
-		// Create two assets, one has a key group the other doesn't.
+		// Create two assets, one has a group key the other doesn't.
 		// We should only get one asset back.
 		{
-			name: "asset with key group",
+			name: "asset with group key",
 			assets: []assetDesc{
 				{
 					assetGen: assetGen.assetGens[0],
@@ -983,19 +1003,20 @@ func TestSelectCommitment(t *testing.T) {
 				},
 				{
 					assetGen: assetGen.assetGens[1],
-					amt:      10,
+					amt:      12,
 
 					anchorPoint: assetGen.anchorPoints[1],
 					noGroupKey:  true,
 				},
 			},
 			constraints: tapfreighter.CommitmentConstraints{
-				GroupKey: assetGen.bindKeyGroup(
+				GroupKey: assetGen.bindGroupKey(
 					0, assetGen.anchorPoints[0],
 				),
 				MinAmt: 1,
 			},
 			numAssets: 1,
+			sum:       10,
 		},
 
 		// Leased assets shouldn't be returned, and neither should other
@@ -1027,6 +1048,48 @@ func TestSelectCommitment(t *testing.T) {
 			numAssets: 0,
 			err:       tapfreighter.ErrMatchingAssetsNotFound,
 		},
+
+		// Create three assets, the first two have a group key but
+		// different asset IDs, the other doesn't have a group key.
+		// We should only get the first two assets back.
+		{
+			name: "multiple different assets with same group key",
+			assets: []assetDesc{
+				{
+					assetGen: assetGen.assetGens[0],
+					amt:      10,
+
+					anchorPoint: assetGen.anchorPoints[0],
+
+					keyGroup: assetGen.groupKeys[0],
+				},
+				{
+					assetGen: assetGen.assetGens[1],
+					amt:      20,
+
+					anchorPoint: assetGen.anchorPoints[0],
+
+					keyGroup:            assetGen.groupKeys[0],
+					groupAnchorGen:      &assetGen.assetGens[0],
+					groupAnchorGenPoint: &assetGen.anchorPoints[0],
+				},
+				{
+					assetGen: assetGen.assetGens[1],
+					amt:      15,
+
+					anchorPoint: assetGen.anchorPoints[1],
+					noGroupKey:  true,
+				},
+			},
+			constraints: tapfreighter.CommitmentConstraints{
+				GroupKey: assetGen.bindGroupKey(
+					0, assetGen.anchorPoints[0],
+				),
+				MinAmt: 1,
+			},
+			numAssets: 2,
+			sum:       30,
+		},
 	}
 
 	ctx := context.Background()
@@ -1053,6 +1116,14 @@ func TestSelectCommitment(t *testing.T) {
 			// properly.
 			require.Equal(t, tc.numAssets, len(selectedAssets))
 
+			// Also verify the expected sum of asset amounts
+			// selected.
+			var sum int64
+			for _, a := range selectedAssets {
+				sum += int64(a.Asset.Amount)
+			}
+			require.Equal(t, tc.sum, sum)
+
 			// If the expectation is to get a single asset, let's
 			// make sure we can fetch the same asset commitment with
 			// the FetchCommitment method.

tapfreighter/coin_select.go

@@ -0,0 +1,174 @@
+package tapfreighter
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"sync"
+	"time"
+
+	"github.com/btcsuite/btcd/wire"
+	"github.com/lightninglabs/taproot-assets/fn"
+)
+
+// NewCoinSelect creates a new CoinSelect.
+func NewCoinSelect(coinLister CoinLister) *CoinSelect {
+	return &CoinSelect{
+		coinLister: coinLister,
+	}
+}
+
+// CoinSelect selects asset coins to spend in order to fund a send
+// transaction.
+type CoinSelect struct {
+	coinLister CoinLister
+
+	// coinLock is a read/write mutex that is used to ensure that only one
+	// goroutine is attempting to call any coin selection related methods at
+	// any time. This is necessary as some of the calls to the store (e.g.
+	// ListEligibleCoins -> LeaseCoin) are called after each other and
+	// cannot be placed within the same database transaction. So calls to
+	// those methods must hold this coin lock.
+	coinLock sync.Mutex
+}
+
+// SelectCoins returns a set of not yet leased coins that satisfy the given
+// constraints and strategy. The coins returned are leased for the default lease
+// duration.
+func (s *CoinSelect) SelectCoins(ctx context.Context,
+	constraints CommitmentConstraints,
+	strategy MultiCommitmentSelectStrategy) ([]*AnchoredCommitment, error) {
+
+	s.coinLock.Lock()
+	defer s.coinLock.Unlock()
+
+	// Before we select any coins, let's do some cleanup of expired leases.
+	if err := s.coinLister.DeleteExpiredLeases(ctx); err != nil {
+		return nil, fmt.Errorf("unable to delete expired leases: %w",
+			err)
+	}
+
+	listConstraints := CommitmentConstraints{
+		GroupKey: constraints.GroupKey,
+		AssetID:  constraints.AssetID,
+		MinAmt:   1,
+	}
+	eligibleCommitments, err := s.coinLister.ListEligibleCoins(
+		ctx, listConstraints,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to list eligible coins: %w", err)
+	}
+
+	log.Infof("Identified %v eligible asset inputs for send of %d to %v",
+		len(eligibleCommitments), constraints)
+
+	selectedCoins, err := s.selectForAmount(
+		constraints.MinAmt, eligibleCommitments, strategy,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to select coins: %w", err)
+	}
+
+	// We now need to lock/lease/reserve those selected coins so
+	// that they can't be used by other processes.
+	expiry := time.Now().Add(defaultCoinLeaseDuration)
+	coinOutPoints := fn.Map(
+		selectedCoins, func(c *AnchoredCommitment) wire.OutPoint {
+			return c.AnchorPoint
+		},
+	)
+	err = s.coinLister.LeaseCoins(
+		ctx, defaultWalletLeaseIdentifier, expiry, coinOutPoints...,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to lease coin: %w", err)
+	}
+
+	return selectedCoins, nil
+}
+
+// LeaseCoins leases/locks/reserves coins for the given lease owner until the
+// given expiry. This is used to prevent multiple concurrent coin selection
+// attempts from selecting the same coin(s).
+func (s *CoinSelect) LeaseCoins(ctx context.Context, leaseOwner [32]byte,
+	expiry time.Time, utxoOutpoints ...wire.OutPoint) error {
+
+	s.coinLock.Lock()
+	defer s.coinLock.Unlock()
+
+	return s.coinLister.LeaseCoins(
+		ctx, leaseOwner, expiry, utxoOutpoints...,
+	)
+}
+
+// ReleaseCoins releases/unlocks coins that were previously leased and makes
+// them available for coin selection again.
+func (s *CoinSelect) ReleaseCoins(ctx context.Context,
+	utxoOutpoints ...wire.OutPoint) error {
+
+	s.coinLock.Lock()
+	defer s.coinLock.Unlock()
+
+	return s.coinLister.ReleaseCoins(ctx, utxoOutpoints...)
+}
+
+// selectForAmount selects a subset of the given eligible commitments which
+// cumulatively sum to at least the minimum required amount. The selection
+// strategy determines how the commitments are selected.
+func (s *CoinSelect) selectForAmount(minTotalAmount uint64,
+	eligibleCommitments []*AnchoredCommitment,
+	strategy MultiCommitmentSelectStrategy) ([]*AnchoredCommitment,
+	error) {
+
+	// Select the first subset of eligible commitments which cumulatively
+	// sum to at least the minimum required amount.
+	var selectedCommitments []*AnchoredCommitment
+	amountSum := uint64(0)
+
+	switch strategy {
+	case PreferMaxAmount:
+		// Sort eligible commitments from the largest amount to
+		// smallest.
+		sort.Slice(
+			eligibleCommitments, func(i, j int) bool {
+				isLess := eligibleCommitments[i].Asset.Amount <
+					eligibleCommitments[j].Asset.Amount
+
+				// Negate the result to sort in descending
+				// order.
+				return !isLess
+			},
+		)
+
+		// Select the first subset of eligible commitments which
+		// cumulatively sum to at least the minimum required amount.
+		for _, anchoredCommitment := range eligibleCommitments {
+			selectedCommitments = append(
+				selectedCommitments, anchoredCommitment,
+			)
+
+			// Keep track of the total amount of assets we've seen
+			// so far.
+			amountSum += anchoredCommitment.Asset.Amount
+			if amountSum >= minTotalAmount {
+				// At this point a target min amount was
+				// specified and has been reached.
+				break
+			}
+		}
+
+	default:
+		return nil, fmt.Errorf("unknown multi coin selection "+
+			"strategy: %v", strategy)
+	}
+
+	// Having examined all the eligible commitments, return an error if the
+	// minimal funding amount was not reached.
+	if amountSum < minTotalAmount {
+		return nil, ErrMatchingAssetsNotFound
+	}
+	return selectedCommitments, nil
+}
+
+var _ CoinSelector = (*CoinSelect)(nil)