sweepbatcher: re-add sweeps after fully confirmed

In case of a reorg sweeps should not go to another batch but stay in the current
batch until it is fully confirmed. Only after that the remaining sweeps are
re-added to another batch.

Field sweep.completed is now set to true only for fully-confirmed sweeps.

Author: starius

sweepbatcher/store.go

@@ -248,7 +248,7 @@ type dbSweep struct {
 	// Amount is the amount of the sweep.
 	Amount btcutil.Amount
 
-	// Completed indicates whether this sweep is completed.
+	// Completed indicates whether this sweep is fully-confirmed.
 	Completed bool
 }
 

sweepbatcher/sweep_batch.go

@@ -1933,7 +1933,6 @@ func getFeePortionPaidBySweep(spendTx *wire.MsgTx, feePortionPerSweep,
 func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 	var (
 		txHash     = spendTx.TxHash()
-		purgeList  = make([]SweepRequest, 0, len(b.sweeps))
 		notifyList = make([]sweep, 0, len(b.sweeps))
 	)
 	b.batchTxid = &txHash
@@ -1943,7 +1942,100 @@ func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 		b.Warnf("transaction %v has no outputs", txHash)
 	}
 
-	// Determine if we should use presigned mode for the batch.
+	// Make a set of confirmed sweeps.
+	confirmedSet := make(map[wire.OutPoint]struct{}, len(spendTx.TxIn))
+	for _, txIn := range spendTx.TxIn {
+		confirmedSet[txIn.PreviousOutPoint] = struct{}{}
+	}
+
+	// As a previous version of the batch transaction may get confirmed,
+	// which does not contain the latest sweeps, we need to detect the
+	// sweeps that did not make it to the confirmed transaction and feed
+	// them back to the batcher. This will ensure that the sweeps will enter
+	// a new batch instead of remaining dangling.
+	var (
+		totalSweptAmt   btcutil.Amount
+		confirmedSweeps = []wire.OutPoint{}
+	)
+	for _, sweep := range b.sweeps {
+		// Skip sweeps that were not included into the confirmed tx.
+		_, found := confirmedSet[sweep.outpoint]
+		if !found {
+			continue
+		}
+
+		totalSweptAmt += sweep.value
+		notifyList = append(notifyList, sweep)
+		confirmedSweeps = append(confirmedSweeps, sweep.outpoint)
+	}
+
+	// Calculate the fee portion that each sweep should pay for the batch.
+	feePortionPaidPerSweep, roundingDifference := getFeePortionForSweep(
+		spendTx, len(notifyList), totalSweptAmt,
+	)
+
+	for _, sweep := range notifyList {
+		// If the sweep's notifier is empty then this means that a swap
+		// is not waiting to read an update from it, so we can skip
+		// the notification part.
+		if sweep.notifier == nil ||
+			*sweep.notifier == (SpendNotifier{}) {
+
+			continue
+		}
+
+		spendDetail := SpendDetail{
+			Tx: spendTx,
+			OnChainFeePortion: getFeePortionPaidBySweep(
+				spendTx, feePortionPaidPerSweep,
+				roundingDifference, &sweep,
+			),
+		}
+
+		// Dispatch the sweep notifier, we don't care about the outcome
+		// of this action so we don't wait for it.
+		go sweep.notifySweepSpend(ctx, &spendDetail)
+	}
+
+	b.Infof("spent, confirmed sweeps: %v", confirmedSweeps)
+
+	// We are no longer able to accept new sweeps, so we mark the batch as
+	// closed and persist on storage.
+	b.state = Closed
+
+	if err := b.persist(ctx); err != nil {
+		return fmt.Errorf("saving batch failed: %w", err)
+	}
+
+	if err := b.monitorConfirmations(ctx); err != nil {
+		return fmt.Errorf("monitorConfirmations failed: %w", err)
+	}
+
+	return nil
+}
+
+// handleConf handles a confirmation notification. This is the final step of the
+// batch. Here we signal to the batcher that this batch was completed.
+func (b *batch) handleConf(ctx context.Context,
+	conf *chainntnfs.TxConfirmation) error {
+
+	spendTx := conf.Tx
+	txHash := spendTx.TxHash()
+	if b.batchTxid == nil || *b.batchTxid != txHash {
+		b.Warnf("Mismatch of batch txid: tx in spend notification had "+
+			"txid %v, but confirmation notification has txif %v. "+
+			"Using the later.", b.batchTxid, txHash)
+	}
+	b.batchTxid = &txHash
+
+	b.Infof("confirmed in txid %s", b.batchTxid)
+	b.state = Confirmed
+
+	if err := b.persist(ctx); err != nil {
+		return fmt.Errorf("saving batch failed: %w", err)
+	}
+
+	// If the batch is in presigned mode, cleanup presignedHelper.
 	presigned, err := b.isPresigned()
 	if err != nil {
 		return fmt.Errorf("failed to determine if the batch %d uses "+
@@ -1965,38 +2057,43 @@ func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 		}
 	}
 
+	// Make a set of confirmed sweeps.
+	confirmedSet := make(map[wire.OutPoint]struct{}, len(spendTx.TxIn))
+	for _, txIn := range spendTx.TxIn {
+		confirmedSet[txIn.PreviousOutPoint] = struct{}{}
+	}
+
 	// As a previous version of the batch transaction may get confirmed,
 	// which does not contain the latest sweeps, we need to detect the
 	// sweeps that did not make it to the confirmed transaction and feed
 	// them back to the batcher. This will ensure that the sweeps will enter
 	// a new batch instead of remaining dangling.
 	var (
-		totalSweptAmt   btcutil.Amount
 		confirmedSweeps = []wire.OutPoint{}
-		purgedSweeps    = []wire.OutPoint{}
-		purgedSwaps     = []lntypes.Hash{}
+		purgeList       = make([]SweepRequest, 0, len(b.sweeps))
 	)
 	for _, sweep := range allSweeps {
-		found := false
-
-		for _, txIn := range spendTx.TxIn {
-			if txIn.PreviousOutPoint == sweep.outpoint {
-				found = true
-				totalSweptAmt += sweep.value
-				notifyList = append(notifyList, sweep)
-				confirmedSweeps = append(
-					confirmedSweeps, sweep.outpoint,
-				)
+		_, found := confirmedSet[sweep.outpoint]
+		if found {
+			// Save the sweep as completed. Note that sweeps are
+			// marked completed after the batch is marked confirmed
+			// because the check in handleSweeps checks sweep's
+			// status first and then checks the batch status.
+			err := b.persistSweep(ctx, sweep, true)
+			if err != nil {
+				return err
 			}
+
+			confirmedSweeps = append(
+				confirmedSweeps, sweep.outpoint,
+			)
+
+			continue
 		}
 
 		// If the sweep's outpoint was not found in the transaction's
 		// inputs this means it was left out. So we delete it from this
 		// batch and feed it back to the batcher.
-		if found {
-			continue
-		}
-
 		newSweep := sweep
 		delete(b.sweeps, sweep.outpoint)
 
@@ -2023,46 +2120,19 @@ func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 			})
 		}
 	}
