htlcswitch: continue threading context through

Here, we address one `context.TODO` in the `memoryMailBox` by adding a
context to the FailAdd method. The rest of the PR ensures that we thread
context to any call to FailAdd.

We add another `context.TODO` in the `interceptedForward` Fail method.

Author: ellemouton

htlcswitch/interceptable_switch.go

@@ -751,9 +751,11 @@ func (f *interceptedForward) ResumeModified(
 // Fail notifies the intention to Fail an existing hold forward with an
 // encrypted failure reason.
 func (f *interceptedForward) Fail(reason []byte) error {
+	ctx := context.TODO()
+
 	obfuscatedReason := f.packet.obfuscator.IntermediateEncrypt(reason)
 
-	return f.resolve(&lnwire.UpdateFailHTLC{
+	return f.resolve(ctx, &lnwire.UpdateFailHTLC{
 		Reason: obfuscatedReason,
 	})
 }
@@ -824,24 +826,29 @@ func (f *interceptedForward) FailWithCode(ctx context.Context,
 		return fmt.Errorf("failed to encrypt failure reason %w", err)
 	}
 
-	return f.resolve(&lnwire.UpdateFailHTLC{
+	return f.resolve(ctx, &lnwire.UpdateFailHTLC{
 		Reason: reason,
 	})
 }
 
 // Settle forwards a settled packet to the switch.
 func (f *interceptedForward) Settle(preimage lntypes.Preimage) error {
+	ctx := context.TODO()
+
 	if !preimage.Matches(f.htlc.PaymentHash) {
 		return errors.New("preimage does not match hash")
 	}
-	return f.resolve(&lnwire.UpdateFulfillHTLC{
+
+	return f.resolve(ctx, &lnwire.UpdateFulfillHTLC{
 		PaymentPreimage: preimage,
 	})
 }
 
 // resolve is used for both Settle and Fail and forwards the message to the
 // switch.
-func (f *interceptedForward) resolve(message lnwire.Message) error {
+func (f *interceptedForward) resolve(ctx context.Context,
+	message lnwire.Message) error {
+
 	pkt := &htlcPacket{
 		incomingChanID: f.packet.incomingChanID,
 		incomingHTLCID: f.packet.incomingHTLCID,
@@ -853,5 +860,8 @@ func (f *interceptedForward) resolve(message lnwire.Message) error {
 		obfuscator:     f.packet.obfuscator,
 		sourceRef:      f.packet.sourceRef,
 	}
-	return f.htlcSwitch.mailOrchestrator.Deliver(pkt.incomingChanID, pkt)
+
+	return f.htlcSwitch.mailOrchestrator.Deliver(
+		ctx, pkt.incomingChanID, pkt,
+	)
 }

htlcswitch/link.go

@@ -1695,7 +1695,7 @@ func (l *channelLink) handleDownstreamUpdateAdd(ctx context.Context,
 	// already sent Stfu, then we can't add new htlcs to the link and we
 	// need to bounce it.
 	if l.IsFlushing(Outgoing) || !l.quiescer.CanSendUpdates() {
-		l.mailBox.FailAdd(pkt)
+		l.mailBox.FailAdd(ctx, pkt)
 
 		return NewDetailedLinkError(
 			&lnwire.FailTemporaryChannelFailure{},
@@ -1718,7 +1718,7 @@ func (l *channelLink) handleDownstreamUpdateAdd(ctx context.Context,
 		l.log.Debugf("Unable to handle downstream HTLC - max fee " +
 			"exposure exceeded")
 
-		l.mailBox.FailAdd(pkt)
+		l.mailBox.FailAdd(ctx, pkt)
 
 		return NewDetailedLinkError(
 			lnwire.NewTemporaryChannelFailure(nil),
@@ -1752,7 +1752,7 @@ func (l *channelLink) handleDownstreamUpdateAdd(ctx context.Context,
 		// the switch, since the circuit was never fully opened,
 		// and the forwarding package shows it as
 		// unacknowledged.
-		l.mailBox.FailAdd(pkt)
+		l.mailBox.FailAdd(ctx, pkt)
 
 		return NewDetailedLinkError(
 			lnwire.NewTemporaryChannelFailure(nil),

htlcswitch/mailbox.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/lightningnetwork/lnd/clock"
+	"github.com/lightningnetwork/lnd/fn/v2"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/lightningnetwork/lnd/lnwire"
@@ -53,7 +54,7 @@ type MailBox interface {
 	// packet from being delivered after the link restarts if the switch has
 	// remained online. The generated LinkError will show an
 	// OutgoingFailureDownstreamHtlcAdd FailureDetail.
-	FailAdd(pkt *htlcPacket)
+	FailAdd(ctx context.Context, pkt *htlcPacket)
 
 	// MessageOutBox returns a channel that any new messages ready for
 	// delivery will be sent on.
@@ -82,7 +83,7 @@ type MailBox interface {
 
 	// Start starts the mailbox and any goroutines it needs to operate
 	// properly.
-	Start()
+	Start(ctx context.Context)
 
 	// Stop signals the mailbox and its goroutines for a graceful shutdown.
 	Stop()
@@ -148,7 +149,9 @@ type memoryMailBox struct {
 
 	wireShutdown chan struct{}
 	pktShutdown  chan struct{}
-	quit         chan struct{}
+
+	cancel fn.Option[context.CancelFunc]
+	quit   chan struct{}
 
 	// feeRate is set when the link receives or sends out fee updates. It
 	// is refreshed when AttachMailBox is called in case a fee update did
@@ -205,10 +208,13 @@ const (
 // Start starts the mailbox and any goroutines it needs to operate properly.
 //
 // NOTE: This method is part of the MailBox interface.
-func (m *memoryMailBox) Start() {
+func (m *memoryMailBox) Start(ctx context.Context) {
 	m.started.Do(func() {
-		go m.wireMailCourier()
-		go m.pktMailCourier()
+		ctx, cancel := context.WithCancel(ctx)
+		m.cancel = fn.Some(cancel)
+
+		go m.wireMailCourier(ctx)
+		go m.pktMailCourier(ctx)
 	})
 }
 
@@ -324,6 +330,7 @@ func (m *memoryMailBox) HasPacket(inKey CircuitKey) bool {
 // NOTE: This method is part of the MailBox interface.
 func (m *memoryMailBox) Stop() {
 	m.stopped.Do(func() {
+		m.cancel.WhenSome(func(fn context.CancelFunc) { fn() })
 		close(m.quit)
 
 		m.signalUntilShutdown(wireCourier)
@@ -372,7 +379,7 @@ func (p *pktWithExpiry) deadline(clock clock.Clock) <-chan time.Time {
 
 // wireMailCourier is a dedicated goroutine whose job is to reliably deliver
 // wire messages.
-func (m *memoryMailBox) wireMailCourier() {
+func (m *memoryMailBox) wireMailCourier(ctx context.Context) {
 	defer close(m.wireShutdown)
 
 	for {
@@ -389,6 +396,9 @@ func (m *memoryMailBox) wireMailCourier() {
 			case <-m.quit:
 				m.wireCond.L.Unlock()
 				return
+			case <-ctx.Done():
+				m.wireCond.L.Unlock()
+				return
 			default:
 			}
 		}
@@ -418,13 +428,15 @@ func (m *memoryMailBox) wireMailCourier() {
 			close(msgDone)
 		case <-m.quit:
 			return
+		case <-ctx.Done():
+			return
 		}
 	}
 }
 
 // pktMailCourier is a dedicated goroutine whose job is to reliably deliver
 // packet messages.
-func (m *memoryMailBox) pktMailCourier() {
+func (m *memoryMailBox) pktMailCourier(ctx context.Context) {
 	defer close(m.pktShutdown)
 
 	for {
@@ -447,6 +459,11 @@ func (m *memoryMailBox) pktMailCourier() {
 			case <-m.quit:
 				m.pktCond.L.Unlock()
 				return
+
+			case <-ctx.Done():
+				m.pktCond.L.Unlock()
+				return
+
 			default:
 			}
 		}
@@ -541,7 +558,7 @@ func (m *memoryMailBox) pktMailCourier() {
 		case <-deadline:
 			log.Debugf("Expiring add htlc with "+
 				"keystone=%v", add.keystone())
-			m.FailAdd(add)
+			m.FailAdd(ctx, add)
 
 		case pktDone := <-m.pktReset:
 			m.pktCond.L.Lock()
@@ -553,6 +570,9 @@ func (m *memoryMailBox) pktMailCourier() {
 
 		case <-m.quit:
 			return
+
+		case <-ctx.Done():
+			return
 		}
 	}
 }
@@ -688,9 +708,7 @@ func (m *memoryMailBox) DustPackets() (lnwire.MilliSatoshi,
 // delivered after the link restarts if the switch has remained online. The
 // generated LinkError will show an OutgoingFailureDownstreamHtlcAdd
 // FailureDetail.
-func (m *memoryMailBox) FailAdd(pkt *htlcPacket) {
-	ctx := context.TODO()
-
+func (m *memoryMailBox) FailAdd(ctx context.Context, pkt *htlcPacket) {
 	// First, remove the packet from mailbox. If we didn't find the packet
 	// because it has already been acked, we'll exit early to avoid sending
 	// a duplicate fail message through the switch.
@@ -844,8 +862,8 @@ func (mo *mailOrchestrator) Stop() {
 
 // GetOrCreateMailBox returns an existing mailbox belonging to `chanID`, or
 // creates and returns a new mailbox if none is found.
-func (mo *mailOrchestrator) GetOrCreateMailBox(chanID lnwire.ChannelID,
-	shortChanID lnwire.ShortChannelID) MailBox {
+func (mo *mailOrchestrator) GetOrCreateMailBox(ctx context.Context,
+	chanID lnwire.ChannelID, shortChanID lnwire.ShortChannelID) MailBox {
 
 	// First, try lookup the mailbox directly using only the shared mutex.
 	mo.mu.RLock()
@@ -859,7 +877,7 @@ func (mo *mailOrchestrator) GetOrCreateMailBox(chanID lnwire.ChannelID,
 	// Otherwise, we will try again with exclusive lock, creating a mailbox
 	// if one still has not been created.
 	mo.mu.Lock()
-	mailbox = mo.exclusiveGetOrCreateMailBox(chanID, shortChanID)
+	mailbox = mo.exclusiveGetOrCreateMailBox(ctx, chanID, shortChanID)
 	mo.mu.Unlock()
 
 	return mailbox
@@ -870,7 +888,7 @@ func (mo *mailOrchestrator) GetOrCreateMailBox(chanID lnwire.ChannelID,
 // recorded.
 //
 // NOTE: This method MUST be invoked with the mailOrchestrator's exclusive lock.
-func (mo *mailOrchestrator) exclusiveGetOrCreateMailBox(
+func (mo *mailOrchestrator) exclusiveGetOrCreateMailBox(ctx context.Context,
 	chanID lnwire.ChannelID, shortChanID lnwire.ShortChannelID) MailBox {
 
 	mailbox, ok := mo.mailboxes[chanID]
@@ -882,7 +900,7 @@ func (mo *mailOrchestrator) exclusiveGetOrCreateMailBox(
 			expiry:            mo.cfg.expiry,
 			failMailboxUpdate: mo.cfg.failMailboxUpdate,
 		})
-		mailbox.Start()
+		mailbox.Start(ctx)
 		mo.mailboxes[chanID] = mailbox
 	}
 
@@ -916,7 +934,7 @@ func (mo *mailOrchestrator) BindLiveShortChanID(mailbox MailBox,
 // to channel_id. If the mailbox is found, the message is delivered directly.
 // Otherwise the packet is recorded as unclaimed, and will be delivered to the
 // mailbox upon the subsequent call to BindLiveShortChanID.
-func (mo *mailOrchestrator) Deliver(
+func (mo *mailOrchestrator) Deliver(ctx context.Context,
 	sid lnwire.ShortChannelID, pkt *htlcPacket) error {
 
 	var (
@@ -961,7 +979,7 @@ func (mo *mailOrchestrator) Deliver(
 		// index should only be set if the mailbox had been initialized
 		// beforehand.  However, this does ensure that this case is
 		// handled properly in the event that it could happen.
-		mailbox = mo.exclusiveGetOrCreateMailBox(chanID, sid)
+		mailbox = mo.exclusiveGetOrCreateMailBox(ctx, chanID, sid)
 		mo.mu.Unlock()
 
 		// Deliver the packet to the mailbox if it was found or created.

htlcswitch/mailbox_test.go

@@ -218,7 +218,7 @@ func newMailboxContextWithClock(t *testing.T,
 		forwardPackets:    ctx.forward,
 		clock:             clock,
 	})
-	ctx.mailbox.Start()
+	ctx.mailbox.Start(context.Background())
 	t.Cleanup(ctx.mailbox.Stop)
 
 	return ctx
@@ -245,7 +245,7 @@ func newMailboxContext(t *testing.T, startTime time.Time,
 		clock:             ctx.clock,
 		expiry:            expiry,
 	})
-	ctx.mailbox.Start()
+	ctx.mailbox.Start(context.Background())
 	t.Cleanup(ctx.mailbox.Stop)
 
 	return ctx
@@ -334,6 +334,7 @@ func (c *mailboxContext) checkFails(adds []*htlcPacket) {
 // TestMailBoxFailAdd asserts that FailAdd returns a response to the switch
 // under various interleavings with other operations on the mailbox.
 func TestMailBoxFailAdd(t *testing.T) {
+	t.Parallel()
 	var (
 		batchDelay       = time.Second
 		expiry           = time.Minute
@@ -346,7 +347,7 @@ func TestMailBoxFailAdd(t *testing.T) {
 
 	failAdds := func(adds []*htlcPacket) {
 		for _, add := range adds {
-			ctx.mailbox.FailAdd(add)
+			ctx.mailbox.FailAdd(context.Background(), add)
 		}
 	}
 
@@ -541,7 +542,7 @@ func TestMailBoxDuplicateAddPacket(t *testing.T) {
 	t.Parallel()
 
 	ctx := newMailboxContext(t, time.Now(), testExpiry)
-	ctx.mailbox.Start()
+	ctx.mailbox.Start(context.Background())
 
 	addTwice := func(t *testing.T, pkt *htlcPacket) {
 		// The first add should succeed.
@@ -697,6 +698,7 @@ func testMailBoxDust(t *testing.T, chantype channeldb.ChannelType) {
 // readily to mailboxes for channels that are already in the live state.
 func TestMailOrchestrator(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	failMailboxUpdate := func(_ context.Context, outScid,
 		mboxScid lnwire.ShortChannelID) lnwire.FailureMessage {
@@ -738,11 +740,11 @@ func TestMailOrchestrator(t *testing.T) {
 		}
 		sentPackets[i] = pkt
 
-		mo.Deliver(pkt.outgoingChanID, pkt)
+		mo.Deliver(ctx, pkt.outgoingChanID, pkt)
 	}
 
 	// Now, initialize a new mailbox for Alice's chanid.
-	mailbox := mo.GetOrCreateMailBox(chanID1, aliceChanID)
+	mailbox := mo.GetOrCreateMailBox(ctx, chanID1, aliceChanID)
 
 	// Verify that no messages are received, since Alice's mailbox has not
 	// been made live.
@@ -787,7 +789,7 @@ func TestMailOrchestrator(t *testing.T) {
 
 	// For the second half of the test, create a new mailbox for Bob and
 	// immediately make it live with an assigned short chan id.
-	mailbox = mo.GetOrCreateMailBox(chanID2, bobChanID)
+	mailbox = mo.GetOrCreateMailBox(ctx, chanID2, bobChanID)
 	mo.BindLiveShortChanID(mailbox, chanID2, bobChanID)
 
 	// Create the second half of our htlcs, and deliver them via the
@@ -806,7 +808,7 @@ func TestMailOrchestrator(t *testing.T) {
 		}
 		sentPackets[i] = pkt
 
-		mo.Deliver(pkt.incomingChanID, pkt)
+		mo.Deliver(ctx, pkt.incomingChanID, pkt)
 
 		timeout := time.After(50 * time.Millisecond)
 		select {