contractcourt: add Launch method to htlc timeout resolver

This commit breaks the `Resolve` into two parts - the first part is
moved into a `Launch` method that handles sending sweep requests, and
the second part remains in `Resolve` which handles waiting for the
spend. Since we are using both utxo nursery and sweeper at the same
time, to make sure this change doesn't break the existing behavior, we
implement the `Launch` as following,
- zero-fee htlc - handled by the sweeper
- direct output from the remote commit - handled by the sweeper
- legacy htlc - handled by the utxo nursery

contractcourt/channel_arbitrator_test.go

@@ -982,6 +982,7 @@ func TestChannelArbitratorLocalForceClosePendingHtlc(t *testing.T) {
 			},
 		},
 	}
+	closeTxid := closeTx.TxHash()
 
 	htlcOp := wire.OutPoint{
 		Hash:  closeTx.TxHash(),
@@ -1117,7 +1118,11 @@ func TestChannelArbitratorLocalForceClosePendingHtlc(t *testing.T) {
 
 	// Notify resolver that the HTLC output of the commitment has been
 	// spent.
-	oldNotifier.SpendChan <- &chainntnfs.SpendDetail{SpendingTx: closeTx}
+	oldNotifier.SpendChan <- &chainntnfs.SpendDetail{
+		SpendingTx:    closeTx,
+		SpentOutPoint: &wire.OutPoint{},
+		SpenderTxHash: &closeTxid,
+	}
 
 	// Finally, we should also receive a resolution message instructing the
 	// switch to cancel back the HTLC.

contractcourt/htlc_timeout_resolver.go

@@ -425,40 +425,25 @@ func checkSizeAndIndex(witness wire.TxWitness, size, index int) bool {
 func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 	// If we're already resolved, then we can exit early.
 	if h.resolved {
+		h.log.Errorf("already resolved")
 		return nil, nil
 	}
 
-	// Start by spending the HTLC output, either by broadcasting the
-	// second-level timeout transaction, or directly if this is the remote
-	// commitment.
-	commitSpend, err := h.spendHtlcOutput()
-	if err != nil {
-		return nil, err
-	}
-
-	// If the spend reveals the pre-image, then we'll enter the clean up
-	// workflow to pass the pre-image back to the incoming link, add it to
-	// the witness cache, and exit.
-	if isPreimageSpend(
-		h.isTaproot(), commitSpend,
-		h.htlcResolution.SignedTimeoutTx != nil,
-	) {
-
-		log.Infof("%T(%v): HTLC has been swept with pre-image by "+
-			"remote party during timeout flow! Adding pre-image to "+
-			"witness cache", h, h.htlc.RHash[:],
-			h.htlcResolution.ClaimOutpoint)
-
-		return nil, h.claimCleanUp(commitSpend)
-	}
-
-	// Depending on whether this was a local or remote commit, we must
-	// handle the spending transaction accordingly.
+	// If this is an output on the remote party's commitment transaction,
+	// use the direct-spend path to sweep the htlc.
 	if h.isRemoteCommitOutput() {
 		return nil, h.resolveRemoteCommitOutput()
 	}
 
-	return nil, h.resolveTimeoutTx()
+	// If this is a zero-fee HTLC, we now handle the spend from our
+	// commitment transaction.
+	if h.isZeroFeeOutput() {
+		return nil, h.resolveTimeoutTx()
+	}
+
+	// If this is an output on our own commitment using pre-anchor channel
+	// type, we will let the utxo nursery handle it.
+	return nil, h.resolveSecondLevelTxLegacy()
 }
 
 // sweepTimeoutTx sends a second level timeout transaction to the sweeper.
@@ -521,11 +506,16 @@ func (h *htlcTimeoutResolver) resolveSecondLevelTxLegacy() error {
 
 	// The utxo nursery will take care of broadcasting the second-level
 	// timeout tx and sweeping its output once it confirms.
-	return h.IncubateOutputs(
+	err := h.IncubateOutputs(
 		h.ChanPoint, fn.Some(h.htlcResolution),
 		fn.None[lnwallet.IncomingHtlcResolution](),
 		h.broadcastHeight, h.incomingHTLCExpiryHeight,
 	)
+	if err != nil {
+		return err
+	}
+
+	return h.resolveTimeoutTx()
 }
 
 // sweepDirectHtlcOutput sends the direct spend of the HTLC output to the
@@ -581,53 +571,6 @@ func (h *htlcTimeoutResolver) sweepDirectHtlcOutput() error {
 	return nil
 }
 
-// spendHtlcOutput handles the initial spend of an HTLC output via the timeout
-// clause. If this is our local commitment, the second-level timeout TX will be
-// used to spend the output into the next stage. If this is the remote
-// commitment, the output will be swept directly without the timeout
-// transaction.
-func (h *htlcTimeoutResolver) spendHtlcOutput() (
-	*chainntnfs.SpendDetail, error) {
-
-	switch {
-	// If we have non-nil SignDetails, this means that have a 2nd level
-	// HTLC transaction that is signed using sighash SINGLE|ANYONECANPAY
-	// (the case for anchor type channels). In this case we can re-sign it
-	// and attach fees at will. We let the sweeper handle this job.
-	case h.isZeroFeeOutput() && !h.outputIncubating:
-		if err := h.sweepTimeoutTx(); err != nil {
-			log.Errorf("Sending timeout tx to sweeper: %v", err)
-
-			return nil, err
-		}
-
-	// If this is a remote commitment there's no second level timeout txn,
-	// and we can just send this directly to the sweeper.
-	case h.isRemoteCommitOutput() && !h.outputIncubating:
-		if err := h.sweepDirectHtlcOutput(); err != nil {
-			log.Errorf("Sending direct spend to sweeper: %v", err)
-
-			return nil, err
-		}
-
-	// If we have a SignedTimeoutTx but no SignDetails, this is a local
-	// commitment for a non-anchor channel, so we'll send it to the utxo
-	// nursery.
-	case h.isLegacyOutput() && !h.outputIncubating:
-		if err := h.resolveSecondLevelTxLegacy(); err != nil {
-			log.Errorf("Sending timeout tx to nursery: %v", err)
-
-			return nil, err
-		}
-	}
-
-	// Now that we've handed off the HTLC to the nursery or sweeper, we'll
-	// watch for a spend of the output, and make our next move off of that.
-	// Depending on if this is our commitment, or the remote party's
-	// commitment, we'll be watching a different outpoint and script.
-	return h.watchHtlcSpend()
-}
-
 // watchHtlcSpend watches for a spend of the HTLC output. For neutrino backend,
 // it will check blocks for the confirmed spend. For btcd and bitcoind, it will
 // check both the mempool and the blocks.
@@ -673,6 +616,9 @@ func (h *htlcTimeoutResolver) waitForConfirmedSpend(op *wire.OutPoint,
 //
 // NOTE: Part of the ContractResolver interface.
 func (h *htlcTimeoutResolver) Stop() {
+	h.log.Debugf("stopping...")
+	defer h.log.Debugf("stopped")
+
 	close(h.quit)
 }
 
@@ -1018,15 +964,6 @@ func (h *htlcTimeoutResolver) isZeroFeeOutput() bool {
 		h.htlcResolution.SignDetails != nil
 }
 
-// isLegacyOutput returns a boolean indicating whether the htlc output is from
-// a non-anchor-enabled channel.
-func (h *htlcTimeoutResolver) isLegacyOutput() bool {
-	// If we have a SignedTimeoutTx but no SignDetails, this is a local
-	// commitment for a non-anchor channel.
-	return h.htlcResolution.SignedTimeoutTx != nil &&
-		h.htlcResolution.SignDetails == nil
-}
-
 // waitHtlcSpendAndCheckPreimage waits for the htlc output to be spent and
 // checks whether the spending reveals the preimage. If the preimage is found,
 // it will be added to the preimage beacon to settle the incoming link, and a
@@ -1296,9 +1233,11 @@ func (h *htlcTimeoutResolver) resolveTimeoutTx() error {
 	h.log.Infof("2nd-level HTLC timeout tx=%v confirmed", spenderTxid)
 
 	// Start the process to sweep the output from the timeout tx.
-	err = h.sweepTimeoutTxOutput()
-	if err != nil {
-		return err
+	if h.isZeroFeeOutput() {
+		err = h.sweepTimeoutTxOutput()
+		if err != nil {
+			return err
+		}
 	}
 
 	// Create a checkpoint since the timeout tx is confirmed and the sweep
@@ -1332,3 +1271,52 @@ func (h *htlcTimeoutResolver) resolveTimeoutTxOutput(op wire.OutPoint) error {
 
 	return h.checkpointClaim(spend)
 }
+
+// Launch creates an input based on the details of the outgoing htlc resolution
+// and offers it to the sweeper.
+func (h *htlcTimeoutResolver) Launch() error {
+	if h.launched {
+		h.log.Tracef("already launched")
+		return nil
+	}
+
+	h.log.Debugf("launching resolver...")
+	h.launched = true
+
+	switch {
+	// If we're already resolved, then we can exit early.
+	case h.resolved:
+		h.log.Errorf("already resolved")
+		return nil
+
+	// If this is an output on the remote party's commitment transaction,
+	// use the direct timeout spend path.
+	//
+	// NOTE: When the outputIncubating is false, it means that the output
+	// has been offered to the utxo nursery as starting in 0.18.4, we
+	// stopped marking this flag for direct timeout spends (#9062). In that
+	// case, we will do nothing and let the utxo nursery handle it.
+	case h.isRemoteCommitOutput() && !h.outputIncubating:
+		return h.sweepDirectHtlcOutput()
+
+	// If this is an anchor type channel, we now sweep either the
+	// second-level timeout tx or the output from the second-level timeout
+	// tx.
+	case h.isZeroFeeOutput():
+		// If the second-level timeout tx has already been swept, we
+		// can go ahead and sweep its output.
+		if h.outputIncubating {
+			return h.sweepTimeoutTxOutput()
+		}
+
+		// Otherwise, sweep the second level tx.
+		return h.sweepTimeoutTx()
+
+	// If this is an output on our own commitment using pre-anchor channel
+	// type, we will let the utxo nursery handle it via Resolve.
+	//
+	// TODO(yy): handle the legacy output by offering it to the sweeper.
+	default:
+		return nil
+	}
+}

contractcourt/htlc_timeout_resolver_test.go

@@ -40,7 +40,7 @@ type mockWitnessBeacon struct {
 func newMockWitnessBeacon() *mockWitnessBeacon {
 	return &mockWitnessBeacon{
 		preImageUpdates: make(chan lntypes.Preimage, 1),
-		newPreimages:    make(chan []lntypes.Preimage),
+		newPreimages:    make(chan []lntypes.Preimage, 1),
 		lookupPreimage:  make(map[lntypes.Hash]lntypes.Preimage),
 	}
 }
@@ -280,7 +280,7 @@ func testHtlcTimeoutResolver(t *testing.T, testCase htlcTimeoutTestCase) {
 
 	notifier := &mock.ChainNotifier{
 		EpochChan: make(chan *chainntnfs.BlockEpoch),
-		SpendChan: make(chan *chainntnfs.SpendDetail),
+		SpendChan: make(chan *chainntnfs.SpendDetail, 1),
 		ConfChan:  make(chan *chainntnfs.TxConfirmation),
 	}
 
@@ -321,6 +321,7 @@ func testHtlcTimeoutResolver(t *testing.T, testCase htlcTimeoutTestCase) {
 
 				return nil
 			},
+			HtlcNotifier: &mockHTLCNotifier{},
 		},
 		PutResolverReport: func(_ kvdb.RwTx,
 			_ *channeldb.ResolverReport) error {
@@ -356,6 +357,7 @@ func testHtlcTimeoutResolver(t *testing.T, testCase htlcTimeoutTestCase) {
 			Amt: testHtlcAmt,
 		},
 	}
+	resolver.initLogger("timeoutResolver")
 
 	var reports []*channeldb.ResolverReport
 
@@ -390,7 +392,12 @@ func testHtlcTimeoutResolver(t *testing.T, testCase htlcTimeoutTestCase) {
 	go func() {
 		defer wg.Done()
 
-		_, err := resolver.Resolve()
+		err := resolver.Launch()
+		if err != nil {
+			resolveErr <- err
+		}
+
+		_, err = resolver.Resolve()
 		if err != nil {
 			resolveErr <- err
 		}
@@ -406,8 +413,7 @@ func testHtlcTimeoutResolver(t *testing.T, testCase htlcTimeoutTestCase) {
 		sweepChan = mockSweeper.sweptInputs
 	}
 
-	// The output should be offered to either the sweeper or
-	// the nursery.
+	// The output should be offered to either the sweeper or the nursery.
 	select {
 	case <-incubateChan:
 	case <-sweepChan:
@@ -431,6 +437,7 @@ func testHtlcTimeoutResolver(t *testing.T, testCase htlcTimeoutTestCase) {
 	case notifier.SpendChan <- &chainntnfs.SpendDetail{
 		SpendingTx:    spendingTx,
 		SpenderTxHash: &spendTxHash,
+		SpentOutPoint: &testChanPoint2,
 	}:
 	case <-time.After(time.Second * 5):
 		t.Fatalf("failed to request spend ntfn")
@@ -487,6 +494,7 @@ func testHtlcTimeoutResolver(t *testing.T, testCase htlcTimeoutTestCase) {
 			case notifier.SpendChan <- &chainntnfs.SpendDetail{
 				SpendingTx:    spendingTx,
 				SpenderTxHash: &spendTxHash,
+				SpentOutPoint: &testChanPoint2,
 			}:
 			case <-time.After(time.Second * 5):
 				t.Fatalf("failed to request spend ntfn")
@@ -549,6 +557,8 @@ func TestHtlcTimeoutResolver(t *testing.T) {
 
 // TestHtlcTimeoutSingleStage tests a remote commitment confirming, and the
 // local node sweeping the HTLC output directly after timeout.
+//
+//nolint:ll
 func TestHtlcTimeoutSingleStage(t *testing.T) {
 	commitOutpoint := wire.OutPoint{Index: 3}
 
@@ -573,6 +583,12 @@ func TestHtlcTimeoutSingleStage(t *testing.T) {
 		SpendTxID:       &sweepTxid,
 	}
 
+	sweepSpend := &chainntnfs.SpendDetail{
+		SpendingTx:    sweepTx,
+		SpentOutPoint: &commitOutpoint,
+		SpenderTxHash: &sweepTxid,
+	}
+
 	checkpoints := []checkpoint{
 		{
 			// We send a confirmation the sweep tx from published
@@ -582,9 +598,10 @@ func TestHtlcTimeoutSingleStage(t *testing.T) {
 
 				// The nursery will create and publish a sweep
 				// tx.
-				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
-					SpendingTx:    sweepTx,
-					SpenderTxHash: &sweepTxid,
+				select {
+				case ctx.notifier.SpendChan <- sweepSpend:
+				case <-time.After(time.Second * 5):
+					t.Fatalf("failed to send spend ntfn")
 				}
 
 				// The resolver should deliver a failure
@@ -620,7 +637,9 @@ func TestHtlcTimeoutSingleStage(t *testing.T) {
 
 // TestHtlcTimeoutSecondStage tests a local commitment being confirmed, and the
 // local node claiming the HTLC output using the second-level timeout tx.
-func TestHtlcTimeoutSecondStage(t *testing.T) {
+//
+//nolint:ll
+func TestHtlcTimeoutSecondStagex(t *testing.T) {
 	commitOutpoint := wire.OutPoint{Index: 2}
 	htlcOutpoint := wire.OutPoint{Index: 3}
 
@@ -678,23 +697,57 @@ func TestHtlcTimeoutSecondStage(t *testing.T) {
 		SpendTxID:       &sweepHash,
 	}
 
+	timeoutSpend := &chainntnfs.SpendDetail{
+		SpendingTx:    timeoutTx,
+		SpentOutPoint: &commitOutpoint,
+		SpenderTxHash: &timeoutTxid,
+	}
+
+	sweepSpend := &chainntnfs.SpendDetail{
+		SpendingTx:    sweepTx,
+		SpentOutPoint: &htlcOutpoint,
+		SpenderTxHash: &sweepHash,
+	}
+
 	checkpoints := []checkpoint{
 		{
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				_ bool) error {
+
+				// Deliver spend of timeout tx.
+				ctx.notifier.SpendChan <- timeoutSpend
+
+				return nil
+			},
+
 			// Output should be handed off to the nursery.
 			incubating: true,
+			reports: []*channeldb.ResolverReport{
+				firstStage,
+			},
 		},
 		{
 			// We send a confirmation for our sweep tx to indicate
 			// that our sweep succeeded.
 			preCheckpoint: func(ctx *htlcResolverTestContext,
-				_ bool) error {
+				resumed bool) error {
 
-				// The nursery will publish the timeout tx.
-				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
-					SpendingTx:    timeoutTx,
-					SpenderTxHash: &timeoutTxid,
+				// When it's reloaded from disk, we need to
+				// re-send the notification to mock the first
+				// `watchHtlcSpend`.
+				if resumed {
+					// Deliver spend of timeout tx.
+					ctx.notifier.SpendChan <- timeoutSpend
+
+					// Deliver spend of timeout tx output.
+					ctx.notifier.SpendChan <- sweepSpend
+
+					return nil
 				}
 
+				// Deliver spend of timeout tx output.
+				ctx.notifier.SpendChan <- sweepSpend
+
 				// The resolver should deliver a failure
 				// resolution message (indicating we
 				// successfully timed out the HTLC).
@@ -707,12 +760,6 @@ func TestHtlcTimeoutSecondStage(t *testing.T) {
 					t.Fatalf("resolution not sent")
 				}
 
-				// Deliver spend of timeout tx.
-				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
-					SpendingTx:    sweepTx,
-					SpenderTxHash: &sweepHash,
-				}
-
 				return nil
 			},
 
@@ -722,7 +769,7 @@ func TestHtlcTimeoutSecondStage(t *testing.T) {
 			incubating: true,
 			resolved:   true,
 			reports: []*channeldb.ResolverReport{
-				firstStage, secondState,
+				secondState,
 			},
 		},
 	}
@@ -796,10 +843,6 @@ func TestHtlcTimeoutSingleStageRemoteSpend(t *testing.T) {
 	}
 
 	checkpoints := []checkpoint{
-		{
-			// Output should be handed off to the nursery.
-			incubating: true,
-		},
 		{
 			// We send a spend notification for a remote spend with
 			// the preimage.
@@ -812,6 +855,7 @@ func TestHtlcTimeoutSingleStageRemoteSpend(t *testing.T) {
 				// the preimage.
 				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
 					SpendingTx:    spendTx,
+					SpentOutPoint: &commitOutpoint,
 					SpenderTxHash: &spendTxHash,
 				}
 
@@ -847,7 +891,7 @@ func TestHtlcTimeoutSingleStageRemoteSpend(t *testing.T) {
 			// After the success tx has confirmed, we expect the
 			// checkpoint to be resolved, and with the above
 			// report.
-			incubating: true,
+			incubating: false,
 			resolved:   true,
 			reports: []*channeldb.ResolverReport{
 				claim,
@@ -914,6 +958,7 @@ func TestHtlcTimeoutSecondStageRemoteSpend(t *testing.T) {
 
 				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
 					SpendingTx:    remoteSuccessTx,
+					SpentOutPoint: &commitOutpoint,
 					SpenderTxHash: &successTxid,
 				}
 
@@ -967,20 +1012,15 @@ func TestHtlcTimeoutSecondStageRemoteSpend(t *testing.T) {
 // TestHtlcTimeoutSecondStageSweeper tests that for anchor channels, when a
 // local commitment confirms, the timeout tx is handed to the sweeper to claim
 // the HTLC output.
+//
+//nolint:ll
 func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
-	commitOutpoint := wire.OutPoint{Index: 2}
 	htlcOutpoint := wire.OutPoint{Index: 3}
 
-	sweepTx := &wire.MsgTx{
-		TxIn:  []*wire.TxIn{{}},
-		TxOut: []*wire.TxOut{{}},
-	}
-	sweepHash := sweepTx.TxHash()
-
 	timeoutTx := &wire.MsgTx{
 		TxIn: []*wire.TxIn{
 			{
-				PreviousOutPoint: commitOutpoint,
+				PreviousOutPoint: htlcOutpoint,
 			},
 		},
 		TxOut: []*wire.TxOut{
@@ -1027,11 +1067,16 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 		},
 	}
 	reSignedHash := reSignedTimeoutTx.TxHash()
-	reSignedOutPoint := wire.OutPoint{
+
+	timeoutTxOutpoint := wire.OutPoint{
 		Hash:  reSignedHash,
 		Index: 1,
 	}
 
+	// Make a copy so `isPreimageSpend` can easily pass.
+	sweepTx := reSignedTimeoutTx.Copy()
+	sweepHash := sweepTx.TxHash()
+
 	// twoStageResolution is a resolution for a htlc on the local
 	// party's commitment, where the timeout tx can be re-signed.
 	twoStageResolution := lnwallet.OutgoingHtlcResolution{
@@ -1045,7 +1090,7 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 	}
 
 	firstStage := &channeldb.ResolverReport{
-		OutPoint:        commitOutpoint,
+		OutPoint:        htlcOutpoint,
 		Amount:          testHtlcAmt.ToSatoshis(),
 		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
 		ResolverOutcome: channeldb.ResolverOutcomeFirstStage,
@@ -1053,12 +1098,45 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 	}
 
 	secondState := &channeldb.ResolverReport{
-		OutPoint:        reSignedOutPoint,
+		OutPoint:        timeoutTxOutpoint,
 		Amount:          btcutil.Amount(testSignDesc.Output.Value),
 		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
 		ResolverOutcome: channeldb.ResolverOutcomeTimeout,
 		SpendTxID:       &sweepHash,
 	}
+	// mockTimeoutTxSpend is a helper closure to mock `waitForSpend` to
+	// return the commit spend in `sweepTimeoutTxOutput`.
+	mockTimeoutTxSpend := func(ctx *htlcResolverTestContext) {
+		select {
+		case ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+			SpendingTx:        reSignedTimeoutTx,
+			SpenderInputIndex: 1,
+			SpenderTxHash:     &reSignedHash,
+			SpendingHeight:    10,
+			SpentOutPoint:     &htlcOutpoint,
+		}:
+
+		case <-time.After(time.Second * 1):
+			t.Fatalf("spend not sent")
+		}
+	}
+
+	// mockSweepTxSpend is a helper closure to mock `waitForSpend` to
+	// return the commit spend in `sweepTimeoutTxOutput`.
+	mockSweepTxSpend := func(ctx *htlcResolverTestContext) {
+		select {
+		case ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+			SpendingTx:        sweepTx,
+			SpenderInputIndex: 1,
+			SpenderTxHash:     &sweepHash,
+			SpendingHeight:    10,
+			SpentOutPoint:     &timeoutTxOutpoint,
+		}:
+
+		case <-time.After(time.Second * 1):
+			t.Fatalf("spend not sent")
+		}
+	}
 
 	checkpoints := []checkpoint{
 		{
@@ -1067,28 +1145,40 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 				_ bool) error {
 
 				resolver := ctx.resolver.(*htlcTimeoutResolver)
-				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
-				op := inp.OutPoint()
-				if op != commitOutpoint {
-					return fmt.Errorf("outpoint %v swept, "+
-						"expected %v", op,
-						commitOutpoint)
+
+				var (
+					inp input.Input
+					ok  bool
+				)
+
+				select {
+				case inp, ok = <-resolver.Sweeper.(*mockSweeper).sweptInputs:
+					require.True(t, ok)
+
+				case <-time.After(1 * time.Second):
+					t.Fatal("expected input to be swept")
 				}
 
-				// Emulat the sweeper spending using the
-				// re-signed timeout tx.
-				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
-					SpendingTx:        reSignedTimeoutTx,
-					SpenderInputIndex: 1,
-					SpenderTxHash:     &reSignedHash,
-					SpendingHeight:    10,
+				op := inp.OutPoint()
+				if op != htlcOutpoint {
+					return fmt.Errorf("outpoint %v swept, "+
+						"expected %v", op, htlcOutpoint)
 				}
 
+				// Mock `waitForSpend` twice, called in,
+				// - `resolveReSignedTimeoutTx`
+				// - `sweepTimeoutTxOutput`.
+				mockTimeoutTxSpend(ctx)
+				mockTimeoutTxSpend(ctx)
+
 				return nil
 			},
 			// incubating=true is used to signal that the
 			// second-level transaction was confirmed.
 			incubating: true,
+			reports: []*channeldb.ResolverReport{
+				firstStage,
+			},
 		},
 		{
 			// We send a confirmation for our sweep tx to indicate
@@ -1096,18 +1186,18 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 			preCheckpoint: func(ctx *htlcResolverTestContext,
 				resumed bool) error {
 
-				// If we are resuming from a checkpoint, we
-				// expect the resolver to re-subscribe to a
-				// spend, hence we must resend it.
+				// Mock `waitForSpend` to return the commit
+				// spend.
 				if resumed {
-					ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
-						SpendingTx:        reSignedTimeoutTx,
-						SpenderInputIndex: 1,
-						SpenderTxHash:     &reSignedHash,
-						SpendingHeight:    10,
-					}
+					mockTimeoutTxSpend(ctx)
+					mockTimeoutTxSpend(ctx)
+					mockSweepTxSpend(ctx)
+
+					return nil
 				}
 
+				mockSweepTxSpend(ctx)
+
 				// The resolver should deliver a failure
 				// resolution message (indicating we
 				// successfully timed out the HTLC).
@@ -1123,7 +1213,20 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 				// The timeout tx output should now be given to
 				// the sweeper.
 				resolver := ctx.resolver.(*htlcTimeoutResolver)
-				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
+
+				var (
+					inp input.Input
+					ok  bool
+				)
+
+				select {
+				case inp, ok = <-resolver.Sweeper.(*mockSweeper).sweptInputs:
+					require.True(t, ok)
+
+				case <-time.After(1 * time.Second):
+					t.Fatal("expected input to be swept")
+				}
+
 				op := inp.OutPoint()
 				exp := wire.OutPoint{
 					Hash:  reSignedHash,
@@ -1133,14 +1236,6 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 					return fmt.Errorf("wrong outpoint swept")
 				}
 
-				// Notify about the spend, which should resolve
-				// the resolver.
-				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
-					SpendingTx:     sweepTx,
-					SpenderTxHash:  &sweepHash,
-					SpendingHeight: 14,
-				}
-
 				return nil
 			},
 
@@ -1150,7 +1245,6 @@ func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
 			incubating: true,
 			resolved:   true,
 			reports: []*channeldb.ResolverReport{
-				firstStage,
 				secondState,
 			},
 		},
@@ -1231,33 +1325,6 @@ func TestHtlcTimeoutSecondStageSweeperRemoteSpend(t *testing.T) {
 	}
 
 	checkpoints := []checkpoint{
-		{
-			// The output should be given to the sweeper.
-			preCheckpoint: func(ctx *htlcResolverTestContext,
-				_ bool) error {
-
-				resolver := ctx.resolver.(*htlcTimeoutResolver)
-				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
-				op := inp.OutPoint()
-				if op != commitOutpoint {
-					return fmt.Errorf("outpoint %v swept, "+
-						"expected %v", op,
-						commitOutpoint)
-				}
-
-				// Emulate the remote sweeping the output with the preimage.
-				// re-signed timeout tx.
-				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
-					SpendingTx:    spendTx,
-					SpenderTxHash: &spendTxHash,
-				}
-
-				return nil
-			},
-			// incubating=true is used to signal that the
-			// second-level transaction was confirmed.
-			incubating: true,
-		},
 		{
 			// We send a confirmation for our sweep tx to indicate
 			// that our sweep succeeded.
@@ -1272,6 +1339,7 @@ func TestHtlcTimeoutSecondStageSweeperRemoteSpend(t *testing.T) {
 					ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
 						SpendingTx:    spendTx,
 						SpenderTxHash: &spendTxHash,
+						SpentOutPoint: &commitOutpoint,
 					}
 				}
 
@@ -1309,7 +1377,7 @@ func TestHtlcTimeoutSecondStageSweeperRemoteSpend(t *testing.T) {
 			// After the sweep has confirmed, we expect the
 			// checkpoint to be resolved, and with the above
 			// reports.
-			incubating: true,
+			incubating: false,
 			resolved:   true,
 			reports: []*channeldb.ResolverReport{
 				claim,
@@ -1334,21 +1402,26 @@ func testHtlcTimeout(t *testing.T, resolution lnwallet.OutgoingHtlcResolution,
 	// for the next portion of the test.
 	ctx := newHtlcResolverTestContext(t,
 		func(htlc channeldb.HTLC, cfg ResolverConfig) ContractResolver {
-			return &htlcTimeoutResolver{
+			r := &htlcTimeoutResolver{
 				contractResolverKit: *newContractResolverKit(cfg),
 				htlc:                htlc,
 				htlcResolution:      resolution,
 			}
+			r.initLogger("htlcTimeoutResolver")
+
+			return r
 		},
 	)
 
 	checkpointedState := runFromCheckpoint(t, ctx, checkpoints)
 
+	t.Log("Running resolver to completion after restart")
+
 	// Now, from every checkpoint created, we re-create the resolver, and
 	// run the test from that checkpoint.
 	for i := range checkpointedState {
 		cp := bytes.NewReader(checkpointedState[i])
-		ctx := newHtlcResolverTestContext(t,
+		ctx := newHtlcResolverTestContextFromReader(t,
 			func(htlc channeldb.HTLC, cfg ResolverConfig) ContractResolver {
 				resolver, err := newTimeoutResolverFromReader(cp, cfg)
 				if err != nil {
@@ -1356,7 +1429,8 @@ func testHtlcTimeout(t *testing.T, resolution lnwallet.OutgoingHtlcResolution,
 				}
 
 				resolver.Supplement(htlc)
-				resolver.htlcResolution = resolution
+				resolver.initLogger("htlcTimeoutResolver")
+
 				return resolver
 			},
 		)