+	var (
+		purgedSweeps = []wire.OutPoint{}
+		purgedSwaps  = []lntypes.Hash{}
+	)
 	for _, sweepReq := range purgeList {
 		purgedSwaps = append(purgedSwaps, sweepReq.SwapHash)
 		for _, input := range sweepReq.Inputs {
 			purgedSweeps = append(purgedSweeps, input.Outpoint)
 		}
 	}
 
-	// Calculate the fee portion that each sweep should pay for the batch.
-	feePortionPaidPerSweep, roundingDifference := getFeePortionForSweep(
-		spendTx, len(notifyList), totalSweptAmt,
-	)
-
-	for _, sweep := range notifyList {
-		// Save the sweep as completed.
-		err := b.persistSweep(ctx, sweep, true)
-		if err != nil {
-			return err
-		}
-
-		// If the sweep's notifier is empty then this means that a swap
-		// is not waiting to read an update from it, so we can skip
-		// the notification part.
-		if sweep.notifier == nil ||
-			*sweep.notifier == (SpendNotifier{}) {
-
-			continue
-		}
-
-		spendDetail := SpendDetail{
-			Tx: spendTx,
-			OnChainFeePortion: getFeePortionPaidBySweep(
-				spendTx, feePortionPaidPerSweep,
-				roundingDifference, &sweep,
-			),
-		}
-
-		// Dispatch the sweep notifier, we don't care about the outcome
-		// of this action so we don't wait for it.
-		go sweep.notifySweepSpend(ctx, &spendDetail)
-	}
+	b.Infof("fully confirmed sweeps: %v, purged sweeps: %v, "+
+		"purged swaps: %v", confirmedSweeps, purgedSweeps, purgedSwaps)
 
 	// Proceed with purging the sweeps. This will feed the sweeps that
 	// didn't make it to the confirmed batch transaction back to the batcher
@@ -2080,46 +2150,6 @@ func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 		}
 	}()
 
