sweepbatcher: customize transaction labels

Previously sweepbatcher used loop/labels.LoopOutBatchSweepSuccess to assign
transactions labels. The value is "BatchOutSweepSuccess -- $batch_id", which
does not fit use cases outside of loop-out. Now there is option WithTxLabeler
which sets a function used to generate the label.

Author: starius

sweepbatcher/sweep_batch.go

@@ -19,7 +19,6 @@ import (
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btclog"
 	"github.com/lightninglabs/lndclient"
-	"github.com/lightninglabs/loop/labels"
 	"github.com/lightninglabs/loop/loopdb"
 	"github.com/lightninglabs/loop/swap"
 	sweeppkg "github.com/lightninglabs/loop/sweep"
@@ -152,6 +151,10 @@ type batchConfig struct {
 	// external source of fee rates (FeeRateProvider).
 	noBumping bool
 
+	// txLabeler is a function generating a transaction label. It is called
+	// before publishing a batch transaction. Batch ID is passed to it.
+	txLabeler func(batchID int32) string
+
 	// customMuSig2Signer is a custom signer. If it is set, it is used to
 	// create musig2 signatures instead of musig2SignSweep and signerClient.
 	// Note that musig2SignSweep must be nil in this case, however signer
@@ -904,7 +907,7 @@ func (b *batch) publishBatch(ctx context.Context) (btcutil.Amount, error) {
 	b.debugLogTx("serialized non-coop sweep", batchTx)
 
 	err = b.wallet.PublishTransaction(
-		ctx, batchTx, labels.LoopOutBatchSweepSuccess(b.id),
+		ctx, batchTx, b.cfg.txLabeler(b.id),
 	)
 	if err != nil {
 		return fee, err
@@ -1053,7 +1056,7 @@ func (b *batch) publishBatchCoop(ctx context.Context) (btcutil.Amount,
 	b.debugLogTx("serialized coop sweep", batchTx)
 
 	err = b.wallet.PublishTransaction(
-		ctx, batchTx, labels.LoopOutBatchSweepSuccess(b.id),
+		ctx, batchTx, b.cfg.txLabeler(b.id),
 	)
 	if err != nil {
 		return fee, err, true

sweepbatcher/sweep_batcher.go

@@ -13,6 +13,7 @@ import (
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightninglabs/lndclient"
+	"github.com/lightninglabs/loop/labels"
 	"github.com/lightninglabs/loop/loopdb"
 	"github.com/lightninglabs/loop/swap"
 	"github.com/lightninglabs/loop/utils"
@@ -276,6 +277,10 @@ type Batcher struct {
 	// ignored and fee bumping by sweepbatcher is disabled.
 	customFeeRate FeeRateProvider
 
+	// txLabeler is a function generating a transaction label. It is called
+	// before publishing a batch transaction. Batch ID is passed to it.
+	txLabeler func(batchID int32) string
+
 	// customMuSig2Signer is a custom signer. If it is set, it is used to
 	// create musig2 signatures instead of musig2SignSweep and signerClient.
 	// Note that musig2SignSweep must be nil in this case, however signer
@@ -307,6 +312,10 @@ type BatcherConfig struct {
 	// ignored and fee bumping by sweepbatcher is disabled.
 	customFeeRate FeeRateProvider
 
+	// txLabeler is a function generating a transaction label. It is called
+	// before publishing a batch transaction. Batch ID is passed to it.
+	txLabeler func(batchID int32) string
+
 	// customMuSig2Signer is a custom signer. If it is set, it is used to
 	// create musig2 signatures instead of musig2SignSweep and signerClient.
 	// Note that musig2SignSweep must be nil in this case, however signer
@@ -357,6 +366,15 @@ func WithCustomFeeRate(customFeeRate FeeRateProvider) BatcherOption {
 	}
 }
 
+// WithTxLabeler sets a function generating a transaction label. It is called
+// before publishing a batch transaction. Batch ID is passed to the function.
+// By default, loop/labels.LoopOutBatchSweepSuccess is used.
+func WithTxLabeler(txLabeler func(batchID int32) string) BatcherOption {
+	return func(cfg *BatcherConfig) {
+		cfg.txLabeler = txLabeler
+	}
+}
+
 // WithCustomSignMuSig2 instructs sweepbatcher to use a custom function to
 // produce MuSig2 signatures. If it is set, it is used to create
 // musig2 signatures instead of musig2SignSweep and signerClient. Note
@@ -376,7 +394,11 @@ func NewBatcher(wallet lndclient.WalletKitClient,
 	store BatcherStore, sweepStore SweepFetcher,
 	opts ...BatcherOption) *Batcher {
 
-	var cfg BatcherConfig
+	cfg := BatcherConfig{
+		// By default, loop/labels.LoopOutBatchSweepSuccess is used
+		// to label sweep transactions.
+		txLabeler: labels.LoopOutBatchSweepSuccess,
+	}
 	for _, opt := range opts {
 		opt(&cfg)
 	}
@@ -409,6 +431,7 @@ func NewBatcher(wallet lndclient.WalletKitClient,
 		initialDelay:       cfg.initialDelay,
 		publishDelay:       cfg.publishDelay,
 		customFeeRate:      cfg.customFeeRate,
+		txLabeler:          cfg.txLabeler,
 		customMuSig2Signer: cfg.customMuSig2Signer,
 	}
 }
@@ -1024,6 +1047,7 @@ func (b *Batcher) newBatchConfig(maxTimeoutDistance int32) batchConfig {
 	return batchConfig{
 		maxTimeoutDistance: maxTimeoutDistance,
 		noBumping:          b.customFeeRate != nil,
+		txLabeler:          b.txLabeler,
 		customMuSig2Signer: b.customMuSig2Signer,
 		clock:              b.clock,
 	}

sweepbatcher/sweep_batcher_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btclog"
+	"github.com/lightninglabs/lndclient"
 	"github.com/lightninglabs/loop/loopdb"
 	"github.com/lightninglabs/loop/test"
 	"github.com/lightninglabs/loop/utils"
@@ -401,6 +402,144 @@ func testFeeBumping(t *testing.T, store testStore,
 	}
 }
 
+// walletKitWrapper wraps a wallet kit and memorizes the label of the most
+// recent published transaction.
+type walletKitWrapper struct {
+	lndclient.WalletKitClient
+
+	lastLabel string
+}
+
+// PublishTransaction publishes the transaction and memorizes its label.
+func (w *walletKitWrapper) PublishTransaction(ctx context.Context,
+	tx *wire.MsgTx, label string) error {
+
+	w.lastLabel = label
+
+	return w.WalletKitClient.PublishTransaction(ctx, tx, label)
+}
+
+// testTxLabeler tests transaction labels.
+func testTxLabeler(t *testing.T, store testStore,
+	batcherStore testBatcherStore) {
+
+	defer test.Guard(t)()
+
+	lnd := test.NewMockLnd()
+	ctx, cancel := context.WithCancel(context.Background())
+
+	sweepStore, err := NewSweepFetcherFromSwapStore(store, lnd.ChainParams)
+	require.NoError(t, err)
+
+	walletKit := &walletKitWrapper{WalletKitClient: lnd.WalletKit}
+
+	batcher := NewBatcher(walletKit, lnd.ChainNotifier, lnd.Signer,
+		testMuSig2SignSweep, testVerifySchnorrSig, lnd.ChainParams,
+		batcherStore, sweepStore)
+
+	var (
+		runErr error
+		wg     sync.WaitGroup
+	)
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		runErr = batcher.Run(ctx)
+	}()
+
+	// Create a sweep request.
+	sweepReq1 := SweepRequest{
+		SwapHash: lntypes.Hash{1, 1, 1},
+		Value:    111,
+		Outpoint: wire.OutPoint{
+			Hash:  chainhash.Hash{1, 1},
+			Index: 1,
+		},
+		Notifier: &dummyNotifier,
+	}
+
+	swap1 := &loopdb.LoopOutContract{
+		SwapContract: loopdb.SwapContract{
+			CltvExpiry:      111,
+			AmountRequested: 111,
+			ProtocolVersion: loopdb.ProtocolVersionMuSig2,
+			HtlcKeys:        htlcKeys,
+		},
+
+		DestAddr:        destAddr,
+		SwapInvoice:     swapInvoice,
+		SweepConfTarget: 111,
+	}
+
+	err = store.CreateLoopOut(ctx, sweepReq1.SwapHash, swap1)
+	require.NoError(t, err)
+	store.AssertLoopOutStored()
+
+	// Deliver sweep request to batcher.
+	require.NoError(t, batcher.AddSweep(&sweepReq1))
+
+	// Eventually request will be consumed and a new batch will spin up.
+	require.Eventually(t, func() bool {
+		return len(batcher.batches) == 1
+	}, test.Timeout, eventuallyCheckFrequency)
+
+	// When batch is successfully created it will execute it's first step,
+	// which leads to a spend monitor of the primary sweep.
+	<-lnd.RegisterSpendChannel
+
+	// Wait for tx to be published.
+	<-lnd.TxPublishChannel
+
+	// Find the batch and assign it to a local variable for easier access.
+	var theBatch *batch
+	for _, btch := range batcher.batches {
+		if btch.primarySweepID == sweepReq1.SwapHash {
+			theBatch = btch
+		}
+	}
+
+	// Now test the label.
+	wantLabel := fmt.Sprintf("BatchOutSweepSuccess -- %d", theBatch.id)
+	require.Equal(t, wantLabel, walletKit.lastLabel)
+
+	// Now make the batcher quit by canceling the context.
+	cancel()
+	wg.Wait()
+	checkBatcherError(t, runErr)
+
+	// Define dummy tx labeler, always returning "test".
+	txLabeler := func(batchID int32) string {
+		return "test"
+	}
+
+	// Now try it with option WithTxLabeler.
+	batcher = NewBatcher(walletKit, lnd.ChainNotifier, lnd.Signer,
+		testMuSig2SignSweep, testVerifySchnorrSig, lnd.ChainParams,
+		batcherStore, sweepStore, WithTxLabeler(txLabeler))
+
+	ctx, cancel = context.WithCancel(context.Background())
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		runErr = batcher.Run(ctx)
+	}()
+
+	// Expect batch to register for spending.
+	<-lnd.RegisterSpendChannel
+
+	// Wait for tx to be published.
+	<-lnd.TxPublishChannel
+
+	// Now test the label.
+	require.Equal(t, "test", walletKit.lastLabel)
+
+	// Now make the batcher quit by canceling the context.
+	cancel()
+	wg.Wait()
+	checkBatcherError(t, runErr)
+}
+
 // testSweepBatcherSimpleLifecycle tests the simple lifecycle of the batches
 // that are created and run by the batcher.
 func testSweepBatcherSimpleLifecycle(t *testing.T, store testStore,
@@ -2767,6 +2906,11 @@ func TestFeeBumping(t *testing.T) {
 	})
 }
 
+// TestTxLabeler tests transaction labels.
+func TestTxLabeler(t *testing.T) {
+	runTests(t, testTxLabeler)
+}
+
 // TestSweepBatcherSimpleLifecycle tests the simple lifecycle of the batches
 // that are created and run by the batcher.
 func TestSweepBatcherSimpleLifecycle(t *testing.T) {