contractcourt/succes+timeout resolver: extract waitForSpend logic

contractcourt/commit_sweep_resolver.go

@@ -108,6 +108,32 @@ func waitForHeight(waitHeight uint32, notifier chainntnfs.ChainNotifier,
 	}
 }
 
+// waitForSpend waits for the given outpoint to be spent, and returns the
+// details of the spending tx.
+func waitForSpend(op *wire.OutPoint, pkScript []byte, heightHint uint32,
+	notifier chainntnfs.ChainNotifier, quit <-chan struct{}) (
+	*chainntnfs.SpendDetail, error) {
+
+	spendNtfn, err := notifier.RegisterSpendNtfn(
+		op, pkScript, heightHint,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	select {
+	case spendDetail, ok := <-spendNtfn.Spend:
+		if !ok {
+			return nil, errResolverShuttingDown
+		}
+
+		return spendDetail, nil
+
+	case <-quit:
+		return nil, errResolverShuttingDown
+	}
+}
+
 // getCommitTxConfHeight waits for confirmation of the commitment tx and returns
 // the confirmation height.
 func (c *commitSweepResolver) getCommitTxConfHeight() (uint32, error) {

contractcourt/htlc_success_resolver.go

@@ -243,33 +243,21 @@ func (h *htlcSuccessResolver) Resolve() (ContractResolver, error) {
 
 	// To wrap this up, we'll wait until the second-level transaction has
 	// been spent, then fully resolve the contract.
-	spendNtfn, err := h.Notifier.RegisterSpendNtfn(
+	log.Infof("%T(%x): waiting for second-level HTLC output to be spent "+
+		"after csv_delay=%v", h, h.htlc.RHash[:], h.htlcResolution.CsvDelay)
+
+	spend, err := waitForSpend(
 		&h.htlcResolution.ClaimOutpoint,
 		h.htlcResolution.SweepSignDesc.Output.PkScript,
-		h.broadcastHeight,
+		h.broadcastHeight, h.Notifier, h.quit,
 	)
 	if err != nil {
 		return nil, err
 	}
 
-	log.Infof("%T(%x): waiting for second-level HTLC output to be spent "+
-		"after csv_delay=%v", h, h.htlc.RHash[:], h.htlcResolution.CsvDelay)
-
-	var spendTxid *chainhash.Hash
-	select {
-	case spend, ok := <-spendNtfn.Spend:
-		if !ok {
-			return nil, errResolverShuttingDown
-		}
-		spendTxid = spend.SpenderTxHash
-
-	case <-h.quit:
-		return nil, errResolverShuttingDown
-	}
-
 	h.resolved = true
 	return nil, h.checkpointClaim(
-		spendTxid, channeldb.ResolverOutcomeClaimed,
+		spend.SpenderTxHash, channeldb.ResolverOutcomeClaimed,
 	)
 }
 

contractcourt/htlc_timeout_resolver.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"io"
 
-	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcutil"
 	"github.com/davecgh/go-spew/spew"
@@ -276,37 +275,6 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 		}
 	}
 
-	var spendTxID *chainhash.Hash
-
-	// waitForOutputResolution waits for the HTLC output to be fully
-	// resolved. The output is considered fully resolved once it has been
-	// spent, and the spending transaction has been fully confirmed.
-	waitForOutputResolution := func() error {
-		// We first need to register to see when the HTLC output itself
-		// has been spent by a confirmed transaction.
-		spendNtfn, err := h.Notifier.RegisterSpendNtfn(
-			&h.htlcResolution.ClaimOutpoint,
-			h.htlcResolution.SweepSignDesc.Output.PkScript,
-			h.broadcastHeight,
-		)
-		if err != nil {
-			return err
-		}
-
-		select {
-		case spendDetail, ok := <-spendNtfn.Spend:
-			if !ok {
-				return errResolverShuttingDown
-			}
-			spendTxID = spendDetail.SpenderTxHash
-
-		case <-h.quit:
-			return errResolverShuttingDown
-		}
-
-		return nil
-	}
-
 	// Now that we've handed off the HTLC to the nursery, 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,
@@ -315,12 +283,6 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 	if err != nil {
 		return nil, err
 	}
-	spendNtfn, err := h.Notifier.RegisterSpendNtfn(
-		outpointToWatch, scriptToWatch, h.broadcastHeight,
-	)
-	if err != nil {
-		return nil, err
-	}
 
 	log.Infof("%T(%v): waiting for HTLC output %v to be spent"+
 		"fully confirmed", h, h.htlcResolution.ClaimOutpoint,
@@ -328,21 +290,16 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 
 	// We'll block here until either we exit, or the HTLC output on the
 	// commitment transaction has been spent.
-	var (
-		spend *chainntnfs.SpendDetail
-		ok    bool
+	spend, err := waitForSpend(
+		outpointToWatch, scriptToWatch, h.broadcastHeight,
+		h.Notifier, h.quit,
 	)
-	select {
-	case spend, ok = <-spendNtfn.Spend:
-		if !ok {
-			return nil, errResolverShuttingDown
-		}
-		spendTxID = spend.SpenderTxHash
-
-	case <-h.quit:
-		return nil, errResolverShuttingDown
+	if err != nil {
+		return nil, err
 	}
 
+	spendTxID := spend.SpenderTxHash
+
 	// 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.
@@ -378,10 +335,18 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 	if h.htlcResolution.SignedTimeoutTx != nil {
 		log.Infof("%T(%v): waiting for nursery to spend CSV delayed "+
 			"output", h, h.htlcResolution.ClaimOutpoint)
-		if err := waitForOutputResolution(); err != nil {
+		sweep, err := waitForSpend(
+			&h.htlcResolution.ClaimOutpoint,
+			h.htlcResolution.SweepSignDesc.Output.PkScript,
+			h.broadcastHeight, h.Notifier, h.quit,
+		)
+		if err != nil {
 			return nil, err
 		}
 
+		// Update the spend txid to the hash of the sweep transaction.
+		spendTxID = sweep.SpenderTxHash
+
 		// Once our timeout tx has confirmed, we add a resolution for
 		// our timeoutTx tx first stage transaction.
 		timeoutTx := h.htlcResolution.SignedTimeoutTx