From 0aadddf14dba6c2a20e6a06e180e6ccbdf9acdaf Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 12 Sep 2018 16:33:46 +0200
Subject: [PATCH 01/15] utxonursery: add test coverage for outgoing htlc

---
 utxonursery_test.go | 461 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 461 insertions(+)

diff --git a/utxonursery_test.go b/utxonursery_test.go
index adff6d3e3b0..5919ed55b2e 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -5,14 +5,20 @@ package main
 import (
 	"bytes"
 	"fmt"
+	"github.com/lightningnetwork/lnd/channeldb"
+	"io/ioutil"
+	"os"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/btcsuite/btclog"
 	"github.com/btcsuite/btcutil"
+	"github.com/lightningnetwork/lnd/chainntnfs"
 	"github.com/lightningnetwork/lnd/lnwallet"
 )
 
@@ -379,3 +385,458 @@ func TestBabyOutputSerialization(t *testing.T) {
 
 	}
 }
+
+var testChanPoint = wire.OutPoint{}
+
+type testWriter struct {
+}
+
+func (w testWriter) Write(p []byte) (n int, err error) {
+	os.Stdout.Write(p)
+
+	return len(p), nil
+}
+
+func init() {
+	backendLog := btclog.NewBackend(testWriter{})
+	utxnLog = backendLog.Logger("TEST")
+	utxnLog.SetLevel(btclog.LevelTrace)
+}
+
+type nurseryTestContext struct {
+	nursery     *utxoNursery
+	notifier    *nurseryMockNotifier
+	publishChan chan chainhash.Hash
+	store       *nurseryStoreInterceptor
+}
+
+func createNurseryTestContext(t *testing.T) *nurseryTestContext {
+	// Create a temporary database and connect nurseryStore to it. The
+	// alternative, mocking nurseryStore, is not chosen because there is
+	// still considerable logic in the store.
+
+	tempDirName, err := ioutil.TempDir("", "channeldb")
+	if err != nil {
+		t.Fatalf("unable to create temp dir: %v", err)
+	}
+
+	cdb, err := channeldb.Open(tempDirName)
+	if err != nil {
+		t.Fatalf("unable to open channeldb: %v", err)
+	}
+
+	store, err := newNurseryStore(&chainhash.Hash{}, cdb)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Wrap the store in an inceptor to be able to wait for events in this
+	// test.
+	storeIntercepter := newNurseryStoreInterceptor(store)
+
+	notifier := newNurseryMockNotifier()
+
+	cfg := NurseryConfig{
+		Notifier: notifier,
+		DB:       cdb,
+		Store:    storeIntercepter,
+		ChainIO:  &mockChainIO{},
+		GenSweepScript: func() ([]byte, error) {
+			return []byte{}, nil
+		},
+		Estimator: &mockFeeEstimator{},
+		Signer:    &nurseryMockSigner{},
+	}
+
+	publishChan := make(chan chainhash.Hash, 1)
+	cfg.PublishTransaction = func(tx *wire.MsgTx) error {
+		utxnLog.Tracef("Publishing tx %v", tx.TxHash())
+		publishChan <- tx.TxHash()
+		return nil
+	}
+
+	nursery := newUtxoNursery(&cfg)
+	nursery.Start()
+
+	return &nurseryTestContext{
+		nursery:     nursery,
+		notifier:    notifier,
+		store:       storeIntercepter,
+		publishChan: publishChan,
+	}
+}
+
+func createOutgoingRes(onLocalCommitment bool) *lnwallet.OutgoingHtlcResolution {
+	// Set up an outgoing htlc resolution to hand off to nursery.
+	closeTx := &wire.MsgTx{}
+
+	htlcOp := wire.OutPoint{
+		Hash:  closeTx.TxHash(),
+		Index: 0,
+	}
+
+	outgoingRes := lnwallet.OutgoingHtlcResolution{
+		Expiry: 125,
+		SweepSignDesc: lnwallet.SignDescriptor{
+			Output: &wire.TxOut{
+				Value: 10000,
+			},
+		},
+		CsvDelay: 2,
+	}
+
+	if onLocalCommitment {
+		timeoutTx := &wire.MsgTx{
+			TxIn: []*wire.TxIn{
+				{
+					PreviousOutPoint: htlcOp,
+					Witness:          [][]byte{{}},
+				},
+			},
+			TxOut: []*wire.TxOut{
+				{},
+			},
+		}
+
+		outgoingRes.SignedTimeoutTx = timeoutTx
+	} else {
+		outgoingRes.ClaimOutpoint = htlcOp
+	}
+
+	return &outgoingRes
+}
+
+func incubateTestOutput(t *testing.T, nursery *utxoNursery,
+	onLocalCommitment bool) *lnwallet.OutgoingHtlcResolution {
+
+	outgoingRes := createOutgoingRes(onLocalCommitment)
+
+	// Hand off to nursery.
+	err := nursery.IncubateOutputs(
+		testChanPoint,
+		nil,
+		[]lnwallet.OutgoingHtlcResolution{*outgoingRes},
+		nil,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// IncubateOutputs is executing synchronously and we expect the output
+	// to immediately show up in the report.
+	expectedStage := uint32(2)
+	if onLocalCommitment {
+		expectedStage = 1
+	}
+
+	// TODO(joostjager): Nursery is currently not reporting this limbo
+	// balance.
+	if onLocalCommitment {
+		assertNurseryReport(t, nursery, 1, expectedStage)
+	}
+
+	return outgoingRes
+}
+
+func assertNurseryReport(t *testing.T, nursery *utxoNursery,
+	expectedNofHtlcs int, expectedStage uint32) {
+	report, err := nursery.NurseryReport(&testChanPoint)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(report.htlcs) != expectedNofHtlcs {
+		t.Fatalf("expected %v outputs to be reported, but report "+
+			"only contains %v", expectedNofHtlcs, len(report.htlcs))
+	}
+	htlcReport := report.htlcs[0]
+	if htlcReport.stage != expectedStage {
+		t.Fatalf("expected htlc be advanced to stage %v, but it is "+
+			"reported in stage %v", expectedStage, htlcReport.stage)
+	}
+}
+
+func assertNurseryReportUnavailable(t *testing.T, nursery *utxoNursery) {
+	_, err := nursery.NurseryReport(&testChanPoint)
+	if err != ErrContractNotFound {
+		t.Fatal("expected report to be unavailable")
+	}
+}
+
+func TestNurserySuccessLocal(t *testing.T) {
+	ctx := createNurseryTestContext(t)
+
+	outgoingRes := incubateTestOutput(t, ctx.nursery, true)
+
+	// Notify arrival of block where HTLC CLTV expires.
+	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
+		Height: 125,
+	}
+
+	// This should trigger nursery to publish the timeout tx.
+	select {
+	case <-ctx.publishChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("tx not published")
+	}
+
+	// Confirm the timeout tx. This should promote the HTLC to KNDR state.
+	timeoutTxHash := outgoingRes.SignedTimeoutTx.TxHash()
+	ctx.notifier.confirmTx(&timeoutTxHash, 126)
+
+	// Wait for output to be promoted in store to KNDR.
+	select {
+	case <-ctx.store.cribToKinderChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("output not promoted to KNDR")
+	}
+
+	// Notify arrival of block where second level HTLC unlocks.
+	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
+		Height: 128,
+	}
+
+	// Check final sweep into wallet.
+	testSweep(t, ctx)
+}
+
+func TestNurserySuccessRemote(t *testing.T) {
+	ctx := createNurseryTestContext(t)
+
+	outgoingRes := incubateTestOutput(t, ctx.nursery, false)
+
+	// Notify confirmation of the commitment tx. Is only listened to when
+	// resolving remote commitment tx.
+	//
+	// TODO(joostjager): This is probably not correct?
+	ctx.notifier.confirmTx(&outgoingRes.ClaimOutpoint.Hash, 124)
+
+	// Wait for output to be promoted from PSCL to KNDR.
+	select {
+	case <-ctx.store.preschoolToKinderChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("output not promoted to KNDR")
+	}
+
+	// Notify arrival of block where HTLC CLTV expires.
+	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
+		Height: 125,
+	}
+
+	// Check final sweep into wallet.
+	testSweep(t, ctx)
+}
+
+func testSweep(t *testing.T, ctx *nurseryTestContext) {
+	// Wait for nursery to publish the sweep tx.
+	var sweepTxHash chainhash.Hash
+	select {
+	case sweepTxHash = <-ctx.publishChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("sweep tx not published")
+	}
+
+	// Verify stage in nursery report. HTLCs should now both still be in
+	// stage two.
+	assertNurseryReport(t, ctx.nursery, 1, 2)
+
+	// Confirm the sweep tx.
+	ctx.notifier.confirmTx(&sweepTxHash, 129)
+
+	// Wait for output to be promoted in store to GRAD.
+	select {
+	case <-ctx.store.graduateKinderChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("output not graduated")
+	}
+
+	// As there only was one output to graduate, we expect the channel to be
+	// closed and no report available anymore.
+	assertNurseryReportUnavailable(t, ctx.nursery)
+}
+
+type nurseryStoreInterceptor struct {
+	ns NurseryStore
+
+	// TODO(joostjager): put more useful info through these channels.
+	cribToKinderChan      chan struct{}
+	cribToRemoteSpendChan chan struct{}
+	graduateKinderChan    chan struct{}
+	preschoolToKinderChan chan struct{}
+}
+
+func newNurseryStoreInterceptor(ns NurseryStore) *nurseryStoreInterceptor {
+	return &nurseryStoreInterceptor{
+		ns:                    ns,
+		cribToKinderChan:      make(chan struct{}),
+		cribToRemoteSpendChan: make(chan struct{}),
+		graduateKinderChan:    make(chan struct{}),
+		preschoolToKinderChan: make(chan struct{}),
+	}
+}
+
+func (i *nurseryStoreInterceptor) Incubate(kidOutputs []kidOutput,
+	babyOutputs []babyOutput) error {
+
+	return i.ns.Incubate(kidOutputs, babyOutputs)
+}
+
+func (i *nurseryStoreInterceptor) CribToKinder(babyOutput *babyOutput) error {
+	err := i.ns.CribToKinder(babyOutput)
+
+	i.cribToKinderChan <- struct{}{}
+
+	return err
+}
+
+func (i *nurseryStoreInterceptor) PreschoolToKinder(kidOutput *kidOutput) error {
+	err := i.ns.PreschoolToKinder(kidOutput)
+
+	i.preschoolToKinderChan <- struct{}{}
+
+	return err
+}
+
+func (i *nurseryStoreInterceptor) GraduateKinder(height uint32) error {
+	err := i.ns.GraduateKinder(height)
+
+	i.graduateKinderChan <- struct{}{}
+
+	return err
+}
+
+func (i *nurseryStoreInterceptor) FetchPreschools() ([]kidOutput, error) {
+	return i.ns.FetchPreschools()
+}
+
+func (i *nurseryStoreInterceptor) FetchClass(height uint32) (*wire.MsgTx, []kidOutput, []babyOutput, error) {
+	return i.ns.FetchClass(height)
+}
+
+func (i *nurseryStoreInterceptor) FinalizeKinder(height uint32, tx *wire.MsgTx) error {
+	return i.ns.FinalizeKinder(height, tx)
+}
+
+func (i *nurseryStoreInterceptor) LastFinalizedHeight() (uint32, error) {
+	return i.ns.LastFinalizedHeight()
+}
+
+func (i *nurseryStoreInterceptor) GraduateHeight(height uint32) error {
+	return i.ns.GraduateHeight(height)
+}
+
+func (i *nurseryStoreInterceptor) LastGraduatedHeight() (uint32, error) {
+	return i.ns.LastGraduatedHeight()
+}
+
+func (i *nurseryStoreInterceptor) HeightsBelowOrEqual(height uint32) ([]uint32, error) {
+	return i.ns.HeightsBelowOrEqual(height)
+}
+
+func (i *nurseryStoreInterceptor) ForChanOutputs(chanPoint *wire.OutPoint,
+	callback func([]byte, []byte) error) error {
+
+	return i.ns.ForChanOutputs(chanPoint, callback)
+}
+
+func (i *nurseryStoreInterceptor) ListChannels() ([]wire.OutPoint, error) {
+	return i.ns.ListChannels()
+}
+
+func (i *nurseryStoreInterceptor) IsMatureChannel(chanPoint *wire.OutPoint) (bool, error) {
+	return i.ns.IsMatureChannel(chanPoint)
+}
+
+func (i *nurseryStoreInterceptor) RemoveChannel(chanPoint *wire.OutPoint) error {
+	return i.ns.RemoveChannel(chanPoint)
+}
+
+type mockFeeEstimator struct{}
+
+func (m *mockFeeEstimator) EstimateFeePerKW(
+	numBlocks uint32) (lnwallet.SatPerKWeight, error) {
+
+	return lnwallet.SatPerKWeight(10000), nil
+}
+
+func (m *mockFeeEstimator) Start() error {
+	return nil
+}
+func (m *mockFeeEstimator) Stop() error {
+	return nil
+}
+
+type nurseryMockSigner struct {
+}
+
+func (m *nurseryMockSigner) SignOutputRaw(tx *wire.MsgTx,
+	signDesc *lnwallet.SignDescriptor) ([]byte, error) {
+
+	return []byte{}, nil
+}
+
+func (m *nurseryMockSigner) ComputeInputScript(tx *wire.MsgTx,
+	signDesc *lnwallet.SignDescriptor) (*lnwallet.InputScript, error) {
+
+	return &lnwallet.InputScript{}, nil
+}
+
+type nurseryMockNotifier struct {
+	confChannel map[chainhash.Hash]chan *chainntnfs.TxConfirmation
+	epochChan   chan *chainntnfs.BlockEpoch
+	spendChan   chan *chainntnfs.SpendDetail
+}
+
+func newNurseryMockNotifier() *nurseryMockNotifier {
+	return &nurseryMockNotifier{
+		confChannel: make(map[chainhash.Hash]chan *chainntnfs.TxConfirmation),
+		epochChan:   make(chan *chainntnfs.BlockEpoch),
+		spendChan:   make(chan *chainntnfs.SpendDetail),
+	}
+}
+
+func (m *nurseryMockNotifier) confirmTx(txid *chainhash.Hash, height uint32) {
+	m.getConfChannel(txid) <- &chainntnfs.TxConfirmation{BlockHeight: height}
+}
+
+func (m *nurseryMockNotifier) RegisterConfirmationsNtfn(txid *chainhash.Hash,
+	_ []byte, numConfs, heightHint uint32) (*chainntnfs.ConfirmationEvent, error) {
+
+	return &chainntnfs.ConfirmationEvent{
+		Confirmed: m.getConfChannel(txid),
+	}, nil
+}
+
+func (m *nurseryMockNotifier) getConfChannel(txid *chainhash.Hash) chan *chainntnfs.TxConfirmation {
+	channel, ok := m.confChannel[*txid]
+	if ok {
+		return channel
+	}
+	channel = make(chan *chainntnfs.TxConfirmation)
+	m.confChannel[*txid] = channel
+	return channel
+}
+
+func (m *nurseryMockNotifier) RegisterBlockEpochNtfn(
+	bestBlock *chainntnfs.BlockEpoch) (*chainntnfs.BlockEpochEvent, error) {
+	return &chainntnfs.BlockEpochEvent{
+		Epochs: m.epochChan,
+		Cancel: func() {},
+	}, nil
+}
+
+func (m *nurseryMockNotifier) Start() error {
+	return nil
+}
+
+func (m *nurseryMockNotifier) Stop() error {
+	return nil
+}
+
+func (m *nurseryMockNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint, _ []byte,
+	heightHint uint32) (*chainntnfs.SpendEvent, error) {
+	return &chainntnfs.SpendEvent{
+		Spend:  m.spendChan,
+		Cancel: func() {},
+	}, nil
+}

From a3c763f5ed0f8b187f7d95cd1c290850c334293b Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 12 Sep 2018 22:23:06 +0200
Subject: [PATCH 02/15] utxonursery: test restart behaviour

---
 nursery_store.go    |   1 +
 utxonursery.go      |  14 +++++
 utxonursery_test.go | 147 +++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 146 insertions(+), 16 deletions(-)

diff --git a/nursery_store.go b/nursery_store.go
index 4b6edbc0189..95e34ce1fd8 100644
--- a/nursery_store.go
+++ b/nursery_store.go
@@ -1457,6 +1457,7 @@ func (ns *nurseryStore) getLastGraduatedHeight(tx *bolt.Tx) (uint32, error) {
 // the last graduated height key.
 func (ns *nurseryStore) putLastGraduatedHeight(tx *bolt.Tx, height uint32) error {
 
+	utxnLog.Infof("Log last graduated height at %v", height)
 	// Ensure that the chain bucket for this nursery store exists.
 	chainBucket, err := tx.CreateBucketIfNotExists(ns.pfxChainKey)
 	if err != nil {
diff --git a/utxonursery.go b/utxonursery.go
index f3bf362ec0e..6a92819c9d8 100644
--- a/utxonursery.go
+++ b/utxonursery.go
@@ -333,6 +333,8 @@ func (u *utxoNursery) Stop() error {
 	close(u.quit)
 	u.wg.Wait()
 
+	utxnLog.Infof("UTXO nursery shut down finished")
+
 	return nil
 }
 
@@ -345,6 +347,18 @@ func (u *utxoNursery) IncubateOutputs(chanPoint wire.OutPoint,
 	outgoingHtlcs []lnwallet.OutgoingHtlcResolution,
 	incomingHtlcs []lnwallet.IncomingHtlcResolution) error {
 
+	// Add to wait group because nursery might shut down during execution of
+	// this function. Otherwise it could happen that nursery thinks it is
+	// shut down, but in this function new goroutines were started and stay
+	// around.
+	u.wg.Add(1)
+	defer u.wg.Done()
+	select {
+	case <-u.quit:
+		return fmt.Errorf("nursery shutting down")
+	default:
+	}
+
 	numHtlcs := len(incomingHtlcs) + len(outgoingHtlcs)
 	var (
 		hasCommit bool
diff --git a/utxonursery_test.go b/utxonursery_test.go
index 5919ed55b2e..3aeda2ab25a 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -406,11 +406,16 @@ func init() {
 type nurseryTestContext struct {
 	nursery     *utxoNursery
 	notifier    *nurseryMockNotifier
-	publishChan chan chainhash.Hash
+	publishChan chan wire.MsgTx
 	store       *nurseryStoreInterceptor
+	restart     func() bool
+	receiveTx   func() wire.MsgTx
+	t           *testing.T
 }
 
-func createNurseryTestContext(t *testing.T) *nurseryTestContext {
+func createNurseryTestContext(t *testing.T,
+	checkStartStop func(func()) bool) *nurseryTestContext {
+
 	// Create a temporary database and connect nurseryStore to it. The
 	// alternative, mocking nurseryStore, is not chosen because there is
 	// still considerable logic in the store.
@@ -448,21 +453,62 @@ func createNurseryTestContext(t *testing.T) *nurseryTestContext {
 		Signer:    &nurseryMockSigner{},
 	}
 
-	publishChan := make(chan chainhash.Hash, 1)
+	publishChan := make(chan wire.MsgTx, 1)
 	cfg.PublishTransaction = func(tx *wire.MsgTx) error {
 		utxnLog.Tracef("Publishing tx %v", tx.TxHash())
-		publishChan <- tx.TxHash()
+		publishChan <- *tx
 		return nil
 	}
 
 	nursery := newUtxoNursery(&cfg)
 	nursery.Start()
 
-	return &nurseryTestContext{
+	ctx := &nurseryTestContext{
 		nursery:     nursery,
 		notifier:    notifier,
 		store:       storeIntercepter,
 		publishChan: publishChan,
+		t:           t,
+	}
+
+	ctx.restart = func() bool {
+		return checkStartStop(func() {
+			ctx.nursery.Stop()
+			// Simulate lnd restart.
+			ctx.nursery = newUtxoNursery(ctx.nursery.cfg)
+			ctx.nursery.Start()
+		})
+	}
+
+	ctx.receiveTx = func() wire.MsgTx {
+		var tx wire.MsgTx
+		select {
+		case tx = <-ctx.publishChan:
+			return tx
+		case <-time.After(5 * time.Second):
+			t.Fatalf("tx not published")
+		}
+		return tx
+	}
+
+	// Start with testing an immediate restart.
+	ctx.restart()
+
+	return ctx
+}
+
+func (ctx *nurseryTestContext) finish() {
+	// Add a final restart point in this state
+	ctx.restart()
+
+	ctx.nursery.Stop()
+
+	// We should have consumed and asserted all published transactions in
+	// our unit tests.
+	select {
+	case <-ctx.publishChan:
+		ctx.t.Fatalf("unexpected transactions published")
+	default:
 	}
 }
 
@@ -563,21 +609,69 @@ func assertNurseryReportUnavailable(t *testing.T, nursery *utxoNursery) {
 	}
 }
 
+// testRestartLoop runs the specified test multiple times and in every run it
+// will attempt to execute a restart action in a different location. This is to
+// assert that the unit under test is recovering correctly from restarts.
+func testRestartLoop(t *testing.T, test func(*testing.T,
+	func(func()) bool)) {
+
+	// Start with running the test without any restarts (index zero)
+	restartIdx := 0
+
+	for {
+		currentStartStopIdx := 0
+
+		// checkStartStop is called at every point in the test where a
+		// restart should be exercised. When this function is called as
+		// many times as the current value of currentStartStopIdx, it
+		// will execute startStopFunc.
+		checkStartStop := func(startStopFunc func()) bool {
+			currentStartStopIdx++
+			if restartIdx == currentStartStopIdx {
+				startStopFunc()
+
+				return true
+			}
+			return false
+		}
+
+		t.Run(fmt.Sprintf("restart_%v", restartIdx),
+			func(t *testing.T) {
+				test(t, checkStartStop)
+			})
+
+		// Exit the loop when all restart points have been tested.
+		if currentStartStopIdx == restartIdx {
+			return
+		}
+		restartIdx++
+	}
+}
+
 func TestNurserySuccessLocal(t *testing.T) {
-	ctx := createNurseryTestContext(t)
+	testRestartLoop(t, testNurserySuccessLocal)
+}
+
+func testNurserySuccessLocal(t *testing.T,
+	checkStartStop func(func()) bool) {
+
+	ctx := createNurseryTestContext(t, checkStartStop)
 
 	outgoingRes := incubateTestOutput(t, ctx.nursery, true)
 
+	ctx.restart()
+
 	// Notify arrival of block where HTLC CLTV expires.
 	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
 		Height: 125,
 	}
 
 	// This should trigger nursery to publish the timeout tx.
-	select {
-	case <-ctx.publishChan:
-	case <-time.After(5 * time.Second):
-		t.Fatalf("tx not published")
+	ctx.receiveTx()
+
+	if ctx.restart() {
+		// Restart should retrigger broadcast of timeout tx.
+		ctx.receiveTx()
 	}
 
 	// Confirm the timeout tx. This should promote the HTLC to KNDR state.
@@ -591,6 +685,8 @@ func TestNurserySuccessLocal(t *testing.T) {
 		t.Fatalf("output not promoted to KNDR")
 	}
 
+	ctx.restart()
+
 	// Notify arrival of block where second level HTLC unlocks.
 	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
 		Height: 128,
@@ -601,10 +697,22 @@ func TestNurserySuccessLocal(t *testing.T) {
 }
 
 func TestNurserySuccessRemote(t *testing.T) {
-	ctx := createNurseryTestContext(t)
+	testRestartLoop(t, testNurserySuccessRemote)
+}
+
+func testNurserySuccessRemote(t *testing.T,
+	checkStartStop func(func()) bool) {
+
+	ctx := createNurseryTestContext(t, checkStartStop)
 
 	outgoingRes := incubateTestOutput(t, ctx.nursery, false)
 
+	// TODO(joostjager): for this restart to work, channel db needs to be
+	// mocked. Waiting for merge of #1847 to completely remove reading
+	// closed channel summary.
+
+	// ctx.restart()
+
 	// Notify confirmation of the commitment tx. Is only listened to when
 	// resolving remote commitment tx.
 	//
@@ -618,6 +726,8 @@ func TestNurserySuccessRemote(t *testing.T) {
 		t.Fatalf("output not promoted to KNDR")
 	}
 
+	ctx.restart()
+
 	// Notify arrival of block where HTLC CLTV expires.
 	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
 		Height: 125,
@@ -629,11 +739,11 @@ func TestNurserySuccessRemote(t *testing.T) {
 
 func testSweep(t *testing.T, ctx *nurseryTestContext) {
 	// Wait for nursery to publish the sweep tx.
-	var sweepTxHash chainhash.Hash
-	select {
-	case sweepTxHash = <-ctx.publishChan:
-	case <-time.After(5 * time.Second):
-		t.Fatalf("sweep tx not published")
+	sweepTx := ctx.receiveTx()
+
+	if ctx.restart() {
+		// Restart will trigger rebroadcast of sweep tx.
+		sweepTx = ctx.receiveTx()
 	}
 
 	// Verify stage in nursery report. HTLCs should now both still be in
@@ -641,6 +751,7 @@ func testSweep(t *testing.T, ctx *nurseryTestContext) {
 	assertNurseryReport(t, ctx.nursery, 1, 2)
 
 	// Confirm the sweep tx.
+	sweepTxHash := sweepTx.TxHash()
 	ctx.notifier.confirmTx(&sweepTxHash, 129)
 
 	// Wait for output to be promoted in store to GRAD.
@@ -650,9 +761,13 @@ func testSweep(t *testing.T, ctx *nurseryTestContext) {
 		t.Fatalf("output not graduated")
 	}
 
+	ctx.restart()
+
 	// As there only was one output to graduate, we expect the channel to be
 	// closed and no report available anymore.
 	assertNurseryReportUnavailable(t, ctx.nursery)
+
+	ctx.finish()
 }
 
 type nurseryStoreInterceptor struct {

From f910a17248ec9a311fbc483dd847fc3c68901ba5 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 12 Sep 2018 16:46:28 +0200
Subject: [PATCH 03/15] utxonursery: report unconfirmed htlc on commit tx

---
 utxonursery.go      | 15 +++++++++++----
 utxonursery_test.go |  7 +------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/utxonursery.go b/utxonursery.go
index 6a92819c9d8..2b45d8bf009 100644
--- a/utxonursery.go
+++ b/utxonursery.go
@@ -548,11 +548,18 @@ func (u *utxoNursery) NurseryReport(
 				case lnwallet.CommitmentTimeLock:
 					report.AddLimboCommitment(&kid)
 
-				// An HTLC output on our commitment transaction
-				// where the second-layer transaction hasn't
-				// yet confirmed.
 				case lnwallet.HtlcAcceptedSuccessSecondLevel:
+					// An HTLC output on our commitment transaction
+					// where the second-layer transaction hasn't
+					// yet confirmed.
 					report.AddLimboStage1SuccessHtlc(&kid)
+
+				case lnwallet.HtlcOfferedRemoteTimeout:
+					// This is an HTLC output on the
+					// commitment transaction of the remote
+					// party. We are waiting for the CLTV
+					// timelock expire.
+					report.AddLimboDirectHtlc(&kid)
 				}
 
 			case bytes.HasPrefix(k, kndrPrefix):
@@ -1510,7 +1517,7 @@ func (c *contractMaturityReport) AddLimboStage1TimeoutHtlc(baby *babyOutput) {
 
 // AddLimboDirectHtlc adds a direct HTLC on the commitment transaction of the
 // remote party to the maturity report. This a CLTV time-locked output that
-// hasn't yet expired.
+// has or hasn't expired yet.
 func (c *contractMaturityReport) AddLimboDirectHtlc(kid *kidOutput) {
 	c.limboBalance += kid.Amount()
 
diff --git a/utxonursery_test.go b/utxonursery_test.go
index 3aeda2ab25a..881521abdbc 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -574,12 +574,7 @@ func incubateTestOutput(t *testing.T, nursery *utxoNursery,
 	if onLocalCommitment {
 		expectedStage = 1
 	}
-
-	// TODO(joostjager): Nursery is currently not reporting this limbo
-	// balance.
-	if onLocalCommitment {
-		assertNurseryReport(t, nursery, 1, expectedStage)
-	}
+	assertNurseryReport(t, nursery, 1, expectedStage)
 
 	return outgoingRes
 }

From 16314a87d7e7547f4802b5df192dd14c5e8bf3c4 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 12 Sep 2018 16:51:24 +0200
Subject: [PATCH 04/15] utxonursery: handle remote spends on local commit

---
 nursery_store.go    | 146 +++++++++++++++++++++++++++++++++++++++++-
 utxonursery.go      | 150 +++++++++++++++++++++++++++++++++++++++++++-
 utxonursery_test.go | 106 +++++++++++++++++++++++++++++--
 3 files changed, 393 insertions(+), 9 deletions(-)

diff --git a/nursery_store.go b/nursery_store.go
index 95e34ce1fd8..007490b1ce5 100644
--- a/nursery_store.go
+++ b/nursery_store.go
@@ -109,6 +109,10 @@ type NurseryStore interface {
 	// kidOutput's CSV delay.
 	CribToKinder(*babyOutput) error
 
+	// CribToRemoteSpend marks an output as spend by the remote party.
+	// TODO(joostjager): Generalize this function to other states too.
+	CribToRemoteSpend(*babyOutput) error
+
 	// PreschoolToKinder atomically moves a kidOutput from the preschool
 	// bucket to the kindergarten bucket. This transition should be
 	// executed after receiving confirmation of the preschool output.
@@ -128,6 +132,10 @@ type NurseryStore interface {
 	// the preschool bucket.
 	FetchPreschools() ([]kidOutput, error)
 
+	// FetchCribs returns a list of all outputs currently stored in the
+	// cribs bucket.
+	FetchCribs() ([]babyOutput, error)
+
 	// FetchClass returns a list of kindergarten and crib outputs whose
 	// timelocks expire at the given height. If the kindergarten class at
 	// this height hash been finalized previously, via FinalizeKinder, it
@@ -231,6 +239,8 @@ var (
 	// this serves as a persistent marker that the nursery should mark the
 	// channel fully closed in the channeldb.
 	gradPrefix = []byte("grad")
+
+	spndPrefix = []byte("spnd")
 )
 
 // prefixChainKey creates the root level keys for the nursery store. The keys
@@ -412,6 +422,67 @@ func (ns *nurseryStore) CribToKinder(bby *babyOutput) error {
 	})
 }
 
+func (ns *nurseryStore) CribToRemoteSpend(bby *babyOutput) error {
+	return ns.db.Update(func(tx *bolt.Tx) error {
+
+		// First, retrieve or create the channel bucket corresponding to
+		// the baby output's origin channel point.
+		chanPoint := bby.OriginChanPoint()
+		chanBucket, err := ns.createChannelBucket(tx, chanPoint)
+		if err != nil {
+			return err
+		}
+
+		// The babyOutput should currently be stored in the crib bucket.
+		// So, we create a key that prefixes the babyOutput's outpoint
+		// with the crib prefix, allowing us to reference it in the
+		// store.
+		pfxOutputKey, err := prefixOutputKey(cribPrefix, bby.OutPoint())
+		if err != nil {
+			return err
+		}
+
+		// Since the babyOutput is being moved to the spend bucket, we
+		// remove the entry from the channel bucket under the
+		// crib-prefixed outpoint key.
+		if err := chanBucket.Delete(pfxOutputKey); err != nil {
+			return err
+		}
+
+		// Remove the crib output's entry in the height index.
+		err = ns.removeOutputFromHeight(tx, bby.expiry, chanPoint,
+			pfxOutputKey)
+		if err != nil {
+			return err
+		}
+
+		// Since we are moving this output from the crib bucket to the
+		// spend bucket, we overwrite the existing prefix of this key
+		// with the spend prefix.
+		copy(pfxOutputKey, spndPrefix)
+
+		// Now, serialize babyOutput's encapsulated kidOutput such that
+		// it can be written to the channel bucket under the new
+		// spend-prefixed key.
+		var kidBuffer bytes.Buffer
+		if err := bby.kidOutput.Encode(&kidBuffer); err != nil {
+			return err
+		}
+		kidBytes := kidBuffer.Bytes()
+
+		// Persist the serialized kidOutput under the spend-prefixed
+		// outpoint key.
+		if err := chanBucket.Put(pfxOutputKey, kidBytes); err != nil {
+			return err
+		}
+
+		utxnLog.Tracef("Transitioning (crib -> spnd) output for "+
+			"chan_point=%v", chanPoint)
+
+		return nil
+	})
+}
+
 // PreschoolToKinder atomically moves a kidOutput from the preschool bucket to
 // the kindergarten bucket. This transition should be executed after receiving
 // confirmation of the preschool output's commitment transaction.
@@ -494,7 +565,7 @@ func (ns *nurseryStore) PreschoolToKinder(kid *kidOutput) error {
 			maturityHeight = lastGradHeight + 1
 		}
 
-		utxnLog.Infof("Transitioning (crib -> kid) output for "+
+		utxnLog.Infof("Transitioning (preschool -> kid) output for "+
 			"chan_point=%v at height_index=%v", chanPoint,
 			maturityHeight)
 
@@ -762,6 +833,79 @@ func (ns *nurseryStore) FetchPreschools() ([]kidOutput, error) {
 	return kids, nil
 }
 
+// FetchPreschools returns a list of all outputs currently stored in the
+// preschool bucket.
+func (ns *nurseryStore) FetchCribs() ([]babyOutput, error) {
+	var babies []babyOutput
+	if err := ns.db.View(func(tx *bolt.Tx) error {
+
+		// Retrieve the existing chain bucket for this nursery store.
+		chainBucket := tx.Bucket(ns.pfxChainKey)
+		if chainBucket == nil {
+			return nil
+		}
+
+		// Load the existing channel index from the chain bucket.
+		chanIndex := chainBucket.Bucket(channelIndexKey)
+		if chanIndex == nil {
+			return nil
+		}
+
+		// Construct a list of all channels in the channel index that
+		// are currently being tracked by the nursery store.
+		var activeChannels [][]byte
+		if err := chanIndex.ForEach(func(chanBytes, _ []byte) error {
+			activeChannels = append(activeChannels, chanBytes)
+			return nil
+		}); err != nil {
+			return err
+		}
+
+		// Iterate over all of the accumulated channels, and do a prefix
+		// scan inside of each channel bucket. Each output found that
+		// has a preschool prefix will be deserialized into a kidOutput,
+		// and added to our list of preschool outputs to return to the
+		// caller.
+		for _, chanBytes := range activeChannels {
+			// Retrieve the channel bucket associated with this
+			// channel.
+			chanBucket := chanIndex.Bucket(chanBytes)
+			if chanBucket == nil {
+				continue
+			}
+
+			// All of the outputs of interest will start with the
+			// "pscl" prefix. So, we will perform a prefix scan of
+			// the channel bucket to efficiently enumerate all the
+			// desired outputs.
+			c := chanBucket.Cursor()
+			for k, v := c.Seek(cribPrefix); bytes.HasPrefix(
+				k, cribPrefix); k, v = c.Next() {
+
+				// Deserialize each output as a kidOutput, since
+				// this should have been the type that was
+				// serialized when it was written to disk.
+				var cribOutput babyOutput
+				cribReader := bytes.NewReader(v)
+				err := cribOutput.Decode(cribReader)
+				if err != nil {
+					return err
+				}
+
+				// Add the deserialized output to our list of
+				// preschool outputs.
+				babies = append(babies, cribOutput)
+			}
+		}
+
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+
+	return babies, nil
+}
+
 // HeightsBelowOrEqual returns a slice of all non-empty heights in the height
 // index at or below the provided upper bound.
 func (ns *nurseryStore) HeightsBelowOrEqual(height uint32) ([]uint32, error) {
diff --git a/utxonursery.go b/utxonursery.go
index 2b45d8bf009..5f2896d039d 100644
--- a/utxonursery.go
+++ b/utxonursery.go
@@ -298,7 +298,7 @@ func (u *utxoNursery) Start() error {
 	// for the force closed commitment txn to confirm, or any second-layer
 	// HTLC success transactions.
 	//
-	// NOTE: The next two steps *may* spawn go routines, thus from this
+	// NOTE: The next three steps *may* spawn go routines, thus from this
 	// point forward, we must close the nursery's quit channel if we detect
 	// any failures during startup to ensure they terminate.
 	if err := u.reloadPreschool(); err != nil {
@@ -307,7 +307,15 @@ func (u *utxoNursery) Start() error {
 		return err
 	}
 
-	// 3. Replay all crib and kindergarten outputs from last pruned to
+	// 3. Restart spend ntfns for any crib outputs, which are waiting
+	// for the CTLV to expire.
+	if err := u.reloadCrib(); err != nil {
+		newBlockChan.Cancel()
+		close(u.quit)
+		return err
+	}
+
+	// 4. Replay all crib and kindergarten outputs from last pruned to
 	// current best height.
 	if err := u.reloadClasses(lastGraduatedHeight); err != nil {
 		newBlockChan.Cancel()
@@ -342,6 +350,9 @@ func (u *utxoNursery) Stop() error {
 // outputs from an existing commitment transaction. Outputs need to incubate if
 // they're CLTV absolute time locked, or if they're CSV relative time locked.
 // Once all outputs reach maturity, they'll be swept back into the wallet.
+//
+// NOTE: For crib outputs, it is assumed that the commit tx is confirmed when
+// IncubateOutputs is called.
 func (u *utxoNursery) IncubateOutputs(chanPoint wire.OutPoint,
 	commitResolution *lnwallet.CommitOutputResolution,
 	outgoingHtlcs []lnwallet.OutgoingHtlcResolution,
@@ -476,6 +487,16 @@ func (u *utxoNursery) IncubateOutputs(chanPoint wire.OutPoint,
 				return err
 			}
 		}
+
+		// Start watch this output for (remote) spends.
+		// TODO(joostjager): This watching should happen for all
+		// commitment outputs and for timeout tx outputs as well. The
+		// remote party cannot only sweep with pre-image of the payment
+		// hash, but with the revocation key too.
+		err := u.registerCribSpend(&babyOutput, uint32(bestHeight))
+		if err != nil {
+			return err
+		}
 	}
 
 	// 3. If we are incubating any preschool outputs, register for a
@@ -524,6 +545,19 @@ func (u *utxoNursery) NurseryReport(
 			// Each crib output represents a stage one htlc, and
 			// will contribute towards the limbo balance.
 			report.AddLimboStage1TimeoutHtlc(&baby)
+		case bytes.HasPrefix(k, spndPrefix):
+			// Cribs outputs are the only kind currently stored as
+			// baby outputs.
+			var kid kidOutput
+			err := kid.Decode(bytes.NewReader(v))
+			if err != nil {
+				return err
+			}
+
+			// Remote spend outputs will be mentioned in the report,
+			// but do not contribute towards the limbo balance
+			// anymore.
+			report.AddLostHtlc(&kid)
 
 		case bytes.HasPrefix(k, psclPrefix),
 			bytes.HasPrefix(k, kndrPrefix),
@@ -669,6 +703,30 @@ func (u *utxoNursery) reloadPreschool() error {
 	return nil
 }
 
+// reloadCrib re-initializes the chain notifier with all of the outputs
+// that had been saved to the "crib" database bucket prior to shutdown.
+func (u *utxoNursery) reloadCrib() error {
+	cribOutputs, err := u.cfg.Store.FetchCribs()
+	if err != nil {
+		return err
+	}
+
+	// For each of the crib outputs stored in the nursery store, load
+	// its close summary from disk so that we can get an accurate height
+	// hint from which to start our range for spend notifications.
+	for i := range cribOutputs {
+		baby := &cribOutputs[i]
+
+		// TODO(joostjager): set height hint after #1847 is merged.
+		err = u.registerCribSpend(baby, 0)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // reloadClasses reinitializes any height-dependent state transitions for which
 // the utxonursery has not received confirmation, and replays the graduation of
 // all kindergarten and crib outputs for heights that have not been finalized.
@@ -1329,6 +1387,78 @@ func (u *utxoNursery) waitForTimeoutConf(baby *babyOutput,
 		"kindergarten", baby.OutPoint())
 }
 
+func (u *utxoNursery) registerCribSpend(baby *babyOutput,
+	heightHint uint32) error {
+
+	outPointToWatch := baby.timeoutTx.TxIn[0].PreviousOutPoint
+	witness := baby.timeoutTx.TxIn[0].Witness
+	scriptToWatch, err := lnwallet.WitnessScriptHash(
+		witness[len(witness)-1],
+	)
+	if err != nil {
+		return err
+	}
+
+	// First, we'll register for a spend notification for this output. If
+	// the remote party sweeps with the pre-image, we'll be notified.
+	spendNtfn, err := u.cfg.Notifier.RegisterSpendNtfn(
+		&outPointToWatch, scriptToWatch, heightHint,
+	)
+	if err != nil {
+		return err
+	}
+
+	u.wg.Add(1)
+	go u.waitForCribSpend(baby, spendNtfn)
+
+	return nil
+}
+
+func (u *utxoNursery) waitForCribSpend(baby *babyOutput,
+	spendNtfn *chainntnfs.SpendEvent) {
+
+	defer u.wg.Done()
+
+	timeoutTxHash := baby.timeoutTx.TxHash()
+
+	select {
+	case spendDetail, ok := <-spendNtfn.Spend:
+		if !ok {
+			utxnLog.Errorf("Notification chan "+
+				"closed, can't monitor output %v",
+				baby.OutPoint())
+			return
+		}
+
+		// Check if it is our own timeout transaction that spends. If
+		// so, no action is needed. Output handling will continue when
+		// transaction sufficiently confirms.
+		if bytes.Equal(spendDetail.SpenderTxHash[:], timeoutTxHash[:]) {
+			utxnLog.Trace("Timeout tx spent output")
+			return
+		}
+	case <-u.quit:
+		spendNtfn.Cancel()
+		return
+	}
+
+	// Remote party has spent the output.
+	utxnLog.Trace("Remote party spend detected!")
+
+	err := u.cfg.Store.CribToRemoteSpend(baby)
+	if err != nil {
+		utxnLog.Errorf("Unable to move %v output "+
+			"to spend bucket", baby.OutPoint())
+
+		// TODO(joostjager): error channel to signal incubator to quit?
+		return
+	}
+
+	// TODO(joostjager): somehow notify the caller of IncubateOutput of the
+	// remote spend. In case of a resolver, it can proceed with extracting
+	// the pre-image and resolving the contract.
+}
+
 // registerPreschoolConf is responsible for subscribing to the confirmation of
 // a commitment transaction, or an htlc success transaction for an incoming
 // HTLC on our commitment transaction.. If successful, the provided preschool
@@ -1469,7 +1599,9 @@ type htlcMaturityReport struct {
 	// stage indicates whether the htlc is in the CLTV-timeout stage (1) or
 	// the CSV-delay stage (2). A stage 1 htlc's maturity height will be set
 	// to its expiry height, while a stage 2 htlc's maturity height will be
-	// set to its confirmation height plus the maturity requirement.
+	// set to its confirmation height plus the maturity requirement. Stage 0
+	// indicates that the htlc was swept by the remote party by either a
+	// revocation key or a payment pre-image.
 	stage uint32
 }
 
@@ -1500,6 +1632,17 @@ func (c *contractMaturityReport) AddRecoveredCommitment(kid *kidOutput) {
 	c.maturityHeight = kid.BlocksToMaturity() + kid.ConfHeight()
 }
 
+// AddLostHtlc adds a lost commitment output to maturity
+// report's htlcs.
+func (c *contractMaturityReport) AddLostHtlc(kid *kidOutput) {
+	c.htlcs = append(c.htlcs, htlcMaturityReport{
+		// TODO: check fields
+		outpoint: *kid.OutPoint(),
+		amount:   kid.Amount(),
+		stage:    0,
+	})
+}
+
 // AddLimboStage1TimeoutHtlc adds an htlc crib output to the maturity report's
 // htlcs, and contributes its amount to the limbo balance.
 func (c *contractMaturityReport) AddLimboStage1TimeoutHtlc(baby *babyOutput) {
@@ -1581,6 +1724,7 @@ func (c *contractMaturityReport) AddRecoveredHtlc(kid *kidOutput) {
 		confHeight:          kid.ConfHeight(),
 		maturityRequirement: kid.BlocksToMaturity(),
 		maturityHeight:      kid.ConfHeight() + kid.BlocksToMaturity(),
+		// TODO(joostjager): set stage?
 	})
 }
 
diff --git a/utxonursery_test.go b/utxonursery_test.go
index 881521abdbc..0255319e82a 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -669,8 +669,14 @@ func testNurserySuccessLocal(t *testing.T,
 		ctx.receiveTx()
 	}
 
-	// Confirm the timeout tx. This should promote the HTLC to KNDR state.
+	// Spend the previous outpoint of the timeout transaction.
+	// Nursery is listening for it.
 	timeoutTxHash := outgoingRes.SignedTimeoutTx.TxHash()
+	ctx.notifier.spendOutpoint(
+		&outgoingRes.SignedTimeoutTx.TxIn[0].PreviousOutPoint,
+		outgoingRes.SignedTimeoutTx)
+
+	// Confirm the timeout tx. This should promote the HTLC to KNDR state.
 	ctx.notifier.confirmTx(&timeoutTxHash, 126)
 
 	// Wait for output to be promoted in store to KNDR.
@@ -745,8 +751,13 @@ func testSweep(t *testing.T, ctx *nurseryTestContext) {
 	// stage two.
 	assertNurseryReport(t, ctx.nursery, 1, 2)
 
-	// Confirm the sweep tx.
 	sweepTxHash := sweepTx.TxHash()
+
+	ctx.notifier.spendOutpoint(
+		&sweepTx.TxIn[0].PreviousOutPoint,
+		&sweepTx)
+
+	// Confirm the sweep tx.
 	ctx.notifier.confirmTx(&sweepTxHash, 129)
 
 	// Wait for output to be promoted in store to GRAD.
@@ -765,6 +776,59 @@ func testSweep(t *testing.T, ctx *nurseryTestContext) {
 	ctx.finish()
 }
 
+func TestNurseryRemoteSpend(t *testing.T) {
+	testRestartLoop(t, testNurseryRemoteSpend)
+}
+
+func testNurseryRemoteSpend(t *testing.T,
+	checkStartStop func(func()) bool) {
+
+	ctx := createNurseryTestContext(t, checkStartStop)
+
+	outgoingRes := incubateTestOutput(t, ctx.nursery, true)
+
+	ctx.restart()
+
+	// Have the remote party spend the commit tx outpoint.
+	remoteSpendTxPreviousOutpoint :=
+		outgoingRes.SignedTimeoutTx.TxIn[0].PreviousOutPoint
+
+	remoteSpendTx := wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: remoteSpendTxPreviousOutpoint,
+			},
+		},
+	}
+	ctx.notifier.spendOutpoint(
+		&remoteSpendTxPreviousOutpoint,
+		&remoteSpendTx)
+
+	// Wait for output to be marked as SPND.
+	select {
+	case <-ctx.store.cribToRemoteSpendChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("not marked as SPND")
+	}
+
+	// Notify arrival of block where HTLC expires.
+	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
+		Height: 125,
+	}
+
+	// This should not trigger the nursery to publish the timeout tx.
+	select {
+	case <-ctx.publishChan:
+		t.Fatalf("tx should not be published anymore")
+	default:
+	}
+
+	// Verify stage in nursery report.
+	assertNurseryReport(t, ctx.nursery, 1, 0)
+
+	ctx.nursery.Stop()
+}
+
 type nurseryStoreInterceptor struct {
 	ns NurseryStore
 
@@ -807,6 +871,14 @@ func (i *nurseryStoreInterceptor) PreschoolToKinder(kidOutput *kidOutput) error
 	return err
 }
 
+func (i *nurseryStoreInterceptor) CribToRemoteSpend(baby *babyOutput) error {
+	err := i.ns.CribToRemoteSpend(baby)
+
+	i.cribToRemoteSpendChan <- struct{}{}
+
+	return err
+}
+
 func (i *nurseryStoreInterceptor) GraduateKinder(height uint32) error {
 	err := i.ns.GraduateKinder(height)
 
@@ -819,6 +891,10 @@ func (i *nurseryStoreInterceptor) FetchPreschools() ([]kidOutput, error) {
 	return i.ns.FetchPreschools()
 }
 
+func (i *nurseryStoreInterceptor) FetchCribs() ([]babyOutput, error) {
+	return i.ns.FetchCribs()
+}
+
 func (i *nurseryStoreInterceptor) FetchClass(height uint32) (*wire.MsgTx, []kidOutput, []babyOutput, error) {
 	return i.ns.FetchClass(height)
 }
@@ -894,14 +970,14 @@ func (m *nurseryMockSigner) ComputeInputScript(tx *wire.MsgTx,
 type nurseryMockNotifier struct {
 	confChannel map[chainhash.Hash]chan *chainntnfs.TxConfirmation
 	epochChan   chan *chainntnfs.BlockEpoch
-	spendChan   chan *chainntnfs.SpendDetail
+	spendChan   map[wire.OutPoint]chan *chainntnfs.SpendDetail
 }
 
 func newNurseryMockNotifier() *nurseryMockNotifier {
 	return &nurseryMockNotifier{
 		confChannel: make(map[chainhash.Hash]chan *chainntnfs.TxConfirmation),
 		epochChan:   make(chan *chainntnfs.BlockEpoch),
-		spendChan:   make(chan *chainntnfs.SpendDetail),
+		spendChan:   make(map[wire.OutPoint]chan *chainntnfs.SpendDetail),
 	}
 }
 
@@ -909,6 +985,16 @@ func (m *nurseryMockNotifier) confirmTx(txid *chainhash.Hash, height uint32) {
 	m.getConfChannel(txid) <- &chainntnfs.TxConfirmation{BlockHeight: height}
 }
 
+func (m *nurseryMockNotifier) spendOutpoint(outpoint *wire.OutPoint,
+	spendingTx *wire.MsgTx) {
+
+	spenderTxHash := spendingTx.TxHash()
+	m.getSpendChannel(outpoint) <- &chainntnfs.SpendDetail{
+		SpenderTxHash: &spenderTxHash,
+		SpendingTx:    spendingTx,
+	}
+}
+
 func (m *nurseryMockNotifier) RegisterConfirmationsNtfn(txid *chainhash.Hash,
 	_ []byte, numConfs, heightHint uint32) (*chainntnfs.ConfirmationEvent, error) {
 
@@ -943,10 +1029,20 @@ func (m *nurseryMockNotifier) Stop() error {
 	return nil
 }
 
+func (m *nurseryMockNotifier) getSpendChannel(outpoint *wire.OutPoint) chan *chainntnfs.SpendDetail {
+	channel, ok := m.spendChan[*outpoint]
+	if ok {
+		return channel
+	}
+	channel = make(chan *chainntnfs.SpendDetail, 1)
+	m.spendChan[*outpoint] = channel
+	return channel
+}
+
 func (m *nurseryMockNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint, _ []byte,
 	heightHint uint32) (*chainntnfs.SpendEvent, error) {
 	return &chainntnfs.SpendEvent{
-		Spend:  m.spendChan,
+		Spend:  m.getSpendChannel(outpoint),
 		Cancel: func() {},
 	}, nil
 }

From d5505b952cb16e8813936adac8d8cc3d299553d6 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Thu, 13 Sep 2018 13:05:03 +0200
Subject: [PATCH 05/15] utxonursery: handle remote spends on remote commit

---
 nursery_store.go    |  97 ++++++++++++++++++++++++++++++---
 utxonursery.go      | 130 +++++++++++++++++++++++++++++++++++++++++++-
 utxonursery_test.go |  85 +++++++++++++++++++++++++++--
 3 files changed, 297 insertions(+), 15 deletions(-)

diff --git a/nursery_store.go b/nursery_store.go
index 007490b1ce5..14a75a36490 100644
--- a/nursery_store.go
+++ b/nursery_store.go
@@ -110,9 +110,11 @@ type NurseryStore interface {
 	CribToKinder(*babyOutput) error
 
 	// CribToRemoteSpend marks an output as spend by the remote party.
-	// TODO(joostjager): Generalize this function to other states too.
 	CribToRemoteSpend(*babyOutput) error
 
+	// KidToRemoteSpend marks an output as spend by the remote party.
+	KidToRemoteSpend(*kidOutput) error
+
 	// PreschoolToKinder atomically moves a kidOutput from the preschool
 	// bucket to the kindergarten bucket. This transition should be
 	// executed after receiving confirmation of the preschool output.
@@ -132,6 +134,10 @@ type NurseryStore interface {
 	// the preschool bucket.
 	FetchPreschools() ([]kidOutput, error)
 
+	// FetchKinder returns a list of all outputs currently stored in
+	// the kindergarten bucket.
+	FetchKinder() ([]kidOutput, error)
+
 	// FetchCribs returns a list of all outputs currently stored in the
 	// cribs bucket.
 	FetchCribs() ([]babyOutput, error)
@@ -483,6 +489,71 @@ func (ns *nurseryStore) CribToRemoteSpend(bby *babyOutput) error {
 	})
 }
 
+func (ns *nurseryStore) KidToRemoteSpend(kid *kidOutput) error {
+	return ns.db.Update(func(tx *bolt.Tx) error {
+
+		// First, retrieve or create the channel bucket corresponding to
+		// the baby output's origin channel point.
+		chanPoint := kid.OriginChanPoint()
+		chanBucket, err := ns.createChannelBucket(tx, chanPoint)
+		if err != nil {
+			return err
+		}
+
+		// The kidOutput should currently be stored in the kndr bucket.
+		// So, we create a key that prefixes the kidOutput's outpoint
+		// with the crib prefix, allowing us to reference it in the
+		// store.
+		pfxOutputKey, err := prefixOutputKey(kndrPrefix, kid.OutPoint())
+		if err != nil {
+			return err
+		}
+
+		// Since the kidOutput is being moved to the spend bucket, we
+		// remove the entry from the channel bucket under the
+		// crib-prefixed outpoint key.
+		if err := chanBucket.Delete(pfxOutputKey); err != nil {
+			return err
+		}
+
+		// Remove the crib output's entry in the height index.
+		//
+		// TODO(joostjager): Here we assume no late graduation into kid
+		// stage has happened! Probably ok, because channel_arb holds
+		// back until commit tx is confirmed.
+		err = ns.removeOutputFromHeight(tx, kid.absoluteMaturity,
+			chanPoint, pfxOutputKey)
+		if err != nil {
+			return err
+		}
+
+		// Since we are moving this output from the crib bucket to the
+		// spend bucket, we overwrite the existing prefix of this key
+		// with the spend prefix.
+		copy(pfxOutputKey, spndPrefix)
+
+		// Now, serialize babyOutput's encapsulated kidOutput such that
+		// it can be written to the channel bucket under the new
+		// spend-prefixed key.
+		var kidBuffer bytes.Buffer
+		if err := kid.Encode(&kidBuffer); err != nil {
+			return err
+		}
+		kidBytes := kidBuffer.Bytes()
+
+		// Persist the serialized kidOutput under the spend-prefixed
+		// outpoint key.
+		if err := chanBucket.Put(pfxOutputKey, kidBytes); err != nil {
+			return err
+		}
+
+		utxnLog.Tracef("Transitioning (kid -> spnd) output for "+
+			"chan_point=%v", chanPoint)
+
+		return nil
+	})
+}
+
 // PreschoolToKinder atomically moves a kidOutput from the preschool bucket to
 // the kindergarten bucket. This transition should be executed after receiving
 // confirmation of the preschool output's commitment transaction.
@@ -763,6 +834,16 @@ func (ns *nurseryStore) FetchClass(
 // FetchPreschools returns a list of all outputs currently stored in the
 // preschool bucket.
 func (ns *nurseryStore) FetchPreschools() ([]kidOutput, error) {
+	return ns.fetchKidsByPrefix(psclPrefix)
+}
+
+// FetchKinder returns a list of all outputs currently stored in the
+// kindergarten bucket.
+func (ns *nurseryStore) FetchKinder() ([]kidOutput, error) {
+	return ns.fetchKidsByPrefix(kndrPrefix)
+}
+
+func (ns *nurseryStore) fetchKidsByPrefix(prefix []byte) ([]kidOutput, error) {
 	var kids []kidOutput
 	if err := ns.db.View(func(tx *bolt.Tx) error {
 
@@ -802,26 +883,26 @@ func (ns *nurseryStore) FetchPreschools() ([]kidOutput, error) {
 			}
 
 			// All of the outputs of interest will start with the
-			// "pscl" prefix. So, we will perform a prefix scan of
+			// specified prefix. So, we will perform a prefix scan of
 			// the channel bucket to efficiently enumerate all the
 			// desired outputs.
 			c := chanBucket.Cursor()
-			for k, v := c.Seek(psclPrefix); bytes.HasPrefix(
-				k, psclPrefix); k, v = c.Next() {
+			for k, v := c.Seek(prefix); bytes.HasPrefix(
+				k, prefix); k, v = c.Next() {
 
 				// Deserialize each output as a kidOutput, since
 				// this should have been the type that was
 				// serialized when it was written to disk.
-				var psclOutput kidOutput
-				psclReader := bytes.NewReader(v)
-				err := psclOutput.Decode(psclReader)
+				var kidOutput kidOutput
+				kidReader := bytes.NewReader(v)
+				err := kidOutput.Decode(kidReader)
 				if err != nil {
 					return err
 				}
 
 				// Add the deserialized output to our list of
 				// preschool outputs.
-				kids = append(kids, psclOutput)
+				kids = append(kids, kidOutput)
 			}
 		}
 
diff --git a/utxonursery.go b/utxonursery.go
index 5f2896d039d..f76d8801ed2 100644
--- a/utxonursery.go
+++ b/utxonursery.go
@@ -298,7 +298,7 @@ func (u *utxoNursery) Start() error {
 	// for the force closed commitment txn to confirm, or any second-layer
 	// HTLC success transactions.
 	//
-	// NOTE: The next three steps *may* spawn go routines, thus from this
+	// NOTE: The next four steps *may* spawn go routines, thus from this
 	// point forward, we must close the nursery's quit channel if we detect
 	// any failures during startup to ensure they terminate.
 	if err := u.reloadPreschool(); err != nil {
@@ -307,7 +307,15 @@ func (u *utxoNursery) Start() error {
 		return err
 	}
 
-	// 3. Restart spend ntfns for any crib outputs, which are waiting
+	// 3. Restart spend ntfns for any kinder outputs, which are waiting
+	// for the CTLV to expire.
+	if err := u.reloadKinder(); err != nil {
+		newBlockChan.Cancel()
+		close(u.quit)
+		return err
+	}
+
+	// 4. Restart spend ntfns for any crib outputs, which are waiting
 	// for the CTLV to expire.
 	if err := u.reloadCrib(); err != nil {
 		newBlockChan.Cancel()
@@ -315,7 +323,7 @@ func (u *utxoNursery) Start() error {
 		return err
 	}
 
-	// 4. Replay all crib and kindergarten outputs from last pruned to
+	// 5. Replay all crib and kindergarten outputs from last pruned to
 	// current best height.
 	if err := u.reloadClasses(lastGraduatedHeight); err != nil {
 		newBlockChan.Cancel()
@@ -703,6 +711,32 @@ func (u *utxoNursery) reloadPreschool() error {
 	return nil
 }
 
+// reloadPreschool re-initializes the chain notifier with all of the outputs
+// that had been saved to the "preschool" database bucket prior to shutdown.
+func (u *utxoNursery) reloadKinder() error {
+	kndrOutputs, err := u.cfg.Store.FetchKinder()
+	if err != nil {
+		return err
+	}
+
+	// For each of the preschool outputs stored in the nursery store, load
+	// its close summary from disk so that we can get an accurate height
+	// hint from which to start our range for spend notifications.
+
+	for i := range kndrOutputs {
+		kid := &kndrOutputs[i]
+
+		// TODO(joostjager): fix height hint
+		heightHint := uint32(0)
+		err = u.registerKinderSpend(kid, heightHint)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // reloadCrib re-initializes the chain notifier with all of the outputs
 // that had been saved to the "crib" database bucket prior to shutdown.
 func (u *utxoNursery) reloadCrib() error {
@@ -1459,6 +1493,85 @@ func (u *utxoNursery) waitForCribSpend(baby *babyOutput,
 	// the pre-image and resolving the contract.
 }
 
+func (u *utxoNursery) registerKinderSpend(kid *kidOutput,
+	heightHint uint32) error {
+
+	// For CLTV locked direct outputs on the remote commit tx, register a
+	// spend notification. These output are still contested and can also be
+	// swept by the remote party. Other kid outputs don't need to be watched
+	// for spends.
+	if kid.WitnessType() != lnwallet.HtlcOfferedRemoteTimeout {
+		return nil
+	}
+
+	outPointToWatch := kid.OutPoint()
+	scriptToWatch := kid.SignDesc().Output.PkScript
+
+	// First, we'll register for a spend notification for this output. If
+	// the remote party sweeps with the pre-image, we'll be notified.
+	spendNtfn, err := u.cfg.Notifier.RegisterSpendNtfn(
+		outPointToWatch, scriptToWatch, heightHint,
+	)
+	if err != nil {
+		return err
+	}
+
+	u.wg.Add(1)
+	go u.waitForKidSpend(kid, spendNtfn)
+
+	return nil
+}
+
+func (u *utxoNursery) waitForKidSpend(kid *kidOutput,
+	spendNtfn *chainntnfs.SpendEvent) {
+
+	defer u.wg.Done()
+
+	select {
+	case spendDetail, ok := <-spendNtfn.Spend:
+		if !ok {
+			utxnLog.Errorf("Notification chan "+
+				"closed, can't monitor output %v",
+				kid.OutPoint())
+			return
+		}
+
+		utxnLog.Tracef("Spend details: %v", spendDetail)
+
+		// Recognize success tx based on script length.
+		isRemoteSuccessTx :=
+			len(spendDetail.SpendingTx.TxIn[0].Witness) == 5
+
+		// If it is not the remote tx, it must be our own sweep tx.
+		// Output handling will continue when transaction sufficiently
+		// confirms.
+
+		if !isRemoteSuccessTx {
+			utxnLog.Trace("Timeout tx spent output")
+			return
+		}
+	case <-u.quit:
+		spendNtfn.Cancel()
+		return
+	}
+
+	// Remote party has spent the output.
+	utxnLog.Trace("Remote party spend detected!")
+
+	err := u.cfg.Store.KidToRemoteSpend(kid)
+	if err != nil {
+		utxnLog.Errorf("Unable to move %v output "+
+			"to spend bucket", kid.OutPoint())
+
+		// TODO(joostjager): error channel to signal incubator to quit?
+		return
+	}
+
+	// TODO(joostjager): somehow notify the caller of IncubateOutput of the
+	// remote spend. In case of a resolver, it can proceed with extracting
+	// the pre-image and resolving the contract.
+}
+
 // registerPreschoolConf is responsible for subscribing to the confirmation of
 // a commitment transaction, or an htlc success transaction for an incoming
 // HTLC on our commitment transaction.. If successful, the provided preschool
@@ -1541,6 +1654,17 @@ func (u *utxoNursery) waitForPreschoolConf(kid *kidOutput,
 			outputType, err)
 		return
 	}
+
+	// Register spend notification for this now confirmed kid output. We
+	// don't handle the case where it is spend before the commit tx is
+	// properly confirmed. This should not happen as channel_arbitrator only
+	// hands off to nursery after confirmation of the commit tx.
+	err = u.registerKinderSpend(kid, kid.ConfHeight())
+	if err != nil {
+		utxnLog.Errorf("Unable to register spend notification for %v "+
+			"kid output: %v", outputType, err)
+		return
+	}
 }
 
 // contractMaturityReport is a report that details the maturity progress of a
diff --git a/utxonursery_test.go b/utxonursery_test.go
index 0255319e82a..b622ba22702 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -588,7 +588,7 @@ func assertNurseryReport(t *testing.T, nursery *utxoNursery,
 
 	if len(report.htlcs) != expectedNofHtlcs {
 		t.Fatalf("expected %v outputs to be reported, but report "+
-			"only contains %v", expectedNofHtlcs, len(report.htlcs))
+			"contains %v", expectedNofHtlcs, len(report.htlcs))
 	}
 	htlcReport := report.htlcs[0]
 	if htlcReport.stage != expectedStage {
@@ -776,11 +776,11 @@ func testSweep(t *testing.T, ctx *nurseryTestContext) {
 	ctx.finish()
 }
 
-func TestNurseryRemoteSpend(t *testing.T) {
-	testRestartLoop(t, testNurseryRemoteSpend)
+func TestNurseryRemoteSpendOnLocal(t *testing.T) {
+	testRestartLoop(t, testNurseryRemoteSpendOnLocal)
 }
 
-func testNurseryRemoteSpend(t *testing.T,
+func testNurseryRemoteSpendOnLocal(t *testing.T,
 	checkStartStop func(func()) bool) {
 
 	ctx := createNurseryTestContext(t, checkStartStop)
@@ -829,6 +829,69 @@ func testNurseryRemoteSpend(t *testing.T,
 	ctx.nursery.Stop()
 }
 
+func TestNurseryRemoteSpendOnRemote(t *testing.T) {
+	testRestartLoop(t, testNurseryRemoteSpendOnRemote)
+}
+
+func testNurseryRemoteSpendOnRemote(t *testing.T,
+	checkStartStop func(func()) bool) {
+
+	ctx := createNurseryTestContext(t, checkStartStop)
+
+	outgoingRes := incubateTestOutput(t, ctx.nursery, false)
+
+	ctx.notifier.confirmTx(&outgoingRes.ClaimOutpoint.Hash, 124)
+
+	// Wait for output to be promoted from PSCL to KNDR.
+	select {
+	case <-ctx.store.preschoolToKinderChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("output not promoted to KNDR")
+	}
+
+	ctx.restart()
+
+	// Remote spend the htlc using a remote success tx. Set witness script
+	// length to 5 so nursery recognizes it properly.
+	remoteSpendTxPreviousOutpoint := outgoingRes.ClaimOutpoint
+
+	remoteSpendTx := wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: remoteSpendTxPreviousOutpoint,
+				Witness:          [][]byte{{}, {}, {}, {}, {}},
+			},
+		},
+	}
+	ctx.notifier.spendOutpoint(
+		&remoteSpendTxPreviousOutpoint,
+		&remoteSpendTx)
+
+	// Wait for output to be marked as SPND.
+	select {
+	case <-ctx.store.kidToRemoteSpendChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("not marked as SPND")
+	}
+
+	// Notify arrival of block where HTLC expires.
+	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
+		Height: 125,
+	}
+
+	// This should not trigger the nursery to publish the timeout tx.
+	select {
+	case <-ctx.publishChan:
+		t.Fatalf("tx should not be published anymore")
+	default:
+	}
+
+	// Verify stage in nursery report.
+	assertNurseryReport(t, ctx.nursery, 1, 0)
+
+	ctx.nursery.Stop()
+}
+
 type nurseryStoreInterceptor struct {
 	ns NurseryStore
 
@@ -837,6 +900,7 @@ type nurseryStoreInterceptor struct {
 	cribToRemoteSpendChan chan struct{}
 	graduateKinderChan    chan struct{}
 	preschoolToKinderChan chan struct{}
+	kidToRemoteSpendChan  chan struct{}
 }
 
 func newNurseryStoreInterceptor(ns NurseryStore) *nurseryStoreInterceptor {
@@ -846,6 +910,7 @@ func newNurseryStoreInterceptor(ns NurseryStore) *nurseryStoreInterceptor {
 		cribToRemoteSpendChan: make(chan struct{}),
 		graduateKinderChan:    make(chan struct{}),
 		preschoolToKinderChan: make(chan struct{}),
+		kidToRemoteSpendChan:  make(chan struct{}),
 	}
 }
 
@@ -879,6 +944,14 @@ func (i *nurseryStoreInterceptor) CribToRemoteSpend(baby *babyOutput) error {
 	return err
 }
 
+func (i *nurseryStoreInterceptor) KidToRemoteSpend(kid *kidOutput) error {
+	err := i.ns.KidToRemoteSpend(kid)
+
+	i.kidToRemoteSpendChan <- struct{}{}
+
+	return err
+}
+
 func (i *nurseryStoreInterceptor) GraduateKinder(height uint32) error {
 	err := i.ns.GraduateKinder(height)
 
@@ -891,6 +964,10 @@ func (i *nurseryStoreInterceptor) FetchPreschools() ([]kidOutput, error) {
 	return i.ns.FetchPreschools()
 }
 
+func (i *nurseryStoreInterceptor) FetchKinder() ([]kidOutput, error) {
+	return i.ns.FetchKinder()
+}
+
 func (i *nurseryStoreInterceptor) FetchCribs() ([]babyOutput, error) {
 	return i.ns.FetchCribs()
 }

From 3488ec81a8de8c889c71342cce3584719e289b59 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Fri, 14 Sep 2018 10:37:59 +0200
Subject: [PATCH 06/15] makefile: dump all goroutines on panic

To provide more diagnostic information when unit tests panic.
---
 Makefile | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Makefile b/Makefile
index c4ee66aa2ef..c8477904de2 100644
--- a/Makefile
+++ b/Makefile
@@ -175,8 +175,8 @@ flakehunter: build
 
 flake-unit:
 	@$(call print, "Flake hunting unit tests.")
-	$(UNIT) -count=1
-	while [ $$? -eq 0 ]; do /bin/sh -c "$(UNIT) -count=1"; done
+	GOTRACEBACK=all $(UNIT) -count=1
+	while [ $$? -eq 0 ]; do /bin/sh -c "GOTRACEBACK=all $(UNIT) -count=1"; done
 
 # ======
 # TRAVIS

From 6d65db947a7a4c67c4f7b8b1417bda1c3e21feb5 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Fri, 14 Sep 2018 10:42:34 +0200
Subject: [PATCH 07/15] utxonursery: assert lingering goroutines

---
 utxonursery_test.go | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/utxonursery_test.go b/utxonursery_test.go
index b622ba22702..ceca05305ea 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -501,6 +501,28 @@ func (ctx *nurseryTestContext) finish() {
 	// Add a final restart point in this state
 	ctx.restart()
 
+	// We assume that when finish is called, nursery has finished all its
+	// goroutines. This implies that the waitgroup is empty.
+	signalChan := make(chan struct{})
+	go func() {
+		ctx.nursery.wg.Wait()
+		close(signalChan)
+	}()
+
+	// The only goroutine that is still expected to be running is
+	// incubator(). Simulate exit of this goroutine.
+	ctx.nursery.wg.Done()
+
+	// We now expect the Wait to succeed.
+	select {
+	case <-signalChan:
+	case <-time.After(time.Second):
+		ctx.t.Fatalf("lingering goroutines detected after test is finished")
+	}
+
+	// Restore waitgroup state to what it was before.
+	ctx.nursery.wg.Add(1)
+
 	ctx.nursery.Stop()
 
 	// We should have consumed and asserted all published transactions in

From 09cfd7a11410ad1cdbae152e25887a0d2618bee4 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Fri, 14 Sep 2018 10:47:13 +0200
Subject: [PATCH 08/15] utxonursery: assert empty database

---
 utxonursery_test.go | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/utxonursery_test.go b/utxonursery_test.go
index ceca05305ea..9b81fd0271a 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"github.com/lightningnetwork/lnd/channeldb"
 	"io/ioutil"
+	"math"
 	"os"
 	"reflect"
 	"testing"
@@ -532,6 +533,25 @@ func (ctx *nurseryTestContext) finish() {
 		ctx.t.Fatalf("unexpected transactions published")
 	default:
 	}
+
+	// Assert that the database is empty. All channels removed and height
+	// index cleared.
+	nurseryChannels, err := ctx.nursery.cfg.Store.ListChannels()
+	if err != nil {
+		ctx.t.Fatal(err)
+	}
+	if len(nurseryChannels) > 0 {
+		ctx.t.Fatalf("Expected all channels to be removed from store")
+	}
+
+	activeHeights, err := ctx.nursery.cfg.Store.HeightsBelowOrEqual(
+		math.MaxUint32)
+	if err != nil {
+		ctx.t.Fatal(err)
+	}
+	if len(activeHeights) > 0 {
+		ctx.t.Fatalf("Expected height index to be empty")
+	}
 }
 
 func createOutgoingRes(onLocalCommitment bool) *lnwallet.OutgoingHtlcResolution {

From e5663d8acde1a6df3a0254942167f50f60df1862 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Fri, 14 Sep 2018 14:07:03 +0200
Subject: [PATCH 09/15] utxonursery: split incubate functions

Previously IncubateOutputs was a sequence of independent actions.
The construct to incubate all outputs of a commit tx in one single
call was not used anymore. This commit splits IncubateOutput into
logical functions.
---
 server.go           |  35 ++++--
 utxonursery.go      | 285 ++++++++++++++++++++++++++------------------
 utxonursery_test.go |   6 +-
 3 files changed, 192 insertions(+), 134 deletions(-)

diff --git a/server.go b/server.go
index 35ef3689b16..73caacc9539 100644
--- a/server.go
+++ b/server.go
@@ -630,20 +630,33 @@ func newServer(listenAddrs []net.Addr, chanDB *channeldb.DB, cc *chainControl,
 			outHtlcRes *lnwallet.OutgoingHtlcResolution,
 			inHtlcRes *lnwallet.IncomingHtlcResolution) error {
 
-			var (
-				inRes  []lnwallet.IncomingHtlcResolution
-				outRes []lnwallet.OutgoingHtlcResolution
-			)
-			if inHtlcRes != nil {
-				inRes = append(inRes, *inHtlcRes)
+			// TODO(joostjager): Propagate function split further
+			// up. To prevent adding three function typed fields in
+			// the config struct, replace by nursery typed field.
+			// But before this can be done, nursery needs to be put
+			// in its own package to fix dependencies.
+			if commitRes == nil &&
+				outHtlcRes == nil && inHtlcRes != nil {
+
+				return s.utxoNursery.IncubateIncomingHtlcOutput(
+					chanPoint, *inHtlcRes)
 			}
-			if outHtlcRes != nil {
-				outRes = append(outRes, *outHtlcRes)
+
+			if commitRes == nil &&
+				outHtlcRes != nil && inHtlcRes == nil {
+
+				return s.utxoNursery.IncubateOutgoingHtlcOutput(
+					chanPoint, *outHtlcRes)
 			}
 
-			return s.utxoNursery.IncubateOutputs(
-				chanPoint, commitRes, outRes, inRes,
-			)
+			if commitRes != nil &&
+				outHtlcRes == nil && inHtlcRes == nil {
+
+				return s.utxoNursery.IncubateCommitOutput(
+					chanPoint, commitRes)
+			}
+
+			panic("invalid combination of inputs")
 		},
 		PreimageDB:   s.witnessBeacon,
 		Notifier:     cc.chainNotifier,
diff --git a/utxonursery.go b/utxonursery.go
index f76d8801ed2..12740eeb7a8 100644
--- a/utxonursery.go
+++ b/utxonursery.go
@@ -354,17 +354,12 @@ func (u *utxoNursery) Stop() error {
 	return nil
 }
 
-// IncubateOutputs sends a request to the utxoNursery to incubate a set of
-// outputs from an existing commitment transaction. Outputs need to incubate if
-// they're CLTV absolute time locked, or if they're CSV relative time locked.
-// Once all outputs reach maturity, they'll be swept back into the wallet.
-//
-// NOTE: For crib outputs, it is assumed that the commit tx is confirmed when
-// IncubateOutputs is called.
-func (u *utxoNursery) IncubateOutputs(chanPoint wire.OutPoint,
-	commitResolution *lnwallet.CommitOutputResolution,
-	outgoingHtlcs []lnwallet.OutgoingHtlcResolution,
-	incomingHtlcs []lnwallet.IncomingHtlcResolution) error {
+// IncubateCommitOutput sends a request to the utxoNursery to incubate the
+// to-self output from an existing commitment transaction. The output needs to
+// incubate it is CSV relative time locked. Once the output reaches maturity,
+// it is swept back into the wallet.
+func (u *utxoNursery) IncubateCommitOutput(chanPoint wire.OutPoint,
+	commitResolution *lnwallet.CommitOutputResolution) error {
 
 	// Add to wait group because nursery might shut down during execution of
 	// this function. Otherwise it could happen that nursery thinks it is
@@ -372,151 +367,203 @@ func (u *utxoNursery) IncubateOutputs(chanPoint wire.OutPoint,
 	// around.
 	u.wg.Add(1)
 	defer u.wg.Done()
-	select {
-	case <-u.quit:
-		return fmt.Errorf("nursery shutting down")
-	default:
-	}
 
-	numHtlcs := len(incomingHtlcs) + len(outgoingHtlcs)
-	var (
-		hasCommit bool
-
-		// Kid outputs can be swept after an initial confirmation
-		// followed by a maturity period.Baby outputs are two stage and
-		// will need to wait for an absolute time out to reach a
-		// confirmation, then require a relative confirmation delay.
-		kidOutputs  = make([]kidOutput, 0, 1+len(incomingHtlcs))
-		babyOutputs = make([]babyOutput, 0, len(outgoingHtlcs))
+	selfOutput := makeKidOutput(
+		&commitResolution.SelfOutPoint,
+		&chanPoint,
+		commitResolution.MaturityDelay,
+		lnwallet.CommitmentTimeLock,
+		&commitResolution.SelfOutputSignDesc,
+		0,
 	)
 
-	// 1. Build all the spendable outputs that we will try to incubate.
-
-	// It could be that our to-self output was below the dust limit. In
-	// that case the commit resolution would be nil and we would not have
-	// that output to incubate.
-	if commitResolution != nil {
-		hasCommit = true
-		selfOutput := makeKidOutput(
-			&commitResolution.SelfOutPoint,
-			&chanPoint,
-			commitResolution.MaturityDelay,
-			lnwallet.CommitmentTimeLock,
-			&commitResolution.SelfOutputSignDesc,
-			0,
-		)
+	// We'll skip any zero valued outputs as this indicates we
+	// don't have a settled balance within the commitment
+	// transaction.
+	if selfOutput.Amount() == 0 {
+		return nil
+	}
 
-		// We'll skip any zero valued outputs as this indicates we
-		// don't have a settled balance within the commitment
-		// transaction.
-		if selfOutput.Amount() > 0 {
-			kidOutputs = append(kidOutputs, selfOutput)
-		}
+	utxnLog.Infof("Incubating Channel(%s) commit output")
+
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	// Persist the output we intend to sweep in the nursery store
+	if err := u.cfg.Store.Incubate([]kidOutput{selfOutput},
+		[]babyOutput{}); err != nil {
+
+		utxnLog.Errorf("unable to begin incubation of Channel(%s): %v",
+			chanPoint, err)
+		return err
 	}
 
-	// TODO(roasbeef): query and see if we already have, if so don't add?
+	err := u.registerPreschoolConf(&selfOutput, u.bestHeight)
+	if err != nil {
+		return err
+	}
 
-	// For each incoming HTLC, we'll register a kid output marked as a
-	// second-layer HTLC output. We effectively skip the baby stage (as the
-	// timelock is zero), and enter the kid stage.
-	for _, htlcRes := range incomingHtlcs {
-		htlcOutput := makeKidOutput(
-			&htlcRes.ClaimOutpoint, &chanPoint, htlcRes.CsvDelay,
-			lnwallet.HtlcAcceptedSuccessSecondLevel,
-			&htlcRes.SweepSignDesc, 0,
-		)
+	return nil
+}
 
-		if htlcOutput.Amount() > 0 {
-			kidOutputs = append(kidOutputs, htlcOutput)
-		}
+// IncubateOutgoingHtlcOutput sends a request to the utxoNursery to incubate a
+// on outgoing htlc output from an existing commitment transaction. This
+// function delegates incubation based on whether the commit tx was published by
+// local or remote.
+func (u *utxoNursery) IncubateOutgoingHtlcOutput(chanPoint wire.OutPoint,
+	outgoingHtlc lnwallet.OutgoingHtlcResolution) error {
+
+	u.wg.Add(1)
+	defer u.wg.Done()
+
+	if outgoingHtlc.SignedTimeoutTx != nil {
+		return u.incubateOutgoingHtlcOnLocal(chanPoint, outgoingHtlc)
 	}
 
-	// For each outgoing HTLC, we'll create a baby output. If this is our
-	// commitment transaction, then we'll broadcast a second-layer
-	// transaction to transition to a kid output. Otherwise, we'll directly
-	// spend once the CLTV delay us up.
-	for _, htlcRes := range outgoingHtlcs {
-		// If this HTLC is on our commitment transaction, then it'll be
-		// a baby output as we need to go to the second level to sweep
-		// it.
-		if htlcRes.SignedTimeoutTx != nil {
-			htlcOutput := makeBabyOutput(&chanPoint, &htlcRes)
+	return u.incubateOutgoingHtlcOnRemote(chanPoint, outgoingHtlc)
+}
 
-			if htlcOutput.Amount() > 0 {
-				babyOutputs = append(babyOutputs, htlcOutput)
-			}
-			continue
-		}
+// incubateOutgoingHtlcOnRemote sends a request to the utxoNursery to incubate
+// an outgoing htlc from an existing remote commitment transaction. The output
+// needs to incubate because it is CLTV absolute time locked. Once the output
+// reaches maturity, it is swept back into the wallet.
+func (u *utxoNursery) incubateOutgoingHtlcOnRemote(chanPoint wire.OutPoint,
+	outgoingHtlc lnwallet.OutgoingHtlcResolution) error {
 
-		// Otherwise, this is actually a kid output as we can sweep it
-		// once the commitment transaction confirms, and the absolute
-		// CLTV lock has expired. We set the CSV delay to zero to
-		// indicate this is actually a CLTV output.
-		htlcOutput := makeKidOutput(
-			&htlcRes.ClaimOutpoint, &chanPoint, 0,
-			lnwallet.HtlcOfferedRemoteTimeout,
-			&htlcRes.SweepSignDesc, htlcRes.Expiry,
-		)
-		kidOutputs = append(kidOutputs, htlcOutput)
+	// We set the CSV delay to zero to indicate this is actually a CLTV
+	// output.
+	htlcOutput := makeKidOutput(
+		&outgoingHtlc.ClaimOutpoint, &chanPoint, 0,
+		lnwallet.HtlcOfferedRemoteTimeout,
+		&outgoingHtlc.SweepSignDesc, outgoingHtlc.Expiry,
+	)
+
+	if htlcOutput.Amount() == 0 {
+		return nil
 	}
 
 	// TODO(roasbeef): if want to handle outgoing on remote commit
 	//  * need ability to cancel in the case that we learn of pre-image or
 	//    remote party pulls
 
-	utxnLog.Infof("Incubating Channel(%s) has-commit=%v, num-htlcs=%d",
-		chanPoint, hasCommit, numHtlcs)
+	utxnLog.Infof("Incubating Channel(%s) outgoing htlc on remote commit",
+		chanPoint)
 
 	u.mu.Lock()
 	defer u.mu.Unlock()
 
-	// 2. Persist the outputs we intended to sweep in the nursery store
-	if err := u.cfg.Store.Incubate(kidOutputs, babyOutputs); err != nil {
+	// Persist the output we intend to sweep in the nursery store
+	if err := u.cfg.Store.Incubate([]kidOutput{htlcOutput},
+		[]babyOutput{}); err != nil {
+
 		utxnLog.Errorf("unable to begin incubation of Channel(%s): %v",
 			chanPoint, err)
 		return err
 	}
 
-	// As an intermediate step, we'll now check to see if any of the baby
-	// outputs has actually _already_ expired. This may be the case if
-	// blocks were mined while we processed this message.
-	_, bestHeight, err := u.cfg.ChainIO.GetBestBlock()
+	err := u.registerPreschoolConf(&htlcOutput, u.bestHeight)
 	if err != nil {
 		return err
 	}
 
-	// We'll examine all the baby outputs just inserted into the database,
-	// if the output has already expired, then we'll *immediately* sweep
-	// it. This may happen if the caller raced a block to call this method.
-	for _, babyOutput := range babyOutputs {
-		if uint32(bestHeight) >= babyOutput.expiry {
-			err = u.sweepCribOutput(uint32(bestHeight), &babyOutput)
-			if err != nil {
-				return err
-			}
-		}
+	return nil
+}
+
+// IncubateOutputs sends a request to the utxoNursery to incubate an outgoing
+// htlc from an existing local commitment transaction. The output needs to
+// incubate because it is CLTV absolute time locked and the second level
+// transaction is CSV time locked. Once the output reaches maturity, it is swept
+// back into the wallet.
+//
+// NOTE: For this crib output, it is assumed that the commit tx is confirmed
+// when incubateOutgoingHtlcOnLocal is called.
+func (u *utxoNursery) incubateOutgoingHtlcOnLocal(chanPoint wire.OutPoint,
+	outgoingHtlc lnwallet.OutgoingHtlcResolution) error {
 
-		// Start watch this output for (remote) spends.
-		// TODO(joostjager): This watching should happen for all
-		// commitment outputs and for timeout tx outputs as well. The
-		// remote party cannot only sweep with pre-image of the payment
-		// hash, but with the revocation key too.
-		err := u.registerCribSpend(&babyOutput, uint32(bestHeight))
+	// This is a HTLC on our commitment transaction,we need to go to the
+	// second level to sweep it.
+	htlcOutput := makeBabyOutput(&chanPoint, &outgoingHtlc)
+
+	if htlcOutput.Amount() == 0 {
+		return nil
+	}
+
+	utxnLog.Infof("Incubating Channel(%s) outgoing htlc on local commit",
+		chanPoint)
+
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	// Persist the output we intend to sweep in the nursery store
+	if err := u.cfg.Store.Incubate([]kidOutput{},
+		[]babyOutput{htlcOutput}); err != nil {
+
+		utxnLog.Errorf("unable to begin incubation of Channel(%s): %v",
+			chanPoint, err)
+		return err
+	}
+
+	// As an intermediate step, we'll now check to see if any of the baby
+	// outputs has actually _already_ expired. This may be the case if
+	// blocks were mined while we processed this message.
+	if uint32(u.bestHeight) >= htlcOutput.expiry {
+		err := u.sweepCribOutput(uint32(u.bestHeight), &htlcOutput)
 		if err != nil {
 			return err
 		}
 	}
 
-	// 3. If we are incubating any preschool outputs, register for a
-	// confirmation notification that will transition it to the
-	// kindergarten bucket.
-	if len(kidOutputs) != 0 {
-		for _, kidOutput := range kidOutputs {
-			err := u.registerPreschoolConf(&kidOutput, u.bestHeight)
-			if err != nil {
-				return err
-			}
-		}
+	// Start watch this output for remote spends.
+	err := u.registerCribSpend(&htlcOutput, uint32(u.bestHeight))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// IncubateOutputs sends a request to the utxoNursery to incubate an incoming
+// htlc output from an existing commitment transaction. The outputs needs to
+// incubate if it is  CSV relative time locked. Once the output reaches
+// maturity, it is swept back into the wallet.
+func (u *utxoNursery) IncubateIncomingHtlcOutput(chanPoint wire.OutPoint,
+	incomingHtlcs lnwallet.IncomingHtlcResolution) error {
+
+	u.wg.Add(1)
+	defer u.wg.Done()
+
+	// TODO(roasbeef): query and see if we already have, if so don't add?
+
+	// For each incoming HTLC, we'll register a kid output marked as a
+	// second-layer HTLC output. We effectively skip the baby stage (as the
+	// timelock is zero), and enter the kid stage.
+	htlcOutput := makeKidOutput(
+		&incomingHtlcs.ClaimOutpoint, &chanPoint, incomingHtlcs.CsvDelay,
+		lnwallet.HtlcAcceptedSuccessSecondLevel,
+		&incomingHtlcs.SweepSignDesc, 0,
+	)
+
+	if htlcOutput.Amount() == 0 {
+		return nil
+	}
+
+	utxnLog.Infof("Incubating Channel(%s) incoming htlc", chanPoint)
+
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	// Persist the output we intend to sweep in the nursery store
+	if err := u.cfg.Store.Incubate([]kidOutput{htlcOutput},
+		[]babyOutput{}); err != nil {
+
+		utxnLog.Errorf("unable to begin incubation of Channel(%s): %v",
+			chanPoint, err)
+		return err
+	}
+
+	err := u.registerPreschoolConf(&htlcOutput, u.bestHeight)
+	if err != nil {
+		return err
 	}
 
 	return nil
diff --git a/utxonursery_test.go b/utxonursery_test.go
index 9b81fd0271a..335bd564c21 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -600,11 +600,9 @@ func incubateTestOutput(t *testing.T, nursery *utxoNursery,
 	outgoingRes := createOutgoingRes(onLocalCommitment)
 
 	// Hand off to nursery.
-	err := nursery.IncubateOutputs(
+	err := nursery.IncubateOutgoingHtlcOutput(
 		testChanPoint,
-		nil,
-		[]lnwallet.OutgoingHtlcResolution{*outgoingRes},
-		nil,
+		*outgoingRes,
 	)
 	if err != nil {
 		t.Fatal(err)

From 79b2885e07f2052cd7da60d2ceba67a778329464 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Mon, 17 Sep 2018 16:29:12 +0200
Subject: [PATCH 10/15] utxonursery: add test for commitment output

---
 utxonursery_test.go | 121 ++++++++++++++++++++++++++++++++++++++------
 1 file changed, 105 insertions(+), 16 deletions(-)

diff --git a/utxonursery_test.go b/utxonursery_test.go
index 335bd564c21..9428c92597a 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -594,6 +594,21 @@ func createOutgoingRes(onLocalCommitment bool) *lnwallet.OutgoingHtlcResolution
 	return &outgoingRes
 }
 
+func createCommitmentRes() *lnwallet.CommitOutputResolution {
+	// Set up a commitment output resolution to hand off to nursery.
+	commitRes := lnwallet.CommitOutputResolution{
+		SelfOutPoint: wire.OutPoint{},
+		SelfOutputSignDesc: lnwallet.SignDescriptor{
+			Output: &wire.TxOut{
+				Value: 10000,
+			},
+		},
+		MaturityDelay: 2,
+	}
+
+	return &commitRes
+}
+
 func incubateTestOutput(t *testing.T, nursery *utxoNursery,
 	onLocalCommitment bool) *lnwallet.OutgoingHtlcResolution {
 
@@ -614,13 +629,14 @@ func incubateTestOutput(t *testing.T, nursery *utxoNursery,
 	if onLocalCommitment {
 		expectedStage = 1
 	}
-	assertNurseryReport(t, nursery, 1, expectedStage)
+	assertNurseryReport(t, nursery, 1, expectedStage, 10000)
 
 	return outgoingRes
 }
 
 func assertNurseryReport(t *testing.T, nursery *utxoNursery,
-	expectedNofHtlcs int, expectedStage uint32) {
+	expectedNofHtlcs int, expectedStage uint32,
+	expectedLimboBalance btcutil.Amount) {
 	report, err := nursery.NurseryReport(&testChanPoint)
 	if err != nil {
 		t.Fatal(err)
@@ -630,10 +646,18 @@ func assertNurseryReport(t *testing.T, nursery *utxoNursery,
 		t.Fatalf("expected %v outputs to be reported, but report "+
 			"contains %v", expectedNofHtlcs, len(report.htlcs))
 	}
-	htlcReport := report.htlcs[0]
-	if htlcReport.stage != expectedStage {
-		t.Fatalf("expected htlc be advanced to stage %v, but it is "+
-			"reported in stage %v", expectedStage, htlcReport.stage)
+
+	if expectedNofHtlcs != 0 {
+		htlcReport := report.htlcs[0]
+		if htlcReport.stage != expectedStage {
+			t.Fatalf("expected htlc be advanced to stage %v, but it is "+
+				"reported in stage %v", expectedStage, htlcReport.stage)
+		}
+	}
+
+	if report.limboBalance != expectedLimboBalance {
+		t.Fatalf("expected limbo balance to be %v, but it is %v instead",
+			expectedLimboBalance, report.limboBalance)
 	}
 }
 
@@ -734,7 +758,9 @@ func testNurserySuccessLocal(t *testing.T,
 	}
 
 	// Check final sweep into wallet.
-	testSweep(t, ctx)
+	testSweepHtlc(t, ctx)
+
+	ctx.finish()
 }
 
 func TestNurserySuccessRemote(t *testing.T) {
@@ -775,10 +801,77 @@ func testNurserySuccessRemote(t *testing.T,
 	}
 
 	// Check final sweep into wallet.
-	testSweep(t, ctx)
+	testSweepHtlc(t, ctx)
+
+	ctx.finish()
+}
+
+func TestNurseryCommitmentOutput(t *testing.T) {
+	testRestartLoop(t, testNurseryCommitmentOutput)
 }
 
-func testSweep(t *testing.T, ctx *nurseryTestContext) {
+func testNurseryCommitmentOutput(t *testing.T,
+	checkStartStop func(func()) bool) {
+
+	ctx := createNurseryTestContext(t, checkStartStop)
+
+	commitRes := createCommitmentRes()
+
+	// Hand off to nursery.
+	err := ctx.nursery.IncubateCommitOutput(
+		testChanPoint,
+		commitRes,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify that commitment output is showing up in nursery report as
+	// limbo balance.
+	assertNurseryReport(t, ctx.nursery, 0, 0, 10000)
+
+	// TODO(joostjager): for this restart to work, channel db needs to be
+	// mocked. Waiting for merge of #1847 to completely remove reading
+	// closed channel summary.
+
+	// ctx.restart()
+
+	// Notify confirmation of the commitment tx.
+	ctx.notifier.confirmTx(&commitRes.SelfOutPoint.Hash, 124)
+
+	// Wait for output to be promoted from PSCL to KNDR.
+	select {
+	case <-ctx.store.preschoolToKinderChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("output not promoted to KNDR")
+	}
+
+	ctx.restart()
+
+	// Notify arrival of block where commit output CSV expires.
+	ctx.notifier.epochChan <- &chainntnfs.BlockEpoch{
+		Height: 126,
+	}
+
+	// Check final sweep into wallet.
+	testSweep(t, ctx, func() {
+		// Check limbo balance after sweep publication
+		assertNurseryReport(t, ctx.nursery, 0, 0, 10000)
+	})
+
+	ctx.finish()
+}
+
+func testSweepHtlc(t *testing.T, ctx *nurseryTestContext) {
+	testSweep(t, ctx, func() {
+		// Verify stage in nursery report. HTLCs should now both still be in
+		// stage two.
+		assertNurseryReport(t, ctx.nursery, 1, 2, 10000)
+	})
+}
+
+func testSweep(t *testing.T, ctx *nurseryTestContext,
+	afterPublishAssert func()) {
 	// Wait for nursery to publish the sweep tx.
 	sweepTx := ctx.receiveTx()
 
@@ -787,9 +880,7 @@ func testSweep(t *testing.T, ctx *nurseryTestContext) {
 		sweepTx = ctx.receiveTx()
 	}
 
-	// Verify stage in nursery report. HTLCs should now both still be in
-	// stage two.
-	assertNurseryReport(t, ctx.nursery, 1, 2)
+	afterPublishAssert()
 
 	sweepTxHash := sweepTx.TxHash()
 
@@ -812,8 +903,6 @@ func testSweep(t *testing.T, ctx *nurseryTestContext) {
 	// As there only was one output to graduate, we expect the channel to be
 	// closed and no report available anymore.
 	assertNurseryReportUnavailable(t, ctx.nursery)
-
-	ctx.finish()
 }
 
 func TestNurseryRemoteSpendOnLocal(t *testing.T) {
@@ -864,7 +953,7 @@ func testNurseryRemoteSpendOnLocal(t *testing.T,
 	}
 
 	// Verify stage in nursery report.
-	assertNurseryReport(t, ctx.nursery, 1, 0)
+	assertNurseryReport(t, ctx.nursery, 1, 0, 0)
 
 	ctx.nursery.Stop()
 }
@@ -927,7 +1016,7 @@ func testNurseryRemoteSpendOnRemote(t *testing.T,
 	}
 
 	// Verify stage in nursery report.
-	assertNurseryReport(t, ctx.nursery, 1, 0)
+	assertNurseryReport(t, ctx.nursery, 1, 0, 0)
 
 	ctx.nursery.Stop()
 }

From ae1b11842d7ec3deaeb06b6b4248409e020a7842 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Mon, 17 Sep 2018 17:31:00 +0200
Subject: [PATCH 11/15] utxonursery: add commit output idempotency test

---
 utxonursery_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/utxonursery_test.go b/utxonursery_test.go
index 9428c92597a..da482e40588 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -862,6 +862,64 @@ func testNurseryCommitmentOutput(t *testing.T,
 	ctx.finish()
 }
 
+func TestNurseryCommitmentOutputIdempotent(t *testing.T) {
+	testRestartLoop(t, testNurseryCommitmentOutputIdempotent)
+}
+
+func testNurseryCommitmentOutputIdempotent(t *testing.T,
+	checkStartStop func(func()) bool) {
+
+	ctx := createNurseryTestContext(t, checkStartStop)
+
+	commitRes := createCommitmentRes()
+
+	// Hand off to nursery.
+	err := ctx.nursery.IncubateCommitOutput(
+		testChanPoint,
+		commitRes,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Notify confirmation of the commitment tx.
+	ctx.notifier.confirmTx(&commitRes.SelfOutPoint.Hash, 124)
+
+	// Wait for output to be promoted from PSCL to KNDR.
+	select {
+	case <-ctx.store.preschoolToKinderChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("output not promoted to KNDR")
+	}
+
+	remoteSpendTx := wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: commitRes.SelfOutPoint,
+			},
+		},
+	}
+
+	ctx.notifier.spendOutpoint(
+		&commitRes.SelfOutPoint,
+		&remoteSpendTx)
+
+	// Wait for output to be promoted in store to GRAD.
+	select {
+	case <-ctx.store.graduateKinderChan:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("output not graduated")
+	}
+
+	ctx.restart()
+
+	// As there only was one output to graduate, we expect the channel to be
+	// closed and no report available anymore.
+	assertNurseryReportUnavailable(t, ctx.nursery)
+
+	ctx.finish()
+}
+
 func testSweepHtlc(t *testing.T, ctx *nurseryTestContext) {
 	testSweep(t, ctx, func() {
 		// Verify stage in nursery report. HTLCs should now both still be in

From 3c28a645379a67949ad8b63c70e0199c4fd5124c Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 19 Sep 2018 08:59:39 -0700
Subject: [PATCH 12/15] cnct: refactor outgoing htlc resolver

---
 contractcourt/contract_resolvers.go | 146 ++++++++++++++++++++--------
 1 file changed, 104 insertions(+), 42 deletions(-)

diff --git a/contractcourt/contract_resolvers.go b/contractcourt/contract_resolvers.go
index 7b8be486663..d68999d7e8f 100644
--- a/contractcourt/contract_resolvers.go
+++ b/contractcourt/contract_resolvers.go
@@ -704,6 +704,15 @@ type htlcOutgoingContestResolver struct {
 	htlcTimeoutResolver
 }
 
+type outgoingState uint8
+
+const (
+	cltvWait uint8 = iota
+	secondLevelConf
+	csvWait
+	sweepWait
+)
+
 // Resolve commences the resolution of this contract. As this contract hasn't
 // yet timed out, we'll wait for one of two things to happen
 //
@@ -787,13 +796,17 @@ func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
 	// the resolution. Otherwise, we fetch this pointer from the input of
 	// the time out transaction.
 	var (
-		outPointToWatch wire.OutPoint
-		scriptToWatch   []byte
-		err             error
+		outPointToWatch      wire.OutPoint
+		scriptToWatch        []byte
+		err                  error
+		secondLevelSpendChan <-chan *chainntnfs.SpendDetail
 	)
 	if h.htlcResolution.SignedTimeoutTx == nil {
 		outPointToWatch = h.htlcResolution.ClaimOutpoint
 		scriptToWatch = h.htlcResolution.SweepSignDesc.Output.PkScript
+
+		// Create dummy channel that will never receive a notification.
+		secondLevelSpendChan = make(chan *chainntnfs.SpendDetail)
 	} else {
 		// If this is the remote party's commitment, then we'll need to
 		// grab watch the output that our timeout transaction points
@@ -808,6 +821,16 @@ func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
 		if err != nil {
 			return nil, err
 		}
+
+		secondLevelOutPointToWatch := h.htlcResolution.ClaimOutpoint
+		secondLevelScriptToWatch := h.htlcResolution.SweepSignDesc.Output.PkScript
+		secondLevelSpendNtfn, err := h.Notifier.RegisterSpendNtfn(
+			&secondLevelOutPointToWatch, secondLevelScriptToWatch, h.broadcastHeight,
+		)
+		if err != nil {
+			return nil, err
+		}
+		secondLevelSpendChan = secondLevelSpendNtfn.Spend
 	}
 
 	// First, we'll register for a spend notification for this output. If
@@ -819,43 +842,42 @@ func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
 		return nil, err
 	}
 
-	// We'll quickly check to see if the output has already been spent.
-	select {
-	// If the output has already been spent, then we can stop early and
-	// sweep the pre-image from the output.
-	case commitSpend, ok := <-spendNtfn.Spend:
-		if !ok {
-			return nil, fmt.Errorf("quitting")
-		}
-
-		// TODO(roasbeef): Checkpoint?
-		return claimCleanUp(commitSpend)
-
-	// If it hasn't, then we'll watch for both the expiration, and the
-	// sweeping out this output.
-	default:
-	}
+	// TODO: Cancel spend notifications on quit?
 
-	// We'll check the current height, if the HTLC has already expired,
-	// then we'll morph immediately into a resolver that can sweep the
-	// HTLC.
-	//
-	// TODO(roasbeef): use grace period instead?
 	_, currentHeight, err := h.ChainIO.GetBestBlock()
 	if err != nil {
 		return nil, err
 	}
 
-	// If the current height is >= expiry-1, then a spend will be valid to
-	// be included in the next block, and we can immediately return the
-	// resolver.
-	if uint32(currentHeight) >= h.htlcResolution.Expiry-1 {
-		log.Infof("%T(%v): HTLC has expired (height=%v, expiry=%v), "+
-			"transforming into timeout resolver", h,
-			h.htlcResolution.ClaimOutpoint)
-		return &h.htlcTimeoutResolver, nil
+	state := cltvWait
+	var secondLevelExpiry int32
+
+	evaluateHeight := func() {
+		switch {
+		case state == cltvWait &&
+			uint32(currentHeight) >= h.htlcResolution.Expiry-1:
+
+			if h.htlcResolution.SignedTimeoutTx == nil {
+				// Start sweep
+
+				state = sweepWait
+			} else {
+				// Publish timeout tx
+
+				state = secondLevelConf
+			}
+
+		case state == csvWait && currentHeight >= secondLevelExpiry-1:
+			// Start sweep
+
+			state = sweepWait
+		}
+
 	}
 
+	// Evaluate current height in case cltv has already expired.
+	evaluateHeight()
+
 	// If we reach this point, then we can't fully act yet, so we'll await
 	// either of our signals triggering: the HTLC expires, or we learn of
 	// the preimage.
@@ -865,6 +887,7 @@ func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
 	}
 	defer blockEpochs.Cancel()
 
+loop:
 	for {
 		select {
 
@@ -878,15 +901,8 @@ func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
 			// If this new height expires the HTLC, then we can
 			// exit early and create a resolver that's capable of
 			// handling the time locked output.
-			newHeight := uint32(newBlock.Height)
-			if newHeight >= h.htlcResolution.Expiry-1 {
-				log.Infof("%T(%v): HTLC has expired "+
-					"(height=%v, expiry=%v), transforming "+
-					"into timeout resolver", h,
-					h.htlcResolution.ClaimOutpoint,
-					newHeight, h.htlcResolution.Expiry)
-				return &h.htlcTimeoutResolver, nil
-			}
+			currentHeight = newBlock.Height
+			evaluateHeight()
 
 		// The output has been spent! This means the preimage has been
 		// revealed on-chain.
@@ -899,12 +915,58 @@ func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
 			// party is by revealing the preimage. So we'll perform
 			// our duties to clean up the contract once it has been
 			// claimed.
-			return claimCleanUp(commitSpend)
+			var isRemoteSpend bool
+			if h.htlcResolution.SignedTimeoutTx != nil {
+				timeoutTxHash := h.htlcResolution.SignedTimeoutTx.TxHash()
+
+				// If spend not by our own timeout tx, it must
+				// be a remote spend.
+				isRemoteSpend = !bytes.Equal(commitSpend.SpenderTxHash[:],
+					timeoutTxHash[:])
+			} else {
+				// Detect remote success tx by witness length.
+				isRemoteSpend = len(commitSpend.SpendingTx.TxIn[0].Witness) == 5
+			}
+
+			if isRemoteSpend {
+				return claimCleanUp(commitSpend)
+			}
+
+			// At this point, the second-level transaction is sufficiently
+			// confirmed, or a transaction directly spending the output is.
+			// Therefore, we can now send back our clean up message.
+			failureMsg := &lnwire.FailPermanentChannelFailure{}
+			if err := h.DeliverResolutionMsg(ResolutionMsg{
+				SourceChan: h.ShortChanID,
+				HtlcIndex:  h.htlcIndex,
+				Failure:    failureMsg,
+			}); err != nil {
+				return nil, err
+			}
+
+			if state == secondLevelConf {
+				state = csvWait
+				secondLevelExpiry = currentHeight +
+					int32(h.htlcResolution.CsvDelay)
+
+				evaluateHeight()
+			} else {
+				break loop
+			}
+
+		case _, ok := <-secondLevelSpendChan:
+			if !ok {
+				return nil, fmt.Errorf("quitting")
+			}
+			break loop
 
 		case <-h.Quit:
 			return nil, fmt.Errorf("resolver cancelled")
 		}
+
 	}
+	h.resolved = true
+	return nil, h.Checkpoint(h)
 }
 
 // Stop signals the resolver to cancel any current resolution processes, and

From acefbfcf2a2560d0dc3a32d857ee7155ed6c0d4c Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 19 Sep 2018 09:12:03 -0700
Subject: [PATCH 13/15] cnct: replace timeoutResolver

---
 contractcourt/briefcase.go          |   12 -
 contractcourt/channel_arbitrator.go |   32 +-
 contractcourt/contract_resolvers.go | 1128 +++++++++++----------------
 3 files changed, 468 insertions(+), 704 deletions(-)

diff --git a/contractcourt/briefcase.go b/contractcourt/briefcase.go
index d7da21ea138..fce16f3a377 100644
--- a/contractcourt/briefcase.go
+++ b/contractcourt/briefcase.go
@@ -355,8 +355,6 @@ func (b *boltArbitratorLog) writeResolver(contractBucket *bolt.Bucket,
 		rType = resolverTimeout
 	case *htlcSuccessResolver:
 		rType = resolverSuccess
-	case *htlcOutgoingContestResolver:
-		rType = resolverOutgoingContest
 	case *htlcIncomingContestResolver:
 		rType = resolverIncomingContest
 	case *commitSweepResolver:
@@ -467,16 +465,6 @@ func (b *boltArbitratorLog) FetchUnresolvedContracts() ([]ContractResolver, erro
 
 				res = successRes
 
-			case resolverOutgoingContest:
-				outContestRes := &htlcOutgoingContestResolver{
-					htlcTimeoutResolver: htlcTimeoutResolver{},
-				}
-				if err := outContestRes.Decode(resReader); err != nil {
-					return err
-				}
-
-				res = outContestRes
-
 			case resolverIncomingContest:
 				inContestRes := &htlcIncomingContestResolver{
 					htlcSuccessResolver: htlcSuccessResolver{},
diff --git a/contractcourt/channel_arbitrator.go b/contractcourt/channel_arbitrator.go
index 55a6e238168..58670bf2e0f 100644
--- a/contractcourt/channel_arbitrator.go
+++ b/contractcourt/channel_arbitrator.go
@@ -1155,7 +1155,7 @@ func (c *ChannelArbitrator) prepContractResolutions(htlcActions ChainActionMap,
 		// If we can timeout the HTLC directly, then we'll create the
 		// proper resolver to do so, who will then cancel the packet
 		// backwards.
-		case HtlcTimeoutAction:
+		case HtlcTimeoutAction, HtlcOutgoingWatchAction:
 			for _, htlc := range htlcs {
 				htlcOp := wire.OutPoint{
 					Hash:  commitHash,
@@ -1212,36 +1212,6 @@ func (c *ChannelArbitrator) prepContractResolutions(htlcActions ChainActionMap,
 				}
 				htlcResolvers = append(htlcResolvers, resolver)
 			}
-
-		// Finally, if this is an outgoing HTLC we've sent, then we'll
-		// launch a resolver to watch for the pre-image (and settle
-		// backwards), or just timeout.
-		case HtlcOutgoingWatchAction:
-			for _, htlc := range htlcs {
-				htlcOp := wire.OutPoint{
-					Hash:  commitHash,
-					Index: uint32(htlc.OutputIndex),
-				}
-
-				resolution, ok := outResolutionMap[htlcOp]
-				if !ok {
-					log.Errorf("ChannelArbitrator(%v) unable to find "+
-						"outgoing resolution: %v",
-						c.cfg.ChanPoint, htlcOp)
-					continue
-				}
-
-				resKit.Quit = make(chan struct{})
-				resolver := &htlcOutgoingContestResolver{
-					htlcTimeoutResolver{
-						htlcResolution:  resolution,
-						broadcastHeight: height,
-						htlcIndex:       htlc.HtlcIndex,
-						ResolverKit:     resKit,
-					},
-				}
-				htlcResolvers = append(htlcResolvers, resolver)
-			}
 		}
 	}
 
diff --git a/contractcourt/contract_resolvers.go b/contractcourt/contract_resolvers.go
index d68999d7e8f..70856d8fac7 100644
--- a/contractcourt/contract_resolvers.go
+++ b/contractcourt/contract_resolvers.go
@@ -134,476 +134,267 @@ func (h *htlcTimeoutResolver) ResolverKey() []byte {
 	return key[:]
 }
 
-// Resolve kicks off full resolution of an outgoing HTLC output. If it's our
-// commitment, it isn't resolved until we see the second level HTLC txn
-// confirmed. If it's the remote party's commitment, we don't resolve until we
-// see a direct sweep via the timeout clause.
+type outgoingState uint8
+
+const (
+	cltvWait uint8 = iota
+	secondLevelConf
+	csvWait
+	sweepWait
+)
+
+// Resolve commences the resolution of this contract. As this contract hasn't
+// yet timed out, we'll wait for one of two things to happen
 //
-// NOTE: Part of the ContractResolver interface.
+//   1. The HTLC expires. In this case, we'll sweep the funds and send a clean
+//      up cancel message to outside sub-systems.
+//
+//   2. The remote party sweeps this HTLC on-chain, in which case we'll add the
+//      pre-image to our global cache, then send a clean up settle message
+//      backwards.
+//
+// When either of these two things happens, we'll create a new resolver which
+// is able to handle the final resolution of the contract. We're only the pivot
+// point.
 func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
-	// If we're already resolved, then we can exit early.
+	// If we're already full resolved, then we don't have anything further
+	// to do.
 	if h.resolved {
 		return nil, nil
 	}
 
-	// If we haven't already sent the output to the utxo nursery, then
-	// we'll do so now.
-	if !h.outputIncubating {
-		log.Tracef("%T(%v): incubating htlc output", h,
-			h.htlcResolution.ClaimOutpoint)
-
-		err := h.IncubateOutputs(h.ChanPoint, nil, &h.htlcResolution, nil)
-		if err != nil {
-			return nil, err
-		}
-
-		h.outputIncubating = true
+	// claimCleanUp is a helper function that's called once the HTLC output
+	// is spent by the remote party. It'll extract the preimage, add it to
+	// the global cache, and finally send the appropriate clean up message.
+	claimCleanUp := func(commitSpend *chainntnfs.SpendDetail) (ContractResolver, error) {
+		// Depending on if this is our commitment or not, then we'll be
+		// looking for a different witness pattern.
+		spenderIndex := commitSpend.SpenderInputIndex
+		spendingInput := commitSpend.SpendingTx.TxIn[spenderIndex]
 
-		if err := h.Checkpoint(h); err != nil {
-			log.Errorf("unable to Checkpoint: %v", err)
-		}
-	}
+		log.Infof("%T(%v): extracting preimage! remote party spent "+
+			"HTLC with tx=%v", h, h.htlcResolution.ClaimOutpoint,
+			spew.Sdump(commitSpend.SpendingTx))
 
-	// 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
+		// If this is the remote party's commitment, then we'll be
+		// looking for them to spend using the second-level success
+		// transaction.
+		var preimage [32]byte
+		if h.htlcResolution.SignedTimeoutTx == nil {
+			// The witness stack when the remote party sweeps the
+			// output to them looks like:
+			//
+			//  * <sender sig> <recvr sig> <preimage> <witness script>
+			copy(preimage[:], spendingInput.Witness[3])
+		} else {
+			// Otherwise, they'll be spending directly from our
+			// commitment output. In which case the witness stack
+			// looks like:
+			//
+			//  * <sig> <preimage> <witness script>
+			copy(preimage[:], spendingInput.Witness[1])
 		}
 
-		select {
-		case _, ok := <-spendNtfn.Spend:
-			if !ok {
-				return fmt.Errorf("notifier quit")
-			}
+		log.Infof("%T(%v): extracting preimage=%x from on-chain "+
+			"spend!", h, h.htlcResolution.ClaimOutpoint, preimage[:])
 
-		case <-h.Quit:
-			return fmt.Errorf("quitting")
+		// With the preimage obtained, we can now add it to the global
+		// cache.
+		if err := h.PreimageDB.AddPreimage(preimage[:]); err != nil {
+			log.Errorf("%T(%v): unable to add witness to cache",
+				h, h.htlcResolution.ClaimOutpoint)
 		}
 
-		return nil
+		// Finally, we'll send the clean up message, mark ourselves as
+		// resolved, then exit.
+		if err := h.DeliverResolutionMsg(ResolutionMsg{
+			SourceChan: h.ShortChanID,
+			HtlcIndex:  h.htlcIndex,
+			PreImage:   &preimage,
+		}); err != nil {
+			return nil, err
+		}
+		h.resolved = true
+		return nil, h.Checkpoint(h)
 	}
 
-	// With the output sent to the nursery, we'll now wait until the output
-	// has been fully resolved before sending the clean up message.
-	//
-	// TODO(roasbeef): need to be able to cancel nursery?
-	//  * if they pull on-chain while we're waiting
+	// Otherwise, we'll watch for two external signals to decide if we'll
+	// morph into another resolver, or fully resolve the contract.
 
-	// If we don't have a second layer transaction, then this is a remote
-	// party's commitment, so we'll watch for a direct spend.
+	// The output we'll be watching for is the *direct* spend from the HTLC
+	// output. If this isn't our commitment transaction, it'll be right on
+	// the resolution. Otherwise, we fetch this pointer from the input of
+	// the time out transaction.
+	var (
+		outPointToWatch      wire.OutPoint
+		scriptToWatch        []byte
+		err                  error
+		secondLevelSpendChan <-chan *chainntnfs.SpendDetail
+	)
 	if h.htlcResolution.SignedTimeoutTx == nil {
-		// We'll block until: the HTLC output has been spent, and the
-		// transaction spending that output is sufficiently confirmed.
-		log.Infof("%T(%v): waiting for nursery to spend CLTV-locked "+
-			"output", h, h.htlcResolution.ClaimOutpoint)
-		if err := waitForOutputResolution(); err != nil {
-			return nil, err
-		}
+		outPointToWatch = h.htlcResolution.ClaimOutpoint
+		scriptToWatch = h.htlcResolution.SweepSignDesc.Output.PkScript
+
+		// Create dummy channel that will never receive a notification.
+		secondLevelSpendChan = make(chan *chainntnfs.SpendDetail)
 	} else {
-		// Otherwise, this is our commitment, so we'll watch for the
-		// second-level transaction to be sufficiently confirmed.
-		secondLevelTXID := h.htlcResolution.SignedTimeoutTx.TxHash()
-		sweepScript := h.htlcResolution.SignedTimeoutTx.TxOut[0].PkScript
-		confNtfn, err := h.Notifier.RegisterConfirmationsNtfn(
-			&secondLevelTXID, sweepScript, 1, h.broadcastHeight,
+		// If this is the remote party's commitment, then we'll need to
+		// grab watch the output that our timeout transaction points
+		// to. We can directly grab the outpoint, then also extract the
+		// witness script (the last element of the witness stack) to
+		// re-construct the pkScipt we need to watch.
+		outPointToWatch = h.htlcResolution.SignedTimeoutTx.TxIn[0].PreviousOutPoint
+		witness := h.htlcResolution.SignedTimeoutTx.TxIn[0].Witness
+		scriptToWatch, err = lnwallet.WitnessScriptHash(
+			witness[len(witness)-1],
 		)
 		if err != nil {
 			return nil, err
 		}
 
-		log.Infof("%T(%v): waiting second-level tx (txid=%v) to be "+
-			"fully confirmed", h, h.htlcResolution.ClaimOutpoint,
-			secondLevelTXID)
-
-		select {
-		case _, ok := <-confNtfn.Confirmed:
-			if !ok {
-				return nil, fmt.Errorf("quitting")
-			}
-
-		case <-h.Quit:
-			return nil, fmt.Errorf("quitting")
+		secondLevelOutPointToWatch := h.htlcResolution.ClaimOutpoint
+		secondLevelScriptToWatch := h.htlcResolution.SweepSignDesc.Output.PkScript
+		secondLevelSpendNtfn, err := h.Notifier.RegisterSpendNtfn(
+			&secondLevelOutPointToWatch, secondLevelScriptToWatch, h.broadcastHeight,
+		)
+		if err != nil {
+			return nil, err
 		}
+		secondLevelSpendChan = secondLevelSpendNtfn.Spend
 	}
 
-	// TODO(roasbeef): need to watch for remote party sweeping with pre-image?
-	//  * have another waiting on spend above, will check the type, if it's
-	//    pre-image, then we'll cancel, and send a clean up back with
-	//    pre-image, also add to preimage cache
-
-	log.Infof("%T(%v): resolving htlc with incoming fail msg, fully "+
-		"confirmed", h, h.htlcResolution.ClaimOutpoint)
-
-	// At this point, the second-level transaction is sufficiently
-	// confirmed, or a transaction directly spending the output is.
-	// Therefore, we can now send back our clean up message.
-	failureMsg := &lnwire.FailPermanentChannelFailure{}
-	if err := h.DeliverResolutionMsg(ResolutionMsg{
-		SourceChan: h.ShortChanID,
-		HtlcIndex:  h.htlcIndex,
-		Failure:    failureMsg,
-	}); err != nil {
+	// First, we'll register for a spend notification for this output. If
+	// the remote party sweeps with the pre-image, we'll be notified.
+	spendNtfn, err := h.Notifier.RegisterSpendNtfn(
+		&outPointToWatch, scriptToWatch, h.broadcastHeight,
+	)
+	if err != nil {
 		return nil, err
 	}
 
-	// Finally, if this was an output on our commitment transaction, we'll
-	// for the second-level HTLC output to be spent, and for that
-	// transaction itself to confirm.
-	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 {
-			return nil, err
-		}
+	// TODO: Cancel spend notifications on quit?
+
+	_, currentHeight, err := h.ChainIO.GetBestBlock()
+	if err != nil {
+		return nil, err
 	}
 
-	// With the clean up message sent, we'll now mark the contract
-	// resolved, and wait.
-	h.resolved = true
-	return nil, h.Checkpoint(h)
-}
+	state := cltvWait
+	var secondLevelExpiry int32
 
-// Stop signals the resolver to cancel any current resolution processes, and
-// suspend.
-//
-// NOTE: Part of the ContractResolver interface.
-func (h *htlcTimeoutResolver) Stop() {
-	close(h.Quit)
-}
+	evaluateHeight := func() {
+		switch {
+		case state == cltvWait &&
+			uint32(currentHeight) >= h.htlcResolution.Expiry-1:
 
-// IsResolved returns true if the stored state in the resolve is fully
-// resolved. In this case the target output can be forgotten.
-//
-// NOTE: Part of the ContractResolver interface.
-func (h *htlcTimeoutResolver) IsResolved() bool {
-	return h.resolved
-}
+			if h.htlcResolution.SignedTimeoutTx == nil {
+				// Start sweep
 
-// Encode writes an encoded version of the ContractResolver into the passed
-// Writer.
-//
-// NOTE: Part of the ContractResolver interface.
-func (h *htlcTimeoutResolver) Encode(w io.Writer) error {
-	// First, we'll write out the relevant fields of the
-	// OutgoingHtlcResolution to the writer.
-	if err := encodeOutgoingResolution(w, &h.htlcResolution); err != nil {
-		return err
-	}
+				state = sweepWait
+			} else {
+				// Publish timeout tx
 
-	// With that portion written, we can now write out the fields specific
-	// to the resolver itself.
-	if err := binary.Write(w, endian, h.outputIncubating); err != nil {
-		return err
-	}
-	if err := binary.Write(w, endian, h.resolved); err != nil {
-		return err
-	}
-	if err := binary.Write(w, endian, h.broadcastHeight); err != nil {
-		return err
-	}
+				state = secondLevelConf
+			}
 
-	if err := binary.Write(w, endian, h.htlcIndex); err != nil {
-		return err
-	}
+		case state == csvWait && currentHeight >= secondLevelExpiry-1:
+			// Start sweep
 
-	return nil
-}
+			state = sweepWait
+		}
 
-// Decode attempts to decode an encoded ContractResolver from the passed Reader
-// instance, returning an active ContractResolver instance.
-//
-// NOTE: Part of the ContractResolver interface.
-func (h *htlcTimeoutResolver) Decode(r io.Reader) error {
-	// First, we'll read out all the mandatory fields of the
-	// OutgoingHtlcResolution that we store.
-	if err := decodeOutgoingResolution(r, &h.htlcResolution); err != nil {
-		return err
 	}
 
-	// With those fields read, we can now read back the fields that are
-	// specific to the resolver itself.
-	if err := binary.Read(r, endian, &h.outputIncubating); err != nil {
-		return err
-	}
-	if err := binary.Read(r, endian, &h.resolved); err != nil {
-		return err
-	}
-	if err := binary.Read(r, endian, &h.broadcastHeight); err != nil {
-		return err
-	}
+	// Evaluate current height in case cltv has already expired.
+	evaluateHeight()
 
-	if err := binary.Read(r, endian, &h.htlcIndex); err != nil {
-		return err
+	// If we reach this point, then we can't fully act yet, so we'll await
+	// either of our signals triggering: the HTLC expires, or we learn of
+	// the preimage.
+	blockEpochs, err := h.Notifier.RegisterBlockEpochNtfn(nil)
+	if err != nil {
+		return nil, err
 	}
+	defer blockEpochs.Cancel()
 
-	return nil
-}
+loop:
+	for {
+		select {
 
-// AttachResolverKit should be called once a resolved is successfully decoded
-// from its stored format. This struct delivers a generic tool kit that
-// resolvers need to complete their duty.
-//
-// NOTE: Part of the ContractResolver interface.
-func (h *htlcTimeoutResolver) AttachResolverKit(r ResolverKit) {
-	h.ResolverKit = r
-}
-
-// A compile time assertion to ensure htlcTimeoutResolver meets the
-// ContractResolver interface.
-var _ ContractResolver = (*htlcTimeoutResolver)(nil)
-
-// htlcSuccessResolver is a resolver that's capable of sweeping an incoming
-// HTLC output on-chain. If this is the remote party's commitment, we'll sweep
-// it directly from the commitment output *immediately*. If this is our
-// commitment, we'll first broadcast the success transaction, then send it to
-// the incubator for sweeping. That's it, no need to send any clean up
-// messages.
-//
-// TODO(roasbeef): don't need to broadcast?
-type htlcSuccessResolver struct {
-	// htlcResolution is the incoming HTLC resolution for this HTLC. It
-	// contains everything we need to properly resolve this HTLC.
-	htlcResolution lnwallet.IncomingHtlcResolution
-
-	// outputIncubating returns true if we've sent the output to the output
-	// incubator (utxo nursery).
-	outputIncubating bool
-
-	// resolved reflects if the contract has been fully resolved or not.
-	resolved bool
-
-	// broadcastHeight is the height that the original contract was
-	// broadcast to the main-chain at. We'll use this value to bound any
-	// historical queries to the chain for spends/confirmations.
-	broadcastHeight uint32
-
-	// payHash is the payment hash of the original HTLC extended to us.
-	payHash [32]byte
-
-	// sweepTx will be non-nil if we've already crafted a transaction to
-	// sweep a direct HTLC output. This is only a concern if we're sweeping
-	// from the commitment transaction of the remote party.
-	//
-	// TODO(roasbeef): send off to utxobundler
-	sweepTx *wire.MsgTx
-
-	ResolverKit
-}
-
-// ResolverKey returns an identifier which should be globally unique for this
-// particular resolver within the chain the original contract resides within.
-//
-// NOTE: Part of the ContractResolver interface.
-func (h *htlcSuccessResolver) ResolverKey() []byte {
-	// The primary key for this resolver will be the outpoint of the HTLC
-	// on the commitment transaction itself. If this is our commitment,
-	// then the output can be found within the signed success tx,
-	// otherwise, it's just the ClaimOutpoint.
-	var op wire.OutPoint
-	if h.htlcResolution.SignedSuccessTx != nil {
-		op = h.htlcResolution.SignedSuccessTx.TxIn[0].PreviousOutPoint
-	} else {
-		op = h.htlcResolution.ClaimOutpoint
-	}
-
-	key := newResolverID(op)
-	return key[:]
-}
-
-// Resolve attempts to resolve an unresolved incoming HTLC that we know the
-// preimage to. If the HTLC is on the commitment of the remote party, then
-// we'll simply sweep it directly. Otherwise, we'll hand this off to the utxo
-// nursery to do its duty.
-//
-// TODO(roasbeef): create multi to batch
-//
-// NOTE: Part of the ContractResolver interface.
-func (h *htlcSuccessResolver) Resolve() (ContractResolver, error) {
-	// If we're already resolved, then we can exit early.
-	if h.resolved {
-		return nil, nil
-	}
-
-	// If we don't have a success transaction, then this means that this is
-	// an output on the remote party's commitment transaction.
-	if h.htlcResolution.SignedSuccessTx == nil {
-		// If we don't already have the sweep transaction constructed,
-		// we'll do so and broadcast it.
-		if h.sweepTx == nil {
-			log.Infof("%T(%x): crafting sweep tx for "+
-				"incoming+remote htlc confirmed", h,
-				h.payHash[:])
-
-			// In this case, we can sweep it directly from the
-			// commitment output. We'll first grab a fresh address
-			// from the wallet to sweep the output.
-			addr, err := h.NewSweepAddr()
-			if err != nil {
-				return nil, err
-			}
-
-			// With our address obtained, we'll query for an
-			// estimate to be confirmed at ease.
-			//
-			// TODO(roasbeef): signal up if fee would be too large
-			// to sweep singly, need to batch
-			feePerKw, err := h.FeeEstimator.EstimateFeePerKW(6)
-			if err != nil {
-				return nil, err
+		// A new block has arrived, we'll check to see if this leads to
+		// HTLC expiration.
+		case newBlock, ok := <-blockEpochs.Epochs:
+			if !ok {
+				return nil, fmt.Errorf("quitting")
 			}
 
-			log.Debugf("%T(%x): using %v sat/kw to sweep htlc"+
-				"incoming+remote htlc confirmed", h,
-				h.payHash[:], int64(feePerKw))
+			// If this new height expires the HTLC, then we can
+			// exit early and create a resolver that's capable of
+			// handling the time locked output.
+			currentHeight = newBlock.Height
+			evaluateHeight()
 
-			// Using a weight estimator, we'll compute the total
-			// fee required, and from that the value we'll end up
-			// with.
-			totalWeight := (&lnwallet.TxWeightEstimator{}).
-				AddWitnessInput(lnwallet.OfferedHtlcSuccessWitnessSize).
-				AddP2WKHOutput().Weight()
-			totalFees := feePerKw.FeeForWeight(int64(totalWeight))
-			sweepAmt := h.htlcResolution.SweepSignDesc.Output.Value -
-				int64(totalFees)
+		// The output has been spent! This means the preimage has been
+		// revealed on-chain.
+		case commitSpend, ok := <-spendNtfn.Spend:
+			if !ok {
+				return nil, fmt.Errorf("quitting")
+			}
 
-			// With the fee computation finished, we'll now
-			// construct the sweep transaction.
-			htlcPoint := h.htlcResolution.ClaimOutpoint
-			h.sweepTx = wire.NewMsgTx(2)
-			h.sweepTx.AddTxIn(&wire.TxIn{
-				PreviousOutPoint: htlcPoint,
-			})
-			h.sweepTx.AddTxOut(&wire.TxOut{
-				PkScript: addr,
-				Value:    sweepAmt,
-			})
+			// The only way this output can be spent by the remote
+			// party is by revealing the preimage. So we'll perform
+			// our duties to clean up the contract once it has been
+			// claimed.
+			var isRemoteSpend bool
+			if h.htlcResolution.SignedTimeoutTx != nil {
+				timeoutTxHash := h.htlcResolution.SignedTimeoutTx.TxHash()
 
-			// With the transaction fully assembled, we can now
-			// generate a valid witness for the transaction.
-			h.htlcResolution.SweepSignDesc.SigHashes = txscript.NewTxSigHashes(
-				h.sweepTx,
-			)
-			h.sweepTx.TxIn[0].Witness, err = lnwallet.SenderHtlcSpendRedeem(
-				h.Signer, &h.htlcResolution.SweepSignDesc, h.sweepTx,
-				h.htlcResolution.Preimage[:],
-			)
-			if err != nil {
-				return nil, err
+				// If spend not by our own timeout tx, it must
+				// be a remote spend.
+				isRemoteSpend = !bytes.Equal(commitSpend.SpenderTxHash[:],
+					timeoutTxHash[:])
+			} else {
+				// Detect remote success tx by witness length.
+				isRemoteSpend = len(commitSpend.SpendingTx.TxIn[0].Witness) == 5
 			}
 
-			log.Infof("%T(%x): crafted sweep tx=%v", h,
-				h.payHash[:], spew.Sdump(h.sweepTx))
-
-			// With the sweep transaction confirmed, we'll now
-			// Checkpoint our state.
-			if err := h.Checkpoint(h); err != nil {
-				log.Errorf("unable to Checkpoint: %v", err)
+			if isRemoteSpend {
+				return claimCleanUp(commitSpend)
 			}
 
-			// Finally, we'll broadcast the sweep transaction to
-			// the network.
-			//
-			// TODO(roasbeef): validate first?
-			if err := h.PublishTx(h.sweepTx); err != nil {
-				log.Infof("%T(%x): unable to publish tx: %v",
-					h, h.payHash[:], err)
+			// At this point, the second-level transaction is sufficiently
+			// confirmed, or a transaction directly spending the output is.
+			// Therefore, we can now send back our clean up message.
+			failureMsg := &lnwire.FailPermanentChannelFailure{}
+			if err := h.DeliverResolutionMsg(ResolutionMsg{
+				SourceChan: h.ShortChanID,
+				HtlcIndex:  h.htlcIndex,
+				Failure:    failureMsg,
+			}); err != nil {
 				return nil, err
 			}
-		}
 
-		// With the sweep transaction broadcast, we'll wait for its
-		// confirmation.
-		sweepTXID := h.sweepTx.TxHash()
-		sweepScript := h.sweepTx.TxOut[0].PkScript
-		confNtfn, err := h.Notifier.RegisterConfirmationsNtfn(
-			&sweepTXID, sweepScript, 1, h.broadcastHeight,
-		)
-		if err != nil {
-			return nil, err
-		}
+			if state == secondLevelConf {
+				state = csvWait
+				secondLevelExpiry = currentHeight +
+					int32(h.htlcResolution.CsvDelay)
 
-		log.Infof("%T(%x): waiting for sweep tx (txid=%v) to be "+
-			"confirmed", h, h.payHash[:], sweepTXID)
+				evaluateHeight()
+			} else {
+				break loop
+			}
 
-		select {
-		case _, ok := <-confNtfn.Confirmed:
+		case _, ok := <-secondLevelSpendChan:
 			if !ok {
 				return nil, fmt.Errorf("quitting")
 			}
-
-		case <-h.Quit:
-			return nil, fmt.Errorf("quitting")
-		}
-
-		// Once the transaction has received a sufficient number of
-		// confirmations, we'll mark ourselves as fully resolved and exit.
-		h.resolved = true
-		return nil, h.Checkpoint(h)
-	}
-
-	log.Infof("%T(%x): broadcasting second-layer transition tx: %v",
-		h, h.payHash[:], spew.Sdump(h.htlcResolution.SignedSuccessTx))
-
-	// We'll now broadcast the second layer transaction so we can kick off
-	// the claiming process.
-	//
-	// TODO(roasbeef): after changing sighashes send to tx bundler
-	if err := h.PublishTx(h.htlcResolution.SignedSuccessTx); err != nil {
-		return nil, err
-	}
-
-	// Otherwise, this is an output on our commitment transaction. In this
-	// case, we'll send it to the incubator, but only if we haven't already
-	// done so.
-	if !h.outputIncubating {
-		log.Infof("%T(%x): incubating incoming htlc output",
-			h, h.payHash[:])
-
-		err := h.IncubateOutputs(h.ChanPoint, nil, nil, &h.htlcResolution)
-		if err != nil {
-			return nil, err
-		}
-
-		h.outputIncubating = true
-
-		if err := h.Checkpoint(h); err != nil {
-			log.Errorf("unable to Checkpoint: %v", err)
-		}
-	}
-
-	// 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(
-		&h.htlcResolution.ClaimOutpoint,
-		h.htlcResolution.SweepSignDesc.Output.PkScript,
-		h.broadcastHeight,
-	)
-	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.payHash[:], h.htlcResolution.CsvDelay)
-
-	select {
-	case _, ok := <-spendNtfn.Spend:
-		if !ok {
-			return nil, fmt.Errorf("quitting")
+			break loop
+
+		case <-h.Quit:
+			return nil, fmt.Errorf("resolver cancelled")
 		}
 
-	case <-h.Quit:
-		return nil, fmt.Errorf("quitting")
 	}
-
 	h.resolved = true
 	return nil, h.Checkpoint(h)
 }
@@ -612,7 +403,7 @@ func (h *htlcSuccessResolver) Resolve() (ContractResolver, error) {
 // suspend.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcSuccessResolver) Stop() {
+func (h *htlcTimeoutResolver) Stop() {
 	close(h.Quit)
 }
 
@@ -620,7 +411,7 @@ func (h *htlcSuccessResolver) Stop() {
 // resolved. In this case the target output can be forgotten.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcSuccessResolver) IsResolved() bool {
+func (h *htlcTimeoutResolver) IsResolved() bool {
 	return h.resolved
 }
 
@@ -628,14 +419,15 @@ func (h *htlcSuccessResolver) IsResolved() bool {
 // Writer.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcSuccessResolver) Encode(w io.Writer) error {
-	// First we'll encode our inner HTLC resolution.
-	if err := encodeIncomingResolution(w, &h.htlcResolution); err != nil {
+func (h *htlcTimeoutResolver) Encode(w io.Writer) error {
+	// First, we'll write out the relevant fields of the
+	// OutgoingHtlcResolution to the writer.
+	if err := encodeOutgoingResolution(w, &h.htlcResolution); err != nil {
 		return err
 	}
 
-	// Next, we'll write out the fields that are specified to the contract
-	// resolver.
+	// With that portion written, we can now write out the fields specific
+	// to the resolver itself.
 	if err := binary.Write(w, endian, h.outputIncubating); err != nil {
 		return err
 	}
@@ -645,7 +437,8 @@ func (h *htlcSuccessResolver) Encode(w io.Writer) error {
 	if err := binary.Write(w, endian, h.broadcastHeight); err != nil {
 		return err
 	}
-	if _, err := w.Write(h.payHash[:]); err != nil {
+
+	if err := binary.Write(w, endian, h.htlcIndex); err != nil {
 		return err
 	}
 
@@ -656,14 +449,15 @@ func (h *htlcSuccessResolver) Encode(w io.Writer) error {
 // instance, returning an active ContractResolver instance.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcSuccessResolver) Decode(r io.Reader) error {
-	// First we'll decode our inner HTLC resolution.
-	if err := decodeIncomingResolution(r, &h.htlcResolution); err != nil {
+func (h *htlcTimeoutResolver) Decode(r io.Reader) error {
+	// First, we'll read out all the mandatory fields of the
+	// OutgoingHtlcResolution that we store.
+	if err := decodeOutgoingResolution(r, &h.htlcResolution); err != nil {
 		return err
 	}
 
-	// Next, we'll read all the fields that are specified to the contract
-	// resolver.
+	// With those fields read, we can now read back the fields that are
+	// specific to the resolver itself.
 	if err := binary.Read(r, endian, &h.outputIncubating); err != nil {
 		return err
 	}
@@ -673,7 +467,8 @@ func (h *htlcSuccessResolver) Decode(r io.Reader) error {
 	if err := binary.Read(r, endian, &h.broadcastHeight); err != nil {
 		return err
 	}
-	if _, err := io.ReadFull(r, h.payHash[:]); err != nil {
+
+	if err := binary.Read(r, endian, &h.htlcIndex); err != nil {
 		return err
 	}
 
@@ -685,286 +480,257 @@ func (h *htlcSuccessResolver) Decode(r io.Reader) error {
 // resolvers need to complete their duty.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcSuccessResolver) AttachResolverKit(r ResolverKit) {
+func (h *htlcTimeoutResolver) AttachResolverKit(r ResolverKit) {
 	h.ResolverKit = r
 }
 
-// A compile time assertion to ensure htlcSuccessResolver meets the
+// A compile time assertion to ensure htlcTimeoutResolver meets the
 // ContractResolver interface.
-var _ ContractResolver = (*htlcSuccessResolver)(nil)
-
-// htlcOutgoingContestResolver is a ContractResolver that's able to resolve an
-// outgoing HTLC that is still contested. An HTLC is still contested, if at the
-// time that we broadcast the commitment transaction, it isn't able to be fully
-// resolved. In this case, we'll either wait for the HTLC to timeout, or for
-// us to learn of the preimage.
-type htlcOutgoingContestResolver struct {
-	// htlcTimeoutResolver is the inner solver that this resolver may turn
-	// into. This only happens if the HTLC expires on-chain.
-	htlcTimeoutResolver
-}
-
-type outgoingState uint8
-
-const (
-	cltvWait uint8 = iota
-	secondLevelConf
-	csvWait
-	sweepWait
-)
+var _ ContractResolver = (*htlcTimeoutResolver)(nil)
 
-// Resolve commences the resolution of this contract. As this contract hasn't
-// yet timed out, we'll wait for one of two things to happen
-//
-//   1. The HTLC expires. In this case, we'll sweep the funds and send a clean
-//      up cancel message to outside sub-systems.
-//
-//   2. The remote party sweeps this HTLC on-chain, in which case we'll add the
-//      pre-image to our global cache, then send a clean up settle message
-//      backwards.
+// htlcSuccessResolver is a resolver that's capable of sweeping an incoming
+// HTLC output on-chain. If this is the remote party's commitment, we'll sweep
+// it directly from the commitment output *immediately*. If this is our
+// commitment, we'll first broadcast the success transaction, then send it to
+// the incubator for sweeping. That's it, no need to send any clean up
+// messages.
 //
-// When either of these two things happens, we'll create a new resolver which
-// is able to handle the final resolution of the contract. We're only the pivot
-// point.
-func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
-	// If we're already full resolved, then we don't have anything further
-	// to do.
-	if h.resolved {
-		return nil, nil
-	}
-
-	// claimCleanUp is a helper function that's called once the HTLC output
-	// is spent by the remote party. It'll extract the preimage, add it to
-	// the global cache, and finally send the appropriate clean up message.
-	claimCleanUp := func(commitSpend *chainntnfs.SpendDetail) (ContractResolver, error) {
-		// Depending on if this is our commitment or not, then we'll be
-		// looking for a different witness pattern.
-		spenderIndex := commitSpend.SpenderInputIndex
-		spendingInput := commitSpend.SpendingTx.TxIn[spenderIndex]
-
-		log.Infof("%T(%v): extracting preimage! remote party spent "+
-			"HTLC with tx=%v", h, h.htlcResolution.ClaimOutpoint,
-			spew.Sdump(commitSpend.SpendingTx))
+// TODO(roasbeef): don't need to broadcast?
+type htlcSuccessResolver struct {
+	// htlcResolution is the incoming HTLC resolution for this HTLC. It
+	// contains everything we need to properly resolve this HTLC.
+	htlcResolution lnwallet.IncomingHtlcResolution
 
-		// If this is the remote party's commitment, then we'll be
-		// looking for them to spend using the second-level success
-		// transaction.
-		var preimage [32]byte
-		if h.htlcResolution.SignedTimeoutTx == nil {
-			// The witness stack when the remote party sweeps the
-			// output to them looks like:
-			//
-			//  * <sender sig> <recvr sig> <preimage> <witness script>
-			copy(preimage[:], spendingInput.Witness[3])
-		} else {
-			// Otherwise, they'll be spending directly from our
-			// commitment output. In which case the witness stack
-			// looks like:
-			//
-			//  * <sig> <preimage> <witness script>
-			copy(preimage[:], spendingInput.Witness[1])
-		}
+	// outputIncubating returns true if we've sent the output to the output
+	// incubator (utxo nursery).
+	outputIncubating bool
 
-		log.Infof("%T(%v): extracting preimage=%x from on-chain "+
-			"spend!", h, h.htlcResolution.ClaimOutpoint, preimage[:])
+	// resolved reflects if the contract has been fully resolved or not.
+	resolved bool
 
-		// With the preimage obtained, we can now add it to the global
-		// cache.
-		if err := h.PreimageDB.AddPreimage(preimage[:]); err != nil {
-			log.Errorf("%T(%v): unable to add witness to cache",
-				h, h.htlcResolution.ClaimOutpoint)
-		}
+	// broadcastHeight is the height that the original contract was
+	// broadcast to the main-chain at. We'll use this value to bound any
+	// historical queries to the chain for spends/confirmations.
+	broadcastHeight uint32
 
-		// Finally, we'll send the clean up message, mark ourselves as
-		// resolved, then exit.
-		if err := h.DeliverResolutionMsg(ResolutionMsg{
-			SourceChan: h.ShortChanID,
-			HtlcIndex:  h.htlcIndex,
-			PreImage:   &preimage,
-		}); err != nil {
-			return nil, err
-		}
-		h.resolved = true
-		return nil, h.Checkpoint(h)
-	}
+	// payHash is the payment hash of the original HTLC extended to us.
+	payHash [32]byte
 
-	// Otherwise, we'll watch for two external signals to decide if we'll
-	// morph into another resolver, or fully resolve the contract.
+	// sweepTx will be non-nil if we've already crafted a transaction to
+	// sweep a direct HTLC output. This is only a concern if we're sweeping
+	// from the commitment transaction of the remote party.
+	//
+	// TODO(roasbeef): send off to utxobundler
+	sweepTx *wire.MsgTx
 
-	// The output we'll be watching for is the *direct* spend from the HTLC
-	// output. If this isn't our commitment transaction, it'll be right on
-	// the resolution. Otherwise, we fetch this pointer from the input of
-	// the time out transaction.
-	var (
-		outPointToWatch      wire.OutPoint
-		scriptToWatch        []byte
-		err                  error
-		secondLevelSpendChan <-chan *chainntnfs.SpendDetail
-	)
-	if h.htlcResolution.SignedTimeoutTx == nil {
-		outPointToWatch = h.htlcResolution.ClaimOutpoint
-		scriptToWatch = h.htlcResolution.SweepSignDesc.Output.PkScript
+	ResolverKit
+}
 
-		// Create dummy channel that will never receive a notification.
-		secondLevelSpendChan = make(chan *chainntnfs.SpendDetail)
+// ResolverKey returns an identifier which should be globally unique for this
+// particular resolver within the chain the original contract resides within.
+//
+// NOTE: Part of the ContractResolver interface.
+func (h *htlcSuccessResolver) ResolverKey() []byte {
+	// The primary key for this resolver will be the outpoint of the HTLC
+	// on the commitment transaction itself. If this is our commitment,
+	// then the output can be found within the signed success tx,
+	// otherwise, it's just the ClaimOutpoint.
+	var op wire.OutPoint
+	if h.htlcResolution.SignedSuccessTx != nil {
+		op = h.htlcResolution.SignedSuccessTx.TxIn[0].PreviousOutPoint
 	} else {
-		// If this is the remote party's commitment, then we'll need to
-		// grab watch the output that our timeout transaction points
-		// to. We can directly grab the outpoint, then also extract the
-		// witness script (the last element of the witness stack) to
-		// re-construct the pkScipt we need to watch.
-		outPointToWatch = h.htlcResolution.SignedTimeoutTx.TxIn[0].PreviousOutPoint
-		witness := h.htlcResolution.SignedTimeoutTx.TxIn[0].Witness
-		scriptToWatch, err = lnwallet.WitnessScriptHash(
-			witness[len(witness)-1],
-		)
-		if err != nil {
-			return nil, err
-		}
-
-		secondLevelOutPointToWatch := h.htlcResolution.ClaimOutpoint
-		secondLevelScriptToWatch := h.htlcResolution.SweepSignDesc.Output.PkScript
-		secondLevelSpendNtfn, err := h.Notifier.RegisterSpendNtfn(
-			&secondLevelOutPointToWatch, secondLevelScriptToWatch, h.broadcastHeight,
-		)
-		if err != nil {
-			return nil, err
-		}
-		secondLevelSpendChan = secondLevelSpendNtfn.Spend
+		op = h.htlcResolution.ClaimOutpoint
 	}
 
-	// First, we'll register for a spend notification for this output. If
-	// the remote party sweeps with the pre-image, we'll be notified.
-	spendNtfn, err := h.Notifier.RegisterSpendNtfn(
-		&outPointToWatch, scriptToWatch, h.broadcastHeight,
-	)
-	if err != nil {
-		return nil, err
+	key := newResolverID(op)
+	return key[:]
+}
+
+// Resolve attempts to resolve an unresolved incoming HTLC that we know the
+// preimage to. If the HTLC is on the commitment of the remote party, then
+// we'll simply sweep it directly. Otherwise, we'll hand this off to the utxo
+// nursery to do its duty.
+//
+// TODO(roasbeef): create multi to batch
+//
+// NOTE: Part of the ContractResolver interface.
+func (h *htlcSuccessResolver) Resolve() (ContractResolver, error) {
+	// If we're already resolved, then we can exit early.
+	if h.resolved {
+		return nil, nil
 	}
 
-	// TODO: Cancel spend notifications on quit?
+	// If we don't have a success transaction, then this means that this is
+	// an output on the remote party's commitment transaction.
+	if h.htlcResolution.SignedSuccessTx == nil {
+		// If we don't already have the sweep transaction constructed,
+		// we'll do so and broadcast it.
+		if h.sweepTx == nil {
+			log.Infof("%T(%x): crafting sweep tx for "+
+				"incoming+remote htlc confirmed", h,
+				h.payHash[:])
 
-	_, currentHeight, err := h.ChainIO.GetBestBlock()
-	if err != nil {
-		return nil, err
-	}
+			// In this case, we can sweep it directly from the
+			// commitment output. We'll first grab a fresh address
+			// from the wallet to sweep the output.
+			addr, err := h.NewSweepAddr()
+			if err != nil {
+				return nil, err
+			}
 
-	state := cltvWait
-	var secondLevelExpiry int32
+			// With our address obtained, we'll query for an
+			// estimate to be confirmed at ease.
+			//
+			// TODO(roasbeef): signal up if fee would be too large
+			// to sweep singly, need to batch
+			feePerKw, err := h.FeeEstimator.EstimateFeePerKW(6)
+			if err != nil {
+				return nil, err
+			}
 
-	evaluateHeight := func() {
-		switch {
-		case state == cltvWait &&
-			uint32(currentHeight) >= h.htlcResolution.Expiry-1:
+			log.Debugf("%T(%x): using %v sat/kw to sweep htlc"+
+				"incoming+remote htlc confirmed", h,
+				h.payHash[:], int64(feePerKw))
 
-			if h.htlcResolution.SignedTimeoutTx == nil {
-				// Start sweep
+			// Using a weight estimator, we'll compute the total
+			// fee required, and from that the value we'll end up
+			// with.
+			totalWeight := (&lnwallet.TxWeightEstimator{}).
+				AddWitnessInput(lnwallet.OfferedHtlcSuccessWitnessSize).
+				AddP2WKHOutput().Weight()
+			totalFees := feePerKw.FeeForWeight(int64(totalWeight))
+			sweepAmt := h.htlcResolution.SweepSignDesc.Output.Value -
+				int64(totalFees)
 
-				state = sweepWait
-			} else {
-				// Publish timeout tx
+			// With the fee computation finished, we'll now
+			// construct the sweep transaction.
+			htlcPoint := h.htlcResolution.ClaimOutpoint
+			h.sweepTx = wire.NewMsgTx(2)
+			h.sweepTx.AddTxIn(&wire.TxIn{
+				PreviousOutPoint: htlcPoint,
+			})
+			h.sweepTx.AddTxOut(&wire.TxOut{
+				PkScript: addr,
+				Value:    sweepAmt,
+			})
 
-				state = secondLevelConf
+			// With the transaction fully assembled, we can now
+			// generate a valid witness for the transaction.
+			h.htlcResolution.SweepSignDesc.SigHashes = txscript.NewTxSigHashes(
+				h.sweepTx,
+			)
+			h.sweepTx.TxIn[0].Witness, err = lnwallet.SenderHtlcSpendRedeem(
+				h.Signer, &h.htlcResolution.SweepSignDesc, h.sweepTx,
+				h.htlcResolution.Preimage[:],
+			)
+			if err != nil {
+				return nil, err
 			}
 
-		case state == csvWait && currentHeight >= secondLevelExpiry-1:
-			// Start sweep
+			log.Infof("%T(%x): crafted sweep tx=%v", h,
+				h.payHash[:], spew.Sdump(h.sweepTx))
 
-			state = sweepWait
-		}
+			// With the sweep transaction confirmed, we'll now
+			// Checkpoint our state.
+			if err := h.Checkpoint(h); err != nil {
+				log.Errorf("unable to Checkpoint: %v", err)
+			}
 
-	}
+			// Finally, we'll broadcast the sweep transaction to
+			// the network.
+			//
+			// TODO(roasbeef): validate first?
+			if err := h.PublishTx(h.sweepTx); err != nil {
+				log.Infof("%T(%x): unable to publish tx: %v",
+					h, h.payHash[:], err)
+				return nil, err
+			}
+		}
 
-	// Evaluate current height in case cltv has already expired.
-	evaluateHeight()
+		// With the sweep transaction broadcast, we'll wait for its
+		// confirmation.
+		sweepTXID := h.sweepTx.TxHash()
+		sweepScript := h.sweepTx.TxOut[0].PkScript
+		confNtfn, err := h.Notifier.RegisterConfirmationsNtfn(
+			&sweepTXID, sweepScript, 1, h.broadcastHeight,
+		)
+		if err != nil {
+			return nil, err
+		}
 
-	// If we reach this point, then we can't fully act yet, so we'll await
-	// either of our signals triggering: the HTLC expires, or we learn of
-	// the preimage.
-	blockEpochs, err := h.Notifier.RegisterBlockEpochNtfn(nil)
-	if err != nil {
-		return nil, err
-	}
-	defer blockEpochs.Cancel()
+		log.Infof("%T(%x): waiting for sweep tx (txid=%v) to be "+
+			"confirmed", h, h.payHash[:], sweepTXID)
 
-loop:
-	for {
 		select {
-
-		// A new block has arrived, we'll check to see if this leads to
-		// HTLC expiration.
-		case newBlock, ok := <-blockEpochs.Epochs:
+		case _, ok := <-confNtfn.Confirmed:
 			if !ok {
 				return nil, fmt.Errorf("quitting")
 			}
 
-			// If this new height expires the HTLC, then we can
-			// exit early and create a resolver that's capable of
-			// handling the time locked output.
-			currentHeight = newBlock.Height
-			evaluateHeight()
+		case <-h.Quit:
+			return nil, fmt.Errorf("quitting")
+		}
 
-		// The output has been spent! This means the preimage has been
-		// revealed on-chain.
-		case commitSpend, ok := <-spendNtfn.Spend:
-			if !ok {
-				return nil, fmt.Errorf("quitting")
-			}
+		// Once the transaction has received a sufficient number of
+		// confirmations, we'll mark ourselves as fully resolved and exit.
+		h.resolved = true
+		return nil, h.Checkpoint(h)
+	}
 
-			// The only way this output can be spent by the remote
-			// party is by revealing the preimage. So we'll perform
-			// our duties to clean up the contract once it has been
-			// claimed.
-			var isRemoteSpend bool
-			if h.htlcResolution.SignedTimeoutTx != nil {
-				timeoutTxHash := h.htlcResolution.SignedTimeoutTx.TxHash()
+	log.Infof("%T(%x): broadcasting second-layer transition tx: %v",
+		h, h.payHash[:], spew.Sdump(h.htlcResolution.SignedSuccessTx))
 
-				// If spend not by our own timeout tx, it must
-				// be a remote spend.
-				isRemoteSpend = !bytes.Equal(commitSpend.SpenderTxHash[:],
-					timeoutTxHash[:])
-			} else {
-				// Detect remote success tx by witness length.
-				isRemoteSpend = len(commitSpend.SpendingTx.TxIn[0].Witness) == 5
-			}
+	// We'll now broadcast the second layer transaction so we can kick off
+	// the claiming process.
+	//
+	// TODO(roasbeef): after changing sighashes send to tx bundler
+	if err := h.PublishTx(h.htlcResolution.SignedSuccessTx); err != nil {
+		return nil, err
+	}
 
-			if isRemoteSpend {
-				return claimCleanUp(commitSpend)
-			}
+	// Otherwise, this is an output on our commitment transaction. In this
+	// case, we'll send it to the incubator, but only if we haven't already
+	// done so.
+	if !h.outputIncubating {
+		log.Infof("%T(%x): incubating incoming htlc output",
+			h, h.payHash[:])
 
-			// At this point, the second-level transaction is sufficiently
-			// confirmed, or a transaction directly spending the output is.
-			// Therefore, we can now send back our clean up message.
-			failureMsg := &lnwire.FailPermanentChannelFailure{}
-			if err := h.DeliverResolutionMsg(ResolutionMsg{
-				SourceChan: h.ShortChanID,
-				HtlcIndex:  h.htlcIndex,
-				Failure:    failureMsg,
-			}); err != nil {
-				return nil, err
-			}
+		err := h.IncubateOutputs(h.ChanPoint, nil, nil, &h.htlcResolution)
+		if err != nil {
+			return nil, err
+		}
 
-			if state == secondLevelConf {
-				state = csvWait
-				secondLevelExpiry = currentHeight +
-					int32(h.htlcResolution.CsvDelay)
+		h.outputIncubating = true
 
-				evaluateHeight()
-			} else {
-				break loop
-			}
+		if err := h.Checkpoint(h); err != nil {
+			log.Errorf("unable to Checkpoint: %v", err)
+		}
+	}
 
-		case _, ok := <-secondLevelSpendChan:
-			if !ok {
-				return nil, fmt.Errorf("quitting")
-			}
-			break loop
+	// 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(
+		&h.htlcResolution.ClaimOutpoint,
+		h.htlcResolution.SweepSignDesc.Output.PkScript,
+		h.broadcastHeight,
+	)
+	if err != nil {
+		return nil, err
+	}
 
-		case <-h.Quit:
-			return nil, fmt.Errorf("resolver cancelled")
+	log.Infof("%T(%x): waiting for second-level HTLC output to be spent "+
+		"after csv_delay=%v", h, h.payHash[:], h.htlcResolution.CsvDelay)
+
+	select {
+	case _, ok := <-spendNtfn.Spend:
+		if !ok {
+			return nil, fmt.Errorf("quitting")
 		}
 
+	case <-h.Quit:
+		return nil, fmt.Errorf("quitting")
 	}
+
 	h.resolved = true
 	return nil, h.Checkpoint(h)
 }
@@ -973,7 +739,7 @@ loop:
 // suspend.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcOutgoingContestResolver) Stop() {
+func (h *htlcSuccessResolver) Stop() {
 	close(h.Quit)
 }
 
@@ -981,7 +747,7 @@ func (h *htlcOutgoingContestResolver) Stop() {
 // resolved. In this case the target output can be forgotten.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcOutgoingContestResolver) IsResolved() bool {
+func (h *htlcSuccessResolver) IsResolved() bool {
 	return h.resolved
 }
 
@@ -989,16 +755,56 @@ func (h *htlcOutgoingContestResolver) IsResolved() bool {
 // Writer.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcOutgoingContestResolver) Encode(w io.Writer) error {
-	return h.htlcTimeoutResolver.Encode(w)
+func (h *htlcSuccessResolver) Encode(w io.Writer) error {
+	// First we'll encode our inner HTLC resolution.
+	if err := encodeIncomingResolution(w, &h.htlcResolution); err != nil {
+		return err
+	}
+
+	// Next, we'll write out the fields that are specified to the contract
+	// resolver.
+	if err := binary.Write(w, endian, h.outputIncubating); err != nil {
+		return err
+	}
+	if err := binary.Write(w, endian, h.resolved); err != nil {
+		return err
+	}
+	if err := binary.Write(w, endian, h.broadcastHeight); err != nil {
+		return err
+	}
+	if _, err := w.Write(h.payHash[:]); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 // Decode attempts to decode an encoded ContractResolver from the passed Reader
 // instance, returning an active ContractResolver instance.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcOutgoingContestResolver) Decode(r io.Reader) error {
-	return h.htlcTimeoutResolver.Decode(r)
+func (h *htlcSuccessResolver) Decode(r io.Reader) error {
+	// First we'll decode our inner HTLC resolution.
+	if err := decodeIncomingResolution(r, &h.htlcResolution); err != nil {
+		return err
+	}
+
+	// Next, we'll read all the fields that are specified to the contract
+	// resolver.
+	if err := binary.Read(r, endian, &h.outputIncubating); err != nil {
+		return err
+	}
+	if err := binary.Read(r, endian, &h.resolved); err != nil {
+		return err
+	}
+	if err := binary.Read(r, endian, &h.broadcastHeight); err != nil {
+		return err
+	}
+	if _, err := io.ReadFull(r, h.payHash[:]); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 // AttachResolverKit should be called once a resolved is successfully decoded
@@ -1006,13 +812,13 @@ func (h *htlcOutgoingContestResolver) Decode(r io.Reader) error {
 // resolvers need to complete their duty.
 //
 // NOTE: Part of the ContractResolver interface.
-func (h *htlcOutgoingContestResolver) AttachResolverKit(r ResolverKit) {
+func (h *htlcSuccessResolver) AttachResolverKit(r ResolverKit) {
 	h.ResolverKit = r
 }
 
-// A compile time assertion to ensure htlcOutgoingContestResolver meets the
+// A compile time assertion to ensure htlcSuccessResolver meets the
 // ContractResolver interface.
-var _ ContractResolver = (*htlcOutgoingContestResolver)(nil)
+var _ ContractResolver = (*htlcSuccessResolver)(nil)
 
 // htlcIncomingContestResolver is a ContractResolver that's able to resolve an
 // incoming HTLC that is still contested. An HTLC is still contested, if at the

From 4bce7ce902ece54979e4b22fdf2779384edd46c2 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 19 Sep 2018 11:35:48 -0700
Subject: [PATCH 14/15] cnct: sweeper

---
 contractcourt/contract_resolvers.go | 37 ++++++++++---
 contractcourt/sweeper.go            | 86 +++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+), 7 deletions(-)
 create mode 100644 contractcourt/sweeper.go

diff --git a/contractcourt/contract_resolvers.go b/contractcourt/contract_resolvers.go
index 70856d8fac7..1c482bdf7f6 100644
--- a/contractcourt/contract_resolvers.go
+++ b/contractcourt/contract_resolvers.go
@@ -112,6 +112,8 @@ type htlcTimeoutResolver struct {
 	htlcIndex uint64
 
 	ResolverKit
+
+	sweeper Sweeper
 }
 
 // ResolverKey returns an identifier which should be globally unique for this
@@ -282,23 +284,36 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 	state := cltvWait
 	var secondLevelExpiry int32
 
+	var sweepChan chan wire.OutPoint
+
+	// If this is the remote party commit tx, we can already announce that
+	// we are going to sweep after expiry. Sweeper will hold back the sweep
+	// tx until our input is in.
+	if h.htlcResolution.SignedTimeoutTx == nil {
+		sweepChan = h.sweeper.AnnounceSweep(int32(h.htlcResolution.Expiry))
+	}
+
 	evaluateHeight := func() {
 		switch {
 		case state == cltvWait &&
 			uint32(currentHeight) >= h.htlcResolution.Expiry-1:
 
 			if h.htlcResolution.SignedTimeoutTx == nil {
-				// Start sweep
+				// Signal sweeper that this input can be picked
+				// up now.
+				sweepChan <- h.htlcResolution.ClaimOutpoint
 
 				state = sweepWait
 			} else {
-				// Publish timeout tx
+				// TODO: Publish timeout tx
 
 				state = secondLevelConf
 			}
 
 		case state == csvWait && currentHeight >= secondLevelExpiry-1:
-			// Start sweep
+			// Signal sweeper that this input can be picked now that
+			// the second level tx time lock is expired.
+			sweepChan <- h.htlcResolution.ClaimOutpoint
 
 			state = sweepWait
 		}
@@ -335,7 +350,7 @@ loop:
 			evaluateHeight()
 
 		// The output has been spent! This means the preimage has been
-		// revealed on-chain.
+		// revealed on-chain or we used the output ourselves.
 		case commitSpend, ok := <-spendNtfn.Spend:
 			if !ok {
 				return nil, fmt.Errorf("quitting")
@@ -359,12 +374,16 @@ loop:
 			}
 
 			if isRemoteSpend {
+				// Cancel sweep
+				close(sweepChan)
+
 				return claimCleanUp(commitSpend)
 			}
 
-			// At this point, the second-level transaction is sufficiently
-			// confirmed, or a transaction directly spending the output is.
-			// Therefore, we can now send back our clean up message.
+			// At this point, the second-level transaction is
+			// sufficiently confirmed, or a transaction directly
+			// spending the output is. Therefore, we can now send
+			// back our clean up message.
 			failureMsg := &lnwire.FailPermanentChannelFailure{}
 			if err := h.DeliverResolutionMsg(ResolutionMsg{
 				SourceChan: h.ShortChanID,
@@ -379,6 +398,10 @@ loop:
 				secondLevelExpiry = currentHeight +
 					int32(h.htlcResolution.CsvDelay)
 
+				// Announce that a sweep will be needed when the
+				// 2nd level tx timelock expires.
+				sweepChan = h.sweeper.AnnounceSweep(secondLevelExpiry)
+
 				evaluateHeight()
 			} else {
 				break loop
diff --git a/contractcourt/sweeper.go b/contractcourt/sweeper.go
new file mode 100644
index 00000000000..21c1b01f6a0
--- /dev/null
+++ b/contractcourt/sweeper.go
@@ -0,0 +1,86 @@
+package contractcourt
+
+import (
+	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/chainntnfs"
+)
+
+// Sweeper is responsible for sweeping outputs back into the wallet
+type Sweeper struct {
+	Notifier chainntnfs.ChainNotifier
+	quit     chan struct{}
+}
+
+// Start starts the process of constructing and publish sweep txes.
+func (s *Sweeper) Start() error {
+	s.quit = make(chan struct{})
+
+	newBlockChan, err := s.Notifier.RegisterBlockEpochNtfn(nil)
+	if err != nil {
+		return err
+	}
+
+	go s.collector(newBlockChan)
+
+	return nil
+}
+
+func (s *Sweeper) collector(newBlockChan *chainntnfs.BlockEpochEvent) {
+	defer newBlockChan.Cancel()
+
+	for {
+		select {
+		case _, ok := <-newBlockChan.Epochs:
+			// If the epoch channel has been closed, then the
+			// ChainNotifier is exiting which means the daemon is
+			// as well. Therefore, we exit early also in order to
+			// ensure the daemon shuts down gracefully, yet
+			// swiftly.
+			if !ok {
+				return
+			}
+
+			// Fetch all input channel from priority queue that have
+			// min target height <= epoch + 1
+
+			// Read outpoints from input channels. If channel is
+			// closed, ignore. This means the sweep for that output
+			// has been cancelled.
+
+			// Construct sweep tx with the outputs.
+
+			// Publish sweep tx.
+
+			// Process publish error (double spend) -> ?
+
+			// Remove from priority queue
+
+			// Log in database outputs that are currently part of an
+			// unconfirmed sweep tx. This is to prevent republishing
+			// in a new sweep tx after restart.
+
+		case <-s.quit:
+			return
+		}
+	}
+
+}
+
+// AnnounceSweep schedules a sweep for a specific minimum target height. The
+// return value is a channel that is expected to provide the outpoint to sweep
+// by the time the block height is high enough. The reason to use a blocking
+// channel is that we want to make sure we collected all inputs before
+// constructing the sweep tx. Otherwise, if block events would reach sweeper and
+// the providers of sweep inputs (like resolvers) in the wrong order, inputs
+// would be unnecessarily delayed until the next block.
+//
+// TODO: Maybe we can already pass in the outpoint at this moment?
+func (s *Sweeper) AnnounceSweep(minTargetHeight int32) chan wire.OutPoint {
+	inputChan := make(chan wire.OutPoint)
+
+	// Insert channel into priority queue keyed with minTargetHeight
+
+	// TODO: publish sweep immediately sometimes?
+
+	return inputChan
+}

From 3ff2d64f236bdba6f7176bd8ba67c96e5cba766c Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Thu, 20 Sep 2018 08:25:03 -0700
Subject: [PATCH 15/15] cnct: less state in sweeper

---
 contractcourt/contract_resolvers.go | 160 +++++++++++++---------------
 contractcourt/sweeper.go            |  19 ++++
 2 files changed, 91 insertions(+), 88 deletions(-)

diff --git a/contractcourt/contract_resolvers.go b/contractcourt/contract_resolvers.go
index 1c482bdf7f6..da971c2eaf1 100644
--- a/contractcourt/contract_resolvers.go
+++ b/contractcourt/contract_resolvers.go
@@ -228,17 +228,14 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 	// the resolution. Otherwise, we fetch this pointer from the input of
 	// the time out transaction.
 	var (
-		outPointToWatch      wire.OutPoint
-		scriptToWatch        []byte
-		err                  error
-		secondLevelSpendChan <-chan *chainntnfs.SpendDetail
+		outPointToWatch     wire.OutPoint
+		scriptToWatch       []byte
+		err                 error
+		firstLevelSpendChan <-chan *chainntnfs.SpendDetail
 	)
 	if h.htlcResolution.SignedTimeoutTx == nil {
-		outPointToWatch = h.htlcResolution.ClaimOutpoint
-		scriptToWatch = h.htlcResolution.SweepSignDesc.Output.PkScript
-
 		// Create dummy channel that will never receive a notification.
-		secondLevelSpendChan = make(chan *chainntnfs.SpendDetail)
+		firstLevelSpendChan = make(chan *chainntnfs.SpendDetail)
 	} else {
 		// If this is the remote party's commitment, then we'll need to
 		// grab watch the output that our timeout transaction points
@@ -254,24 +251,13 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 			return nil, err
 		}
 
-		secondLevelOutPointToWatch := h.htlcResolution.ClaimOutpoint
-		secondLevelScriptToWatch := h.htlcResolution.SweepSignDesc.Output.PkScript
 		secondLevelSpendNtfn, err := h.Notifier.RegisterSpendNtfn(
-			&secondLevelOutPointToWatch, secondLevelScriptToWatch, h.broadcastHeight,
+			&outPointToWatch, scriptToWatch, h.broadcastHeight,
 		)
 		if err != nil {
 			return nil, err
 		}
-		secondLevelSpendChan = secondLevelSpendNtfn.Spend
-	}
-
-	// First, we'll register for a spend notification for this output. If
-	// the remote party sweeps with the pre-image, we'll be notified.
-	spendNtfn, err := h.Notifier.RegisterSpendNtfn(
-		&outPointToWatch, scriptToWatch, h.broadcastHeight,
-	)
-	if err != nil {
-		return nil, err
+		firstLevelSpendChan = secondLevelSpendNtfn.Spend
 	}
 
 	// TODO: Cancel spend notifications on quit?
@@ -283,75 +269,90 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 
 	state := cltvWait
 	var secondLevelExpiry int32
+	sweepDoneChan := make(chan struct{})
 
-	var sweepChan chan wire.OutPoint
+	evaluateSweep := func(sweeperCall SweeperCall) bool {
+		height := sweeperCall.TargetHeight
 
-	// If this is the remote party commit tx, we can already announce that
-	// we are going to sweep after expiry. Sweeper will hold back the sweep
-	// tx until our input is in.
-	if h.htlcResolution.SignedTimeoutTx == nil {
-		sweepChan = h.sweeper.AnnounceSweep(int32(h.htlcResolution.Expiry))
-	}
-
-	evaluateHeight := func() {
 		switch {
-		case state == cltvWait &&
-			uint32(currentHeight) >= h.htlcResolution.Expiry-1:
-
-			if h.htlcResolution.SignedTimeoutTx == nil {
-				// Signal sweeper that this input can be picked
-				// up now.
-				sweepChan <- h.htlcResolution.ClaimOutpoint
-
-				state = sweepWait
-			} else {
-				// TODO: Publish timeout tx
-
-				state = secondLevelConf
+		case state == cltvWait:
+			// Nothing to sweep if not expired or second level
+			// required.
+			if h.htlcResolution.SignedTimeoutTx != nil {
+				return false
 			}
 
-		case state == csvWait && currentHeight >= secondLevelExpiry-1:
-			// Signal sweeper that this input can be picked now that
-			// the second level tx time lock is expired.
-			sweepChan <- h.htlcResolution.ClaimOutpoint
+			if uint32(height) < h.htlcResolution.Expiry-1 {
+				return false
+			}
+		case state == csvWait:
+			// Nothing to sweep if second level not expired yet.
+			if height < secondLevelExpiry-1 {
+				return false
+			}
+		default:
+			return false
+		}
 
-			state = sweepWait
+		// Signal sweeper that this input can be picked
+		// up now.
+		sweeperCall.InputChan <- SweepInput{
+			OutPoint:   h.htlcResolution.ClaimOutpoint,
+			ResultChan: sweepDoneChan,
 		}
+		close(sweeperCall.InputChan)
 
+		state = sweepWait
+
+		return true
 	}
 
-	// Evaluate current height in case cltv has already expired.
-	evaluateHeight()
+	// defer sweeperCallChan.Close() ?
 
-	// If we reach this point, then we can't fully act yet, so we'll await
-	// either of our signals triggering: the HTLC expires, or we learn of
-	// the preimage.
+	sweeperCallChan := h.sweeper.RegisterForCalls()
 	blockEpochs, err := h.Notifier.RegisterBlockEpochNtfn(nil)
 	if err != nil {
 		return nil, err
 	}
 	defer blockEpochs.Cancel()
 
+	// Early publish of timeout tx
+
 loop:
 	for {
 		select {
 
 		// A new block has arrived, we'll check to see if this leads to
-		// HTLC expiration.
+		// HTLC expiration and the need to publish the timeout tx.
 		case newBlock, ok := <-blockEpochs.Epochs:
 			if !ok {
 				return nil, fmt.Errorf("quitting")
 			}
 
+			if state == cltvWait &&
+				uint32(newBlock.Height) >= h.htlcResolution.Expiry-1 &&
+				h.htlcResolution.SignedTimeoutTx != nil {
+
+				// TODO: Publish timeout tx
+
+				state = secondLevelConf
+			}
+
+		// Sweeper is constructing a new sweep. Evaluate if we have
+		// anything to add.
+		case newBlock, ok := <-sweeperCallChan:
+			if !ok {
+				return nil, fmt.Errorf("quitting")
+			}
+
 			// If this new height expires the HTLC, then we can
 			// exit early and create a resolver that's capable of
 			// handling the time locked output.
-			currentHeight = newBlock.Height
-			evaluateHeight()
+			evaluateSweep(newBlock)
 
 		// The output has been spent! This means the preimage has been
 		// revealed on-chain or we used the output ourselves.
-		case commitSpend, ok := <-spendNtfn.Spend:
+		case commitSpend, ok := <-firstLevelSpendChan:
 			if !ok {
 				return nil, fmt.Errorf("quitting")
 			}
@@ -360,23 +361,14 @@ loop:
 			// party is by revealing the preimage. So we'll perform
 			// our duties to clean up the contract once it has been
 			// claimed.
-			var isRemoteSpend bool
-			if h.htlcResolution.SignedTimeoutTx != nil {
-				timeoutTxHash := h.htlcResolution.SignedTimeoutTx.TxHash()
-
-				// If spend not by our own timeout tx, it must
-				// be a remote spend.
-				isRemoteSpend = !bytes.Equal(commitSpend.SpenderTxHash[:],
-					timeoutTxHash[:])
-			} else {
-				// Detect remote success tx by witness length.
-				isRemoteSpend = len(commitSpend.SpendingTx.TxIn[0].Witness) == 5
-			}
+			timeoutTxHash := h.htlcResolution.SignedTimeoutTx.TxHash()
 
-			if isRemoteSpend {
-				// Cancel sweep
-				close(sweepChan)
+			// If spend not by our own timeout tx, it must
+			// be a remote spend.
+			isRemoteSpend := !bytes.Equal(commitSpend.SpenderTxHash[:],
+				timeoutTxHash[:])
 
+			if isRemoteSpend {
 				return claimCleanUp(commitSpend)
 			}
 
@@ -393,26 +385,18 @@ loop:
 				return nil, err
 			}
 
-			if state == secondLevelConf {
-				state = csvWait
-				secondLevelExpiry = currentHeight +
-					int32(h.htlcResolution.CsvDelay)
+			state = csvWait
+			secondLevelExpiry = currentHeight +
+				int32(h.htlcResolution.CsvDelay)
+		case <-sweepDoneChan:
+			// Check for remote spend
 
-				// Announce that a sweep will be needed when the
-				// 2nd level tx timelock expires.
-				sweepChan = h.sweeper.AnnounceSweep(secondLevelExpiry)
+			// Detect remote success tx by witness length.
+			// isRemoteSpend = len(commitSpend.SpendingTx.TxIn[0].Witness) == 5
 
-				evaluateHeight()
-			} else {
-				break loop
-			}
+			// Deliver failure message
 
-		case _, ok := <-secondLevelSpendChan:
-			if !ok {
-				return nil, fmt.Errorf("quitting")
-			}
 			break loop
-
 		case <-h.Quit:
 			return nil, fmt.Errorf("resolver cancelled")
 		}
diff --git a/contractcourt/sweeper.go b/contractcourt/sweeper.go
index 21c1b01f6a0..c67578add02 100644
--- a/contractcourt/sweeper.go
+++ b/contractcourt/sweeper.go
@@ -25,6 +25,25 @@ func (s *Sweeper) Start() error {
 	return nil
 }
 
+// SweepInput is what is pushed to the sweeper to describe the input to sweep.
+type SweepInput struct {
+	OutPoint   wire.OutPoint
+	ResultChan chan struct{}
+}
+
+// SweeperCall is push in the channel to collect sweep inputs.
+type SweeperCall struct {
+	InputChan    chan SweepInput
+	TargetHeight int32
+}
+
+// RegisterForCalls returns a channel on which the caller can listen for sweeper
+// events.
+func (s *Sweeper) RegisterForCalls() chan SweeperCall {
+
+	return make(chan SweeperCall)
+}
+
 func (s *Sweeper) collector(newBlockChan *chainntnfs.BlockEpochEvent) {
 	defer newBlockChan.Cancel()