htlcswitch: add ctx, cancel to memoryMailbox

And remove one context.TODO

Author: ellemouton

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.
@@ -149,6 +150,7 @@ type memoryMailBox struct {
 	wireShutdown chan struct{}
 	pktShutdown  chan struct{}
 	quit         chan struct{}
+	cancel       fn.Option[context.CancelFunc]
 
 	// 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
@@ -207,8 +209,11 @@ const (
 // NOTE: This method is part of the MailBox interface.
 func (m *memoryMailBox) Start() {
 	m.started.Do(func() {
+		ctx, cancel := context.WithCancel(context.Background())
+		m.cancel = fn.Some(cancel)
+
 		go m.wireMailCourier()
-		go m.pktMailCourier()
+		go m.pktMailCourier(ctx)
 	})
 }
 
@@ -324,6 +329,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)
@@ -424,7 +430,7 @@ func (m *memoryMailBox) wireMailCourier() {
 
 // 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 {
@@ -541,7 +547,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()
@@ -688,9 +694,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.

htlcswitch/mailbox_test.go

@@ -334,6 +334,8 @@ 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 +348,7 @@ func TestMailBoxFailAdd(t *testing.T) {
 
 	failAdds := func(adds []*htlcPacket) {
 		for _, add := range adds {
-			ctx.mailbox.FailAdd(add)
+			ctx.mailbox.FailAdd(context.Background(), add)
 		}
 	}