htlcswitch: continue threading context through

Here, we focus on the `interceptForward.FailWithCode` method which also
makes a call to the switch's `FetchLastChannelUpdate` method which will
later take a context since it is a call to the graph DB.

Author: ellemouton

htlcswitch/held_htlc_set.go

@@ -1,6 +1,7 @@
 package htlcswitch
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -39,13 +40,15 @@ func (h *heldHtlcSet) popAll(cb func(InterceptedForward)) {
 
 // popAutoFails calls the callback for each forward that has an auto-fail height
 // equal or less then the specified pop height and removes them from the set.
-func (h *heldHtlcSet) popAutoFails(height uint32, cb func(InterceptedForward)) {
+func (h *heldHtlcSet) popAutoFails(ctx context.Context, height uint32,
+	cb func(context.Context, InterceptedForward)) {
+
 	for key, fwd := range h.set {
 		if uint32(fwd.Packet().AutoFailHeight) > height {
 			continue
 		}
 
-		cb(fwd)
+		cb(ctx, fwd)
 
 		delete(h.set, key)
 	}

htlcswitch/held_htlc_set_test.go

@@ -1,6 +1,7 @@
 package htlcswitch
 
 import (
+	"context"
 	"testing"
 
 	"github.com/lightningnetwork/lnd/graph/db/models"
@@ -80,6 +81,8 @@ func TestHeldHtlcSet(t *testing.T) {
 }
 
 func TestHeldHtlcSetAutoFails(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
 	set := newHeldHtlcSet()
 
 	key := models.CircuitKey{
@@ -98,17 +101,17 @@ func TestHeldHtlcSetAutoFails(t *testing.T) {
 	// Test popping auto fails up to one block before the auto-fail height
 	// of our forward.
 	set.popAutoFails(
-		autoFailHeight-1,
-		func(_ InterceptedForward) {
+		ctx, autoFailHeight-1,
+		func(_ context.Context, _ InterceptedForward) {
 			require.Fail(t, "unexpected fwd")
 		},
 	)
 
 	// Popping succeeds at the auto-fail height.
 	cbCalled := false
 	set.popAutoFails(
-		autoFailHeight,
-		func(poppedFwd InterceptedForward) {
+		ctx, autoFailHeight,
+		func(_ context.Context, poppedFwd InterceptedForward) {
 			cbCalled = true
 
 			require.Equal(t, fwd, poppedFwd)
@@ -118,8 +121,8 @@ func TestHeldHtlcSetAutoFails(t *testing.T) {
 
 	// After this, there should be nothing more to pop.
 	set.popAutoFails(
-		autoFailHeight,
-		func(_ InterceptedForward) {
+		ctx, autoFailHeight,
+		func(_ context.Context, _ InterceptedForward) {
 			require.Fail(t, "unexpected fwd")
 		},
 	)

htlcswitch/interceptable_switch.go

@@ -1,6 +1,7 @@
 package htlcswitch
 
 import (
+	"context"
 	"crypto/sha256"
 	"fmt"
 	"sync"
@@ -89,8 +90,9 @@ type InterceptableSwitch struct {
 	// currentHeight is the currently best known height.
 	currentHeight int32
 
-	wg   sync.WaitGroup
-	quit chan struct{}
+	wg     sync.WaitGroup
+	quit   chan struct{}
+	cancel fn.Option[context.CancelFunc]
 }
 
 type interceptedPackets struct {
@@ -222,12 +224,14 @@ func (s *InterceptableSwitch) SetInterceptor(
 	}
 }
 
-func (s *InterceptableSwitch) Start() error {
+func (s *InterceptableSwitch) Start(ctx context.Context) error {
 	log.Info("InterceptableSwitch starting...")
 
 	if s.started.Swap(true) {
 		return fmt.Errorf("InterceptableSwitch started more than once")
 	}
+	ctx, cancel := context.WithCancel(ctx)
+	s.cancel = fn.Some(cancel)
 
 	blockEpochStream, err := s.notifier.RegisterBlockEpochNtfn(nil)
 	if err != nil {
@@ -239,7 +243,7 @@ func (s *InterceptableSwitch) Start() error {
 	go func() {
 		defer s.wg.Done()
 
-		err := s.run()
+		err := s.run(ctx)
 		if err != nil {
 			log.Errorf("InterceptableSwitch stopped: %v", err)
 		}
@@ -257,6 +261,7 @@ func (s *InterceptableSwitch) Stop() error {
 		return fmt.Errorf("InterceptableSwitch stopped more than once")
 	}
 
+	s.cancel.WhenSome(func(fn context.CancelFunc) { fn() })
 	close(s.quit)
 	s.wg.Wait()
 
@@ -271,7 +276,7 @@ func (s *InterceptableSwitch) Stop() error {
 	return nil
 }
 
-func (s *InterceptableSwitch) run() error {
+func (s *InterceptableSwitch) run(ctx context.Context) error {
 	// The block epoch stream will immediately stream the current height.
 	// Read it out here.
 	select {
@@ -298,7 +303,7 @@ func (s *InterceptableSwitch) run() error {
 			var notIntercepted []*htlcPacket
 			for _, p := range packets.packets {
 				intercepted, err := s.interceptForward(
-					p, packets.isReplay,
+					ctx, p, packets.isReplay,
 				)
 				if err != nil {
 					return err
@@ -325,12 +330,12 @@ func (s *InterceptableSwitch) run() error {
 			// already intercepted in the off-chain flow. And even
 			// if not, it is safe to signal replay so that we won't
 			// unexpectedly skip over this htlc.
-			if _, err := s.forward(fwd, true); err != nil {
+			if _, err := s.forward(ctx, fwd, true); err != nil {
 				return err
 			}
 
 		case res := <-s.resolutionChan:
-			res.errChan <- s.resolve(res.resolution)
+			res.errChan <- s.resolve(ctx, res.resolution)
 
 		case currentBlock, ok := <-s.blockEpochStream.Epochs:
 			if !ok {
@@ -341,20 +346,23 @@ func (s *InterceptableSwitch) run() error {
 
 			// A new block is appended. Fail any held htlcs that
 			// expire at this height to prevent channel force-close.
-			s.failExpiredHtlcs()
+			s.failExpiredHtlcs(ctx)
 
 		case <-s.quit:
 			return nil
+
+		case <-ctx.Done():
+			return ctx.Err()
 		}
 	}
 }
 
-func (s *InterceptableSwitch) failExpiredHtlcs() {
+func (s *InterceptableSwitch) failExpiredHtlcs(ctx context.Context) {
 	s.heldHtlcSet.popAutoFails(
-		uint32(s.currentHeight),
-		func(fwd InterceptedForward) {
+		ctx, uint32(s.currentHeight),
+		func(ctx context.Context, fwd InterceptedForward) {
 			err := fwd.FailWithCode(
-				lnwire.CodeTemporaryChannelFailure,
+				ctx, lnwire.CodeTemporaryChannelFailure,
 			)
 			if err != nil {
 				log.Errorf("Cannot fail packet: %v", err)
@@ -407,7 +415,9 @@ func (s *InterceptableSwitch) setInterceptor(interceptor ForwardInterceptor) {
 
 // resolve processes a HTLC given the resolution type specified by the
 // intercepting client.
-func (s *InterceptableSwitch) resolve(res *FwdResolution) error {
+func (s *InterceptableSwitch) resolve(ctx context.Context,
+	res *FwdResolution) error {
+
 	intercepted, err := s.heldHtlcSet.pop(res.Key)
 	if err != nil {
 		return err
@@ -431,7 +441,7 @@ func (s *InterceptableSwitch) resolve(res *FwdResolution) error {
 			return intercepted.Fail(res.FailureMessage)
 		}
 
-		return intercepted.FailWithCode(res.FailureCode)
+		return intercepted.FailWithCode(ctx, res.FailureCode)
 
 	default:
 		return fmt.Errorf("unrecognized action %v", res.Action)
@@ -503,8 +513,8 @@ func (s *InterceptableSwitch) ForwardPacket(
 
 // interceptForward forwards the packet to the external interceptor after
 // checking the interception criteria.
-func (s *InterceptableSwitch) interceptForward(packet *htlcPacket,
-	isReplay bool) (bool, error) {
+func (s *InterceptableSwitch) interceptForward(ctx context.Context,
+	packet *htlcPacket, isReplay bool) (bool, error) {
 
 	switch htlc := packet.htlc.(type) {
 	case *lnwire.UpdateAddHTLC:
@@ -522,7 +532,7 @@ func (s *InterceptableSwitch) interceptForward(packet *htlcPacket,
 		}
 
 		// Handle forwards that are too close to expiry.
-		handled, err := s.handleExpired(intercepted)
+		handled, err := s.handleExpired(ctx, intercepted)
 		if err != nil {
 			log.Errorf("Error handling intercepted htlc "+
 				"that expires too soon: circuit=%v, "+
@@ -542,15 +552,15 @@ func (s *InterceptableSwitch) interceptForward(packet *htlcPacket,
 			return true, nil
 		}
 
-		return s.forward(intercepted, isReplay)
+		return s.forward(ctx, intercepted, isReplay)
 
 	default:
 		return false, nil
 	}
 }
 
 // forward records the intercepted htlc and forwards it to the interceptor.
-func (s *InterceptableSwitch) forward(
+func (s *InterceptableSwitch) forward(ctx context.Context,
 	fwd InterceptedForward, isReplay bool) (bool, error) {
 
 	inKey := fwd.Packet().IncomingCircuit
@@ -573,7 +583,7 @@ func (s *InterceptableSwitch) forward(
 		// yet. This limits the backlog of htlcs when the interceptor is down.
 		if !isReplay {
 			err := fwd.FailWithCode(
-				lnwire.CodeTemporaryChannelFailure,
+				ctx, lnwire.CodeTemporaryChannelFailure,
 			)
 			if err != nil {
 				log.Errorf("Cannot fail packet: %v", err)
@@ -605,8 +615,8 @@ func (s *InterceptableSwitch) forward(
 
 // handleExpired checks that the htlc isn't too close to the channel
 // force-close broadcast height. If it is, it is cancelled back.
-func (s *InterceptableSwitch) handleExpired(fwd *interceptedForward) (
-	bool, error) {
+func (s *InterceptableSwitch) handleExpired(ctx context.Context,
+	fwd *interceptedForward) (bool, error) {
 
 	height := uint32(s.currentHeight)
 	if fwd.packet.incomingTimeout >= height+s.cltvInterceptDelta {
@@ -620,7 +630,7 @@ func (s *InterceptableSwitch) handleExpired(fwd *interceptedForward) (
 		fwd.packet.incomingTimeout)
 
 	err := fwd.FailWithCode(
-		lnwire.CodeExpiryTooSoon,
+		ctx, lnwire.CodeExpiryTooSoon,
 	)
 	if err != nil {
 		return false, err
@@ -747,7 +757,9 @@ func (f *interceptedForward) Fail(reason []byte) error {
 
 // FailWithCode notifies the intention to fail an existing hold forward with the
 // specified failure code.
-func (f *interceptedForward) FailWithCode(code lnwire.FailCode) error {
+func (f *interceptedForward) FailWithCode(_ context.Context,
+	code lnwire.FailCode) error {
+
 	shaOnionBlob := func() [32]byte {
 		return sha256.Sum256(f.htlc.OnionBlob[:])
 	}

htlcswitch/interfaces.go

@@ -454,7 +454,7 @@ type InterceptedForward interface {
 
 	// FailWithCode notifies the intention to fail an existing hold forward
 	// with the specified failure code.
-	FailWithCode(code lnwire.FailCode) error
+	FailWithCode(ctx context.Context, code lnwire.FailCode) error
 }
 
 // htlcNotifier is an interface which represents the input side of the

htlcswitch/switch_test.go

@@ -3893,6 +3893,7 @@ func (c *interceptableSwitchTestContext) createSettlePacket(
 
 func TestSwitchHoldForward(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	c := newInterceptableSwitchTestContext(t)
 	defer c.finish()
@@ -3911,7 +3912,7 @@ func TestSwitchHoldForward(t *testing.T) {
 		},
 	)
 	require.NoError(t, err)
-	require.NoError(t, switchForwardInterceptor.Start())
+	require.NoError(t, switchForwardInterceptor.Start(ctx))
 
 	switchForwardInterceptor.SetInterceptor(c.forwardInterceptor.InterceptForwardHtlc)
 	linkQuit := make(chan struct{})
@@ -4115,7 +4116,7 @@ func TestSwitchHoldForward(t *testing.T) {
 		},
 	)
 	require.NoError(t, err)
-	require.NoError(t, switchForwardInterceptor.Start())
+	require.NoError(t, switchForwardInterceptor.Start(ctx))
 
 	// Forward a fresh packet. It is expected to be failed immediately,
 	// because there is no interceptor registered.
@@ -4183,6 +4184,7 @@ func TestSwitchHoldForward(t *testing.T) {
 
 func TestInterceptableSwitchWatchDog(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	c := newInterceptableSwitchTestContext(t)
 	defer c.finish()
@@ -4202,7 +4204,7 @@ func TestInterceptableSwitchWatchDog(t *testing.T) {
 		},
 	)
 	require.NoError(t, err)
-	require.NoError(t, switchForwardInterceptor.Start())
+	require.NoError(t, switchForwardInterceptor.Start(ctx))
 
 	// Set interceptor.
 	switchForwardInterceptor.SetInterceptor(
@@ -5439,6 +5441,7 @@ func TestSwitchAliasInterceptFail(t *testing.T) {
 
 func testSwitchAliasInterceptFail(t *testing.T, zeroConf bool) {
 	t.Parallel()
+	ctx := context.Background()
 
 	chanID, aliceScid := genID()
 
@@ -5524,14 +5527,14 @@ func testSwitchAliasInterceptFail(t *testing.T, zeroConf bool) {
 		},
 	)
 	require.NoError(t, err)
-	require.NoError(t, interceptSwitch.Start())
+	require.NoError(t, interceptSwitch.Start(ctx))
 	interceptSwitch.SetInterceptor(forwardInterceptor.InterceptForwardHtlc)
 
 	err = interceptSwitch.ForwardPackets(nil, false, ogPacket)
 	require.NoError(t, err)
 
 	inCircuit := forwardInterceptor.getIntercepted().IncomingCircuit
-	require.NoError(t, interceptSwitch.resolve(&FwdResolution{
+	require.NoError(t, interceptSwitch.resolve(ctx, &FwdResolution{
 		Action:      FwdActionFail,
 		Key:         inCircuit,
 		FailureCode: lnwire.CodeTemporaryChannelFailure,

intercepted_forward.go

@@ -1,6 +1,7 @@
 package lnd
 
 import (
+	"context"
 	"errors"
 
 	"github.com/lightningnetwork/lnd/fn/v2"
@@ -72,7 +73,9 @@ func (f *interceptedForward) Fail(_ []byte) error {
 
 // FailWithCode notifies the intention to fail an existing hold forward with the
 // specified failure code.
-func (f *interceptedForward) FailWithCode(_ lnwire.FailCode) error {
+func (f *interceptedForward) FailWithCode(_ context.Context,
+	_ lnwire.FailCode) error {
+
 	return ErrCannotFail
 }
 

server.go

@@ -2353,7 +2353,7 @@ func (s *server) Start(ctx context.Context) error {
 		}
 
 		cleanup = cleanup.add(s.interceptableSwitch.Stop)
-		if err := s.interceptableSwitch.Start(); err != nil {
+		if err := s.interceptableSwitch.Start(ctx); err != nil {
 			startErr = err
 			return
 		}