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

@@ -1939,7 +1939,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
@@ -1949,7 +1948,106 @@ 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 func() {
+			// Make sure this context doesn't expire so we
+			// successfully notify the caller.
+			ctx := context.WithoutCancel(ctx)
+
+			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 "+
@@ -1967,40 +2065,43 @@ func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 			b.id, err)
 	}
 
+	// 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,
-				)
-
-				break
+		_, 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)
 
@@ -2032,52 +2133,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 func() {
-			// Make sure this context doesn't expire so we
-			// successfully notify the caller.
-			ctx := context.WithoutCancel(ctx)
-
-			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
@@ -2099,50 +2167,6 @@ func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 		}
 	}()
 
-	b.Infof("spent, confirmed sweeps: %v, purged sweeps: %v, "+
-		"purged swaps: %v, purged groups: %v", confirmedSweeps,
-		purgedSweeps, purgedSwaps, len(purgeList))
-
-	// 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. We also
-// cleanup up presigned transactions whose primarySweepID is one of the sweeps
-// that were spent and fully confirmed: such a transaction can't be broadcasted
-// since it is either in a block or double-spends one of spent outputs.
-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")
 
@@ -2158,13 +2182,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

@@ -273,7 +273,7 @@ type addSweepsRequest struct {
 	notifier *SpendNotifier
 
 	// completed is set if the sweep is spent and the spending transaction
-	// is confirmed.
+	// is fully confirmed.
 	completed bool
 
 	// parentBatch is the parent batch of this sweep. It is loaded ony if
@@ -777,8 +777,8 @@ func (b *Batcher) AddSweep(ctx context.Context, sweepReq *SweepRequest) error {
 	}
 
 	infof("Batcher adding sweep group of %d sweeps with primarySweep %x, "+
-		"presigned=%v, completed=%v", len(sweeps), sweep.swapHash[:6],
-		sweep.presigned, completed)
+		"presigned=%v, fully_confirmed=%v", len(sweeps),
+		sweep.swapHash[:6], sweep.presigned, completed)
 
 	req := &addSweepsRequest{
 		sweeps:      sweeps,
@@ -838,14 +838,10 @@ func (b *Batcher) handleSweeps(ctx context.Context, sweeps []*sweep,
 	// 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{}) {
-		// The parent batch is indeed confirmed, meaning it is complete
-		// and we won't be able to attach this sweep to it.
-		if parentBatch.Confirmed {
-			return b.monitorSpendAndNotify(
-				ctx, sweep, parentBatch.ID, notifier,
-			)
-		}
+	if completed && parentBatch.Confirmed {
+		return b.monitorSpendAndNotify(
+			ctx, sweep, parentBatch.ID, notifier,
+		)
 	}
 
 	sweep.notifier = notifier
@@ -1129,6 +1125,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.