-	b.Infof("spent, confirmed sweeps: %v, purged sweeps: %v, "+
-		"purged swaps: %v", confirmedSweeps, purgedSweeps, purgedSwaps)
-
-	// We are no longer able to accept new sweeps, so we mark the batch as
-	// closed and persist on storage.
-	b.state = Closed
-
-	if err := b.persist(ctx); err != nil {
-		return fmt.Errorf("saving batch failed: %w", err)
-	}
-
-	err = b.monitorConfirmations(ctx)
-	if err != nil {
-		return fmt.Errorf("monitorConfirmations failed: %w", err)
-	}
-
-	return nil
-}
-
-// handleConf handles a confirmation notification. This is the final step of the
-// batch. Here we signal to the batcher that this batch was completed.
-func (b *batch) handleConf(ctx context.Context,
-	conf *chainntnfs.TxConfirmation) error {
-
-	spendTx := conf.Tx
-	txHash := spendTx.TxHash()
-	if b.batchTxid == nil || *b.batchTxid != txHash {
-		b.Warnf("Mismatch of batch txid: tx in spend notification had "+
-			"txid %v, but confirmation notification has txif %v. "+
-			"Using the later.", b.batchTxid, txHash)
-	}
-	b.batchTxid = &txHash
-
-	// If the batch is in presigned mode, cleanup presignedHelper.
-	presigned, err := b.isPresigned()
-	if err != nil {
-		return fmt.Errorf("failed to determine if the batch %d uses "+
-			"presigned mode: %w", b.id, err)
-	}
-
 	if presigned {
 		b.Infof("Cleaning up presigned store")
 
@@ -2135,13 +2165,6 @@ func (b *batch) handleConf(ctx context.Context,
 		}
 	}
 
-	b.Infof("confirmed in txid %s", b.batchTxid)
-	b.state = Confirmed
-
-	if err := b.store.ConfirmBatch(ctx, b.id); err != nil {
-		return fmt.Errorf("failed to store confirmed state: %w", err)
-	}
-
 	// Send the confirmation to all the notifiers.
 	for _, s := range b.sweeps {
 		// If the sweep's notifier is empty then this means that

sweepbatcher/sweep_batcher.go

@@ -756,13 +756,13 @@ func (b *Batcher) handleSweeps(ctx context.Context, sweeps []*sweep,
 		return err
 	}
 
-	infof("Batcher handling sweep %x, presigned=%v, completed=%v",
+	infof("Batcher handling sweep %x, presigned=%v, fully_confirmed=%v",
 		sweep.swapHash[:6], sweep.presigned, completed)
 
 	// If the sweep has already been completed in a confirmed batch then we
 	// can't attach its notifier to the batch as that is no longer running.
 	// Instead we directly detect and return the spend here.
-	if completed && *notifier != (SpendNotifier{}) {
+	if completed {
 		// Verify that the parent batch is confirmed. Note that a batch
 		// is only considered confirmed after it has received three
 		// on-chain confirmations to prevent issues caused by reorgs.
@@ -1064,6 +1064,11 @@ func (b *Batcher) FetchUnconfirmedBatches(ctx context.Context) ([]*batch,
 func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweep *sweep,
 	parentBatchID int32, notifier *SpendNotifier) error {
 
+	// If the caller has not provided a notifier, stop.
+	if notifier == nil || *notifier == (SpendNotifier{}) {
+		return nil
+	}
+
 	spendCtx, cancel := context.WithCancel(ctx)
 
 	// Then we get the total amount that was swept by the batch.

sweepbatcher/sweep_batcher_test.go

@@ -2302,22 +2302,15 @@ func testSweepBatcherSweepReentry(t *testing.T, store testStore,
 		return b.state == Closed
 	}, test.Timeout, eventuallyCheckFrequency)
 
-	// Since second batch was created we check that it registered for its
-	// primary sweep's spend.
-	<-lnd.RegisterSpendChannel
-
-	// While handling the spend notification the batch should detect that
-	// some sweeps did not appear in the spending tx, therefore it redirects
-	// them back to the batcher and the batcher inserts them in a new batch.
-	require.Eventually(t, func() bool {
-		return batcher.numBatches(ctx) == 2
-	}, test.Timeout, eventuallyCheckFrequency)
-
 	// We mock the confirmation notification.
 	lnd.ConfChannel <- &chainntnfs.TxConfirmation{
 		Tx: spendingTx,
 	}
 
+	// Since second batch was created we check that it registered for its
+	// primary sweep's spend.
+	<-lnd.RegisterSpendChannel
+
 	// Wait for tx to be published.
 	// Here is a race condition, which is unlikely to cause a crash: if we
 	// wait for publish tx before sending a conf notification (previous