Merge pull request #2501 from cfromknecht/batch-preimage-writes

htlcswitch: batch preimage writes/consistency fix

htlcswitch/link.go

@@ -333,6 +333,12 @@ type channelLink struct {
 	// commitment fee every time it fires.
 	updateFeeTimer *time.Timer
 
+	// uncommittedPreimages stores a list of all preimages that have been
+	// learned since receiving the last CommitSig from the remote peer. The
+	// batch will be flushed just before accepting the subsequent CommitSig
+	// or on shutdown to avoid doing a write for each preimage received.
+	uncommittedPreimages []lntypes.Preimage
+
 	sync.RWMutex
 
 	wg   sync.WaitGroup
@@ -449,6 +455,18 @@ func (l *channelLink) Stop() {
 
 	close(l.quit)
 	l.wg.Wait()
+
+	// As a final precaution, we will attempt to flush any uncommitted
+	// preimages to the preimage cache. The preimages should be re-delivered
+	// after channel reestablishment, however this adds an extra layer of
+	// protection in case the peer never returns. Without this, we will be
+	// unable to settle any contracts depending on the preimages even though
+	// we had learned them at some point.
+	err := l.cfg.PreimageCache.AddPreimages(l.uncommittedPreimages...)
+	if err != nil {
+		log.Errorf("Unable to add preimages=%v to cache: %v",
+			l.uncommittedPreimages, err)
+	}
 }
 
 // WaitForShutdown blocks until the link finishes shutting down, which includes
@@ -1412,17 +1430,11 @@ func (l *channelLink) handleUpstreamMsg(msg lnwire.Message) {
 
 		// TODO(roasbeef): pipeline to switch
 
-		// As we've learned of a new preimage for the first time, we'll
-		// add it to our preimage cache. By doing this, we ensure
-		// any contested contracts watched by any on-chain arbitrators
-		// can now sweep this HTLC on-chain.
-		go func() {
-			err := l.cfg.PreimageCache.AddPreimage(pre[:])
-			if err != nil {
-				l.errorf("unable to add preimage=%x to "+
-					"cache", pre[:])
-			}
-		}()
+		// Add the newly discovered preimage to our growing list of
+		// uncommitted preimage. These will be written to the witness
+		// cache just before accepting the next commitment signature
+		// from the remote peer.
+		l.uncommittedPreimages = append(l.uncommittedPreimages, pre)
 
 	case *lnwire.UpdateFailMalformedHTLC:
 		// Convert the failure type encoded within the HTLC fail
@@ -1475,10 +1487,39 @@ func (l *channelLink) handleUpstreamMsg(msg lnwire.Message) {
 		}
 
 	case *lnwire.CommitSig:
+		// Since we may have learned new preimages for the first time,
+		// we'll add them to our preimage cache. By doing this, we
+		// ensure any contested contracts watched by any on-chain
+		// arbitrators can now sweep this HTLC on-chain. We delay
+		// committing the preimages until just before accepting the new
+		// remote commitment, as afterwards the peer won't resend the
+		// Settle messages on the next channel reestablishment. Doing so
+		// allows us to more effectively batch this operation, instead
+		// of doing a single write per preimage.
+		err := l.cfg.PreimageCache.AddPreimages(
+			l.uncommittedPreimages...,
+		)
+		if err != nil {
+			l.fail(
+				LinkFailureError{code: ErrInternalError},
+				"unable to add preimages=%v to cache: %v",
+				l.uncommittedPreimages, err,
+			)
+			return
+		}
+
+		// Instead of truncating the slice to conserve memory
+		// allocations, we simply set the uncommitted preimage slice to
+		// nil so that a new one will be initialized if any more
+		// witnesses are discovered. We do this maximum size of the
+		// slice can occupy 15KB, and want to ensure we release that
+		// memory back to the runtime.
+		l.uncommittedPreimages = nil
+
 		// We just received a new updates to our local commitment
 		// chain, validate this new commitment, closing the link if
 		// invalid.
-		err := l.channel.ReceiveNewCommitment(msg.CommitSig, msg.HtlcSigs)
+		err = l.channel.ReceiveNewCommitment(msg.CommitSig, msg.HtlcSigs)
 		if err != nil {
 			// If we were unable to reconstruct their proposed
 			// commitment, then we'll examine the type of error. If

htlcswitch/link_test.go

@@ -3,6 +3,7 @@ package htlcswitch
 import (
 	"bytes"
 	"crypto/rand"
+	"crypto/sha256"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -28,6 +29,7 @@ import (
 	"github.com/lightningnetwork/lnd/htlcswitch/hodl"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lnpeer"
+	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/ticker"
@@ -1532,10 +1534,7 @@ func newSingleLinkTestHarness(chanAmt, chanReserve btcutil.Amount) (
 		invoiceRegistry = newMockRegistry(globalPolicy.TimeLockDelta)
 	)
 
-	pCache := &mockPreimageCache{
-		// hash -> preimage
-		preimageMap: make(map[[32]byte][]byte),
-	}
+	pCache := newMockPreimageCache()
 
 	aliceDb := aliceChannel.State().Db
 	aliceSwitch, err := initSwitchWithDB(testStartingHeight, aliceDb)
@@ -4042,10 +4041,7 @@ func restartLink(aliceChannel *lnwallet.LightningChannel, aliceSwitch *Switch,
 
 		invoiceRegistry = newMockRegistry(globalPolicy.TimeLockDelta)
 
-		pCache = &mockPreimageCache{
-			// hash -> preimage
-			preimageMap: make(map[[32]byte][]byte),
-		}
+		pCache = newMockPreimageCache()
 	)
 
 	aliceDb := aliceChannel.State().Db
@@ -4120,6 +4116,29 @@ func restartLink(aliceChannel *lnwallet.LightningChannel, aliceSwitch *Switch,
 // gnerateHtlc generates a simple payment from Bob to Alice.
 func generateHtlc(t *testing.T, coreLink *channelLink,
 	bobChannel *lnwallet.LightningChannel, id uint64) *lnwire.UpdateAddHTLC {
+
+	t.Helper()
+
+	htlc, invoice := generateHtlcAndInvoice(t, id)
+
+	// We must add the invoice to the registry, such that Alice
+	// expects this payment.
+	err := coreLink.cfg.Registry.(*mockInvoiceRegistry).AddInvoice(
+		*invoice)
+	if err != nil {
+		t.Fatalf("unable to add invoice to registry: %v", err)
+	}
+
+	return htlc
+}
+
+// generateHtlcAndInvoice generates an invoice and a single hop htlc to send to
+// the receiver.
+func generateHtlcAndInvoice(t *testing.T,
+	id uint64) (*lnwire.UpdateAddHTLC, *channeldb.Invoice) {
+
+	t.Helper()
+
 	htlcAmt := lnwire.NewMSatFromSatoshis(10000)
 	hops := []ForwardingInfo{
 		{
@@ -4130,27 +4149,28 @@ func generateHtlc(t *testing.T, coreLink *channelLink,
 		},
 	}
 	blob, err := generateRoute(hops...)
+	if err != nil {
+		t.Fatalf("unable to generate route: %v", err)
+	}
+
 	invoice, htlc, err := generatePayment(htlcAmt, htlcAmt, 144,
 		blob)
 	if err != nil {
 		t.Fatalf("unable to create payment: %v", err)
 	}
 
-	// We must add the invoice to the registry, such that Alice
-	// expects this payment.
-	err = coreLink.cfg.Registry.(*mockInvoiceRegistry).AddInvoice(
-		*invoice)
-	if err != nil {
-		t.Fatalf("unable to add invoice to registry: %v", err)
-	}
 	htlc.ID = id
-	return htlc
+
+	return htlc, invoice
 }
 
 // sendHtlcBobToAlice sends an HTLC from Bob to Alice, that pays to a preimage
 // already in Alice's registry.
 func sendHtlcBobToAlice(t *testing.T, aliceLink ChannelLink,
 	bobChannel *lnwallet.LightningChannel, htlc *lnwire.UpdateAddHTLC) {
+
+	t.Helper()
+
 	_, err := bobChannel.AddHTLC(htlc, nil)
 	if err != nil {
 		t.Fatalf("bob failed adding htlc: %v", err)
@@ -4159,10 +4179,70 @@ func sendHtlcBobToAlice(t *testing.T, aliceLink ChannelLink,
 	aliceLink.HandleChannelUpdate(htlc)
 }
 
+// sendHtlcAliceToBob sends an HTLC from Alice to Bob, by first committing the
+// HTLC in the circuit map, then delivering the outgoing packet to Alice's link.
+// The HTLC will be sent to Bob via Alice's message stream.
+func sendHtlcAliceToBob(t *testing.T, aliceLink ChannelLink, htlcID int,
+	htlc *lnwire.UpdateAddHTLC) {
+
+	t.Helper()
+
+	circuitMap := aliceLink.(*channelLink).cfg.Switch.circuits
+	fwdActions, err := circuitMap.CommitCircuits(
+		&PaymentCircuit{
+			Incoming: CircuitKey{
+				HtlcID: uint64(htlcID),
+			},
+			PaymentHash: htlc.PaymentHash,
+		},
+	)
+	if err != nil {
+		t.Fatalf("unable to commit circuit: %v", err)
+	}
+
+	if len(fwdActions.Adds) != 1 {
+		t.Fatalf("expected 1 adds, found %d", len(fwdActions.Adds))
+	}
+
+	aliceLink.HandleSwitchPacket(&htlcPacket{
+		incomingHTLCID: uint64(htlcID),
+		htlc:           htlc,
+	})
+
+}
+
+// receiveHtlcAliceToBob pulls the next message from Alice's message stream,
+// asserts that it is an UpdateAddHTLC, then applies it to Bob's state machine.
+func receiveHtlcAliceToBob(t *testing.T, aliceMsgs <-chan lnwire.Message,
+	bobChannel *lnwallet.LightningChannel) {
+
+	t.Helper()
+
+	var msg lnwire.Message
+	select {
+	case msg = <-aliceMsgs:
+	case <-time.After(15 * time.Second):
+		t.Fatalf("did not received htlc from alice")
+	}
+
+	htlcAdd, ok := msg.(*lnwire.UpdateAddHTLC)
+	if !ok {
+		t.Fatalf("expected UpdateAddHTLC, got %T", msg)
+	}
+
+	_, err := bobChannel.ReceiveHTLC(htlcAdd)
+	if err != nil {
+		t.Fatalf("bob failed receiving htlc: %v", err)
+	}
+}
+
 // sendCommitSigBobToAlice makes Bob sign a new commitment and send it to
 // Alice, asserting that it signs expHtlcs number of HTLCs.
 func sendCommitSigBobToAlice(t *testing.T, aliceLink ChannelLink,
 	bobChannel *lnwallet.LightningChannel, expHtlcs int) {
+
+	t.Helper()
+
 	sig, htlcSigs, err := bobChannel.SignNextCommitment()
 	if err != nil {
 		t.Fatalf("error signing commitment: %v", err)
@@ -4186,6 +4266,9 @@ func sendCommitSigBobToAlice(t *testing.T, aliceLink ChannelLink,
 func receiveRevAndAckAliceToBob(t *testing.T, aliceMsgs chan lnwire.Message,
 	aliceLink ChannelLink,
 	bobChannel *lnwallet.LightningChannel) {
+
+	t.Helper()
+
 	var msg lnwire.Message
 	select {
 	case msg = <-aliceMsgs:
@@ -4239,6 +4322,9 @@ func receiveCommitSigAliceToBob(t *testing.T, aliceMsgs chan lnwire.Message,
 // the RevokeAndAck to Alice.
 func sendRevAndAckBobToAlice(t *testing.T, aliceLink ChannelLink,
 	bobChannel *lnwallet.LightningChannel) {
+
+	t.Helper()
+
 	rev, _, err := bobChannel.RevokeCurrentCommitment()
 	if err != nil {
 		t.Fatalf("unable to revoke commitment: %v", err)
@@ -4273,6 +4359,28 @@ func receiveSettleAliceToBob(t *testing.T, aliceMsgs chan lnwire.Message,
 	}
 }
 
+// sendSettleBobToAlice settles an HTLC on Bob's state machine, then sends an
+// UpdateFulfillHTLC message to Alice's upstream inbox.
+func sendSettleBobToAlice(t *testing.T, aliceLink ChannelLink,
+	bobChannel *lnwallet.LightningChannel, htlcID uint64,
+	preimage lntypes.Preimage) {
+
+	t.Helper()
+
+	err := bobChannel.SettleHTLC(preimage, htlcID, nil, nil, nil)
+	if err != nil {
+		t.Fatalf("alice failed settling htlc id=%d hash=%x",
+			htlcID, sha256.Sum256(preimage[:]))
+	}
+
+	settle := &lnwire.UpdateFulfillHTLC{
+		ID:              htlcID,
+		PaymentPreimage: preimage,
+	}
+
+	aliceLink.HandleChannelUpdate(settle)
+}
+
 // receiveSettleAliceToBob waits for Alice to send a HTLC settle message to
 // Bob, then hands this to Bob.
 func receiveFailAliceToBob(t *testing.T, aliceMsgs chan lnwire.Message,
@@ -4389,6 +4497,26 @@ func TestChannelLinkNoMoreUpdates(t *testing.T) {
 	}
 }
 
+// checkHasPreimages inspects Alice's preimage cache, and asserts whether the
+// preimages for the provided HTLCs are known and unknown, and that all of them
+// match the expected status of expOk.
+func checkHasPreimages(t *testing.T, coreLink *channelLink,
+	htlcs []*lnwire.UpdateAddHTLC, expOk bool) {
+
+	t.Helper()
+
+	for i := range htlcs {
+		_, ok := coreLink.cfg.PreimageCache.LookupPreimage(
+			htlcs[i].PaymentHash,
+		)
+		if ok != expOk {
+			t.Fatalf("expected to find witness: %v, "+
+				"got %v for hash=%x", expOk, ok,
+				htlcs[i].PaymentHash)
+		}
+	}
+}
+
 // TestChannelLinkWaitForRevocation tests that we will keep accepting updates
 // to our commitment transaction, even when we are waiting for a revocation
 // from the remote node.
@@ -4500,6 +4628,135 @@ func TestChannelLinkWaitForRevocation(t *testing.T) {
 	}
 }
 
+// TestChannelLinkBatchPreimageWrite asserts that a link will batch preimage
+// writes when just as it receives a CommitSig to lock in any Settles, and also
+// if the link is aware of any uncommitted preimages if the link is stopped,
+// i.e. due to a disconnection or shutdown.
+func TestChannelLinkBatchPreimageWrite(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		disconnect bool
+	}{
+		{
+			name:       "flush on commit sig",
+			disconnect: false,
+		},
+		{
+			name:       "flush on disconnect",
+			disconnect: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			testChannelLinkBatchPreimageWrite(t, test.disconnect)
+		})
+	}
+}
+
+func testChannelLinkBatchPreimageWrite(t *testing.T, disconnect bool) {
+	const chanAmt = btcutil.SatoshiPerBitcoin * 5
+	const chanReserve = btcutil.SatoshiPerBitcoin * 1
+	aliceLink, bobChannel, batchTicker, startUp, cleanUp, _, err :=
+		newSingleLinkTestHarness(chanAmt, chanReserve)
+	if err != nil {
+		t.Fatalf("unable to create link: %v", err)
+	}
+	defer cleanUp()
+
+	if err := startUp(); err != nil {
+		t.Fatalf("unable to start test harness: %v", err)
+	}
+
+	var (
+		coreLink  = aliceLink.(*channelLink)
+		aliceMsgs = coreLink.cfg.Peer.(*mockPeer).sentMsgs
+	)
+
+	// We will send 10 HTLCs in total, from Bob to Alice.
+	numHtlcs := 10
+	var htlcs []*lnwire.UpdateAddHTLC
+	var invoices []*channeldb.Invoice
+	for i := 0; i < numHtlcs; i++ {
+		htlc, invoice := generateHtlcAndInvoice(t, uint64(i))
+		htlcs = append(htlcs, htlc)
+		invoices = append(invoices, invoice)
+	}
+
+	// First, send a batch of Adds from Alice to Bob.
+	for i, htlc := range htlcs {
+		sendHtlcAliceToBob(t, aliceLink, i, htlc)
+		receiveHtlcAliceToBob(t, aliceMsgs, bobChannel)
+	}
+
+	// Assert that no preimages exist for these htlcs in Alice's cache.
+	checkHasPreimages(t, coreLink, htlcs, false)
+
+	// Force alice's link to sign a commitment covering the htlcs sent thus
+	// far.
+	select {
+	case batchTicker <- time.Now():
+	case <-time.After(15 * time.Second):
+		t.Fatalf("could not force commit sig")
+	}
+
+	// Do a commitment dance to lock in the Adds, we expect numHtlcs htlcs
+	// to be on each party's commitment transactions.
+	receiveCommitSigAliceToBob(
+		t, aliceMsgs, aliceLink, bobChannel, numHtlcs,
+	)
+	sendRevAndAckBobToAlice(t, aliceLink, bobChannel)
+	sendCommitSigBobToAlice(t, aliceLink, bobChannel, numHtlcs)
+	receiveRevAndAckAliceToBob(t, aliceMsgs, aliceLink, bobChannel)
+
+	// Check again that no preimages exist for these htlcs in Alice's cache.
+	checkHasPreimages(t, coreLink, htlcs, false)
+
+	// Now, have Bob settle the HTLCs back to Alice using the preimages in
+	// the invoice corresponding to each of the HTLCs.
+	for i, invoice := range invoices {
+		sendSettleBobToAlice(
+			t, aliceLink, bobChannel, uint64(i),
+			invoice.Terms.PaymentPreimage,
+		)
+	}
+
+	// Assert that Alice has not yet written the preimages, even though she
+	// has received them in the UpdateFulfillHTLC messages.
+	checkHasPreimages(t, coreLink, htlcs, false)
+
+	// If this is the disconnect run, we will having Bob send Alice his
+	// CommitSig, and simply stop Alice's link. As she exits, we should
+	// detect that she has uncommitted preimages and write them to disk.
+	if disconnect {
+		aliceLink.Stop()
+		checkHasPreimages(t, coreLink, htlcs, true)
+		return
+	}
+
+	// Otherwise, we are testing that Alice commits the preimages after
+	// receiving a CommitSig from Bob. Bob's commitment should now have 0
+	// HTLCs.
+	sendCommitSigBobToAlice(t, aliceLink, bobChannel, 0)
+
+	// Since Alice will process the CommitSig asynchronously, we wait until
+	// she replies with her RevokeAndAck to ensure the tests reliably
+	// inspect her cache after advancing her state.
+	select {
+
+	// Received Alice's RevokeAndAck, assert that she has written all of the
+	// uncommitted preimages learned in this commitment.
+	case <-aliceMsgs:
+		checkHasPreimages(t, coreLink, htlcs, true)
+
+	// Alice didn't send her RevokeAndAck, something is wrong.
+	case <-time.After(15 * time.Second):
+		t.Fatalf("alice did not send her revocation")
+	}
+}
+
 // TestChannelLinkCleanupSpuriousResponses tests that we properly cleanup
 // references in the event that internal retransmission continues as a result of
 // not properly cleaning up Add/SettleFailRefs.

htlcswitch/mock.go

@@ -32,25 +32,32 @@ import (
 
 type mockPreimageCache struct {
 	sync.Mutex
-	preimageMap map[[32]byte][]byte
+	preimageMap map[lntypes.Hash]lntypes.Preimage
 }
 
-func (m *mockPreimageCache) LookupPreimage(hash []byte) ([]byte, bool) {
+func newMockPreimageCache() *mockPreimageCache {
+	return &mockPreimageCache{
+		preimageMap: make(map[lntypes.Hash]lntypes.Preimage),
+	}
+}
+
+func (m *mockPreimageCache) LookupPreimage(
+	hash lntypes.Hash) (lntypes.Preimage, bool) {
+
 	m.Lock()
 	defer m.Unlock()
 
-	var h [32]byte
-	copy(h[:], hash)
-
-	p, ok := m.preimageMap[h]
+	p, ok := m.preimageMap[hash]
 	return p, ok
 }
 
-func (m *mockPreimageCache) AddPreimage(preimage []byte) error {
+func (m *mockPreimageCache) AddPreimages(preimages ...lntypes.Preimage) error {
 	m.Lock()
 	defer m.Unlock()
 
-	m.preimageMap[sha256.Sum256(preimage[:])] = preimage
+	for _, preimage := range preimages {
+		m.preimageMap[preimage.Hash()] = preimage
+	}
 
 	return nil
 }

htlcswitch/test_utils.go

@@ -367,10 +367,7 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 	aliceSigner := &mockSigner{aliceKeyPriv}
 	bobSigner := &mockSigner{bobKeyPriv}
 
-	pCache := &mockPreimageCache{
-		// hash -> preimage
-		preimageMap: make(map[[32]byte][]byte),
-	}
+	pCache := newMockPreimageCache()
 
 	alicePool := lnwallet.NewSigPool(runtime.NumCPU(), aliceSigner)
 	channelAlice, err := lnwallet.NewLightningChannel(
@@ -982,10 +979,7 @@ type hopNetwork struct {
 func newHopNetwork() *hopNetwork {
 	defaultDelta := uint32(6)
 
-	pCache := &mockPreimageCache{
-		// hash -> preimage
-		preimageMap: make(map[[32]byte][]byte),
-	}
+	pCache := newMockPreimageCache()
 
 	globalPolicy := ForwardingPolicy{
 		MinHTLC:       lnwire.NewMSatFromSatoshis(5),