Repository: lightningnetwork/lnd
Pull request #9344 (2024-12-10 12:13:53)

htlcswitch+go.mod: use updated fn.ContextGuard
This commit updates the fn dep to the version containing the updates to the ContextGuard implementation. Only the htlcswitch/link uses the guard at the moment so this is updated to make use of the new implementation.

Depends on https://github.com/lightningnetwork/lnd/pull/9342 (so only look at the last commit) 


coderabbitai[bot]: <!-- This is an auto-generated comment: summarize by coderabbit.ai -->
<!-- This is an auto-generated comment: skip review by coderabbit.ai -->

> [!IMPORTANT]
> ## Review skipped
> 
> Auto reviews are limited to specific labels.
> 
> <details>
> <summary>:label: Labels to auto review (1)</summary>
> 
> * llm-review
> 
> </details>
> 
> Please check the settings in the CodeRabbit UI or the `.coderabbit.yaml` file in this repository. To trigger a single review, invoke the `@coderabbitai review` command.
> 
> You can disable this status message by setting the `reviews.review_status` to `false` in the CodeRabbit configuration file.

<!-- end of auto-generated comment: skip review by coderabbit.ai -->

<!-- tips_start -->

---

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

<details>
<summary>❤️ Share</summary>

- [X](https://twitter.com/intent/tweet?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A&url=https%3A//coderabbit.ai)
- [Mastodon](https://mastodon.social/share?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A%20https%3A%2F%2Fcoderabbit.ai)
- [Reddit](https://www.reddit.com/submit?title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&text=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code.%20Check%20it%20out%3A%20https%3A//coderabbit.ai)
- [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https%3A%2F%2Fcoderabbit.ai&mini=true&title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&summary=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code)

</details>

<details>
<summary>🪧 Tips</summary>

### Chat

There are 3 ways to chat with [CodeRabbit](https://coderabbit.ai):

- Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  - `I pushed a fix in commit <commit_id>, please review it.`
  - `Generate unit testing code for this file.`
  - `Open a follow-up GitHub issue for this discussion.`
- Files and specific lines of code (under the "Files changed" tab): Tag `@coderabbitai` in a new review comment at the desired location with your query. Examples:
  - `@coderabbitai generate unit testing code for this file.`
  -	`@coderabbitai modularize this function.`
- PR comments: Tag `@coderabbitai` in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  - `@coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.`
  - `@coderabbitai read src/utils.ts and generate unit testing code.`
  - `@coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.`
  - `@coderabbitai help me debug CodeRabbit configuration file.`

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

### CodeRabbit Commands (Invoked using PR comments)

- `@coderabbitai pause` to pause the reviews on a PR.
- `@coderabbitai resume` to resume the paused reviews.
- `@coderabbitai review` to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
- `@coderabbitai full review` to do a full review from scratch and review all the files again.
- `@coderabbitai summary` to regenerate the summary of the PR.
- `@coderabbitai generate docstrings` to [generate docstrings](https://docs.coderabbit.ai/finishing-touches/docstrings) for this PR. ([Experiment](https://docs.coderabbit.ai/early-access#experiments))
- `@coderabbitai resolve` resolve all the CodeRabbit review comments.
- `@coderabbitai configuration` to show the current CodeRabbit configuration for the repository.
- `@coderabbitai help` to get help.

### Other keywords and placeholders

- Add `@coderabbitai ignore` anywhere in the PR description to prevent this PR from being reviewed.
- Add `@coderabbitai summary` to generate the high-level summary at a specific location in the PR description.
- Add `@coderabbitai` anywhere in the PR title to generate the title automatically.

### CodeRabbit Configuration File (`.coderabbit.yaml`)

- You can programmatically configure CodeRabbit by adding a `.coderabbit.yaml` file to the root of your repository.
- Please see the [configuration documentation](https://docs.coderabbit.ai/guides/configure-coderabbit) for more information.
- If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: `# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json`

### Documentation and Community

- Visit our [Documentation](https://docs.coderabbit.ai) for detailed information on how to use CodeRabbit.
- Join our [Discord Community](http://discord.gg/coderabbit) to get help, request features, and share feedback.
- Follow us on [X/Twitter](https://twitter.com/coderabbitai) for updates and announcements.

</details>

<!-- tips_end -->
lightninglabs-deploy: @ellemouton, remember to re-request review from reviewers when ready

Patch:
diff --git a/go.mod b/go.mod
index 8ed20be32..5f4219afa 100644
--- a/go.mod
+++ b/go.mod
@@ -36,7 +36,7 @@ require (
 	github.com/lightningnetwork/lightning-onion v1.2.1-0.20240712235311-98bd56499dfb
 	github.com/lightningnetwork/lnd/cert v1.2.2
 	github.com/lightningnetwork/lnd/clock v1.1.1
-	github.com/lightningnetwork/lnd/fn/v2 v2.0.4
+	github.com/lightningnetwork/lnd/fn/v2 v2.0.8
 	github.com/lightningnetwork/lnd/healthcheck v1.2.6
 	github.com/lightningnetwork/lnd/kvdb v1.4.12
 	github.com/lightningnetwork/lnd/queue v1.1.1
diff --git a/go.sum b/go.sum
index 0433b87e5..6e744d679 100644
--- a/go.sum
+++ b/go.sum
@@ -456,8 +456,8 @@ github.com/lightningnetwork/lnd/cert v1.2.2 h1:71YK6hogeJtxSxw2teq3eGeuy4rHGKcFf
 github.com/lightningnetwork/lnd/cert v1.2.2/go.mod h1:jQmFn/Ez4zhDgq2hnYSw8r35bqGVxViXhX6Cd7HXM6U=
 github.com/lightningnetwork/lnd/clock v1.1.1 h1:OfR3/zcJd2RhH0RU+zX/77c0ZiOnIMsDIBjgjWdZgA0=
 github.com/lightningnetwork/lnd/clock v1.1.1/go.mod h1:mGnAhPyjYZQJmebS7aevElXKTFDuO+uNFFfMXK1W8xQ=
-github.com/lightningnetwork/lnd/fn/v2 v2.0.4 h1:DiC/AEa7DhnY4qOEQBISu1cp+1+51LjbVDzNLVBwNjI=
-github.com/lightningnetwork/lnd/fn/v2 v2.0.4/go.mod h1:TOzwrhjB/Azw1V7aa8t21ufcQmdsQOQMDtxVOQWNl8s=
+github.com/lightningnetwork/lnd/fn/v2 v2.0.8 h1:r2SLz7gZYQPVc3IZhU82M66guz3Zk2oY+Rlj9QN5S3g=
+github.com/lightningnetwork/lnd/fn/v2 v2.0.8/go.mod h1:TOzwrhjB/Azw1V7aa8t21ufcQmdsQOQMDtxVOQWNl8s=
 github.com/lightningnetwork/lnd/healthcheck v1.2.6 h1:1sWhqr93GdkWy4+6U7JxBfcyZIE78MhIHTJZfPx7qqI=
 github.com/lightningnetwork/lnd/healthcheck v1.2.6/go.mod h1:Mu02um4CWY/zdTOvFje7WJgJcHyX2zq/FG3MhOAiGaQ=
 github.com/lightningnetwork/lnd/kvdb v1.4.12 h1:Y0WY5Tbjyjn6eCYh068qkWur5oFtioJlfxc8w5SlJeQ=
diff --git a/htlcswitch/link.go b/htlcswitch/link.go
index 6e67e87de..f0e958992 100644
--- a/htlcswitch/link.go
+++ b/htlcswitch/link.go
@@ -2,6 +2,7 @@ package htlcswitch
 
 import (
 	"bytes"
+	"context"
 	crand "crypto/rand"
 	"crypto/sha256"
 	"errors"
@@ -408,10 +409,10 @@ type channelLink struct {
 	// the result.
 	quiescenceReqs chan StfuReq
 
-	// ContextGuard is a helper that encapsulates a wait group and quit
-	// channel and allows contexts that either block or cancel on those
-	// depending on the use case.
-	*fn.ContextGuard
+	// cg is a helper that encapsulates a wait group and quit channel and
+	// allows contexts that either block or cancel on those depending on
+	// the use case.
+	cg *fn.ContextGuard
 }
 
 // hookMap is a data structure that is used to track the hooks that need to be
@@ -517,7 +518,7 @@ func NewChannelLink(cfg ChannelLinkConfig,
 		incomingCommitHooks: newHookMap(),
 		quiescer:            qsm,
 		quiescenceReqs:      quiescenceReqs,
-		ContextGuard:        fn.NewContextGuard(),
+		cg:                  fn.NewContextGuard(),
 	}
 }
 
@@ -596,8 +597,8 @@ func (l *channelLink) Start() error {
 
 	l.updateFeeTimer = time.NewTimer(l.randomFeeUpdateTimeout())
 
-	l.Wg.Add(1)
-	go l.htlcManager()
+	l.cg.WgAdd(1)
+	go l.htlcManager(context.TODO())
 
 	return nil
 }
@@ -636,8 +637,8 @@ func (l *channelLink) Stop() {
 		l.hodlQueue.Stop()
 	}
 
-	close(l.Quit)
-	l.Wg.Wait()
+	l.cg.Quit()
+	l.cg.WgWait()
 
 	// Now that the htlcManager has completely exited, reset the packet
 	// courier. This allows the mailbox to revaluate any lingering Adds that
@@ -662,7 +663,7 @@ func (l *channelLink) Stop() {
 // WaitForShutdown blocks until the link finishes shutting down, which includes
 // termination of all dependent goroutines.
 func (l *channelLink) WaitForShutdown() {
-	l.Wg.Wait()
+	l.cg.WgWait()
 }
 
 // EligibleToForward returns a bool indicating if the channel is able to
@@ -740,7 +741,7 @@ func (l *channelLink) IsFlushing(linkDirection LinkDirection) bool {
 func (l *channelLink) OnFlushedOnce(hook func()) {
 	select {
 	case l.flushHooks.newTransients <- hook:
-	case <-l.Quit:
+	case <-l.cg.Done():
 	}
 }
 
@@ -759,7 +760,7 @@ func (l *channelLink) OnCommitOnce(direction LinkDirection, hook func()) {
 
 	select {
 	case queue <- hook:
-	case <-l.Quit:
+	case <-l.cg.Done():
 	}
 }
 
@@ -777,7 +778,7 @@ func (l *channelLink) InitStfu() <-chan fn.Result[lntypes.ChannelParty] {
 
 	select {
 	case l.quiescenceReqs <- req:
-	case <-l.Quit:
+	case <-l.cg.Done():
 		req.Resolve(fn.Err[lntypes.ChannelParty](ErrLinkShuttingDown))
 	}
 
@@ -887,7 +888,7 @@ func (l *channelLink) createFailureWithUpdate(incoming bool,
 // This method is to be called upon reconnection after the initial funding
 // flow. We'll compare out commitment chains with the remote party, and re-send
 // either a danging commit signature, a revocation, or both.
-func (l *channelLink) syncChanStates() error {
+func (l *channelLink) syncChanStates(ctx context.Context) error {
 	chanState := l.channel.State()
 
 	l.log.Infof("Attempting to re-synchronize channel: %v", chanState)
@@ -989,7 +990,7 @@ func (l *channelLink) syncChanStates() error {
 		// We've just received a ChanSync message from the remote
 		// party, so we'll process the message  in order to determine
 		// if we need to re-transmit any messages to the remote party.
-		ctx, cancel := l.WithCtxQuitNoTimeout()
+		ctx, cancel := l.cg.Create(ctx)
 		defer cancel()
 		msgsToReSend, openedCircuits, closedCircuits, err =
 			l.channel.ProcessChanSyncMsg(ctx, remoteChanSyncMsg)
@@ -1021,7 +1022,7 @@ func (l *channelLink) syncChanStates() error {
 			l.cfg.Peer.SendMessage(false, msg)
 		}
 
-	case <-l.Quit:
+	case <-l.cg.Done():
 		return ErrLinkShuttingDown
 	}
 
@@ -1033,7 +1034,7 @@ func (l *channelLink) syncChanStates() error {
 // we previously received are reinstated in memory, and forwarded to the switch
 // if necessary. After a restart, this will also delete any previously
 // completed packages.
-func (l *channelLink) resolveFwdPkgs() error {
+func (l *channelLink) resolveFwdPkgs(ctx context.Context) error {
 	fwdPkgs, err := l.channel.LoadFwdPkgs()
 	if err != nil {
 		return err
@@ -1050,7 +1051,7 @@ func (l *channelLink) resolveFwdPkgs() error {
 	// If any of our reprocessing steps require an update to the commitment
 	// txn, we initiate a state transition to capture all relevant changes.
 	if l.channel.NumPendingUpdates(lntypes.Local, lntypes.Remote) > 0 {
-		return l.updateCommitTx()
+		return l.updateCommitTx(ctx)
 	}
 
 	return nil
@@ -1111,7 +1112,7 @@ func (l *channelLink) resolveFwdPkg(fwdPkg *channeldb.FwdPkg) error {
 //
 // NOTE: This MUST be run as a goroutine.
 func (l *channelLink) fwdPkgGarbager() {
-	defer l.Wg.Done()
+	defer l.cg.WgDone()
 
 	l.cfg.FwdPkgGCTicker.Resume()
 	defer l.cfg.FwdPkgGCTicker.Stop()
@@ -1128,7 +1129,7 @@ func (l *channelLink) fwdPkgGarbager() {
 					err)
 				continue
 			}
-		case <-l.Quit:
+		case <-l.cg.Done():
 			return
 		}
 	}
@@ -1248,10 +1249,12 @@ func (l *channelLink) handleChanSyncErr(err error) {
 // and also the htlc trickle queue+timer for this active channels.
 //
 // NOTE: This MUST be run as a goroutine.
-func (l *channelLink) htlcManager() {
+//
+//nolint:funlen
+func (l *channelLink) htlcManager(ctx context.Context) {
 	defer func() {
 		l.cfg.BatchTicker.Stop()
-		l.Wg.Done()
+		l.cg.WgDone()
 		l.log.Infof("exited")
 	}()
 
@@ -1271,7 +1274,7 @@ func (l *channelLink) htlcManager() {
 	// re-synchronize state with the remote peer. settledHtlcs is a map of
 	// HTLC's that we re-settled as part of the channel state sync.
 	if l.cfg.SyncStates {
-		err := l.syncChanStates()
+		err := l.syncChanStates(ctx)
 		if err != nil {
 			l.handleChanSyncErr(err)
 			return
@@ -1322,7 +1325,7 @@ func (l *channelLink) htlcManager() {
 	// the channel is not pending, otherwise we should have no htlcs to
 	// reforward.
 	if l.ShortChanID() != hop.Source {
-		err := l.resolveFwdPkgs()
+		err := l.resolveFwdPkgs(ctx)
 		switch err {
 		// No error was encountered, success.
 		case nil:
@@ -1345,7 +1348,7 @@ func (l *channelLink) htlcManager() {
 		// With our link's in-memory state fully reconstructed, spawn a
 		// goroutine to manage the reclamation of disk space occupied by
 		// completed forwarding packages.
-		l.Wg.Add(1)
+		l.cg.WgAdd(1)
 		go l.fwdPkgGarbager()
 	}
 
@@ -1447,7 +1450,8 @@ func (l *channelLink) htlcManager() {
 
 			// If we do, then we'll send a new UpdateFee message to
 			// the remote party, to be locked in with a new update.
-			if err := l.updateChannelFee(newCommitFee); err != nil {
+			err = l.updateChannelFee(ctx, newCommitFee)
+			if err != nil {
 				l.log.Errorf("unable to update fee rate: %v",
 					err)
 				continue
@@ -1475,7 +1479,7 @@ func (l *channelLink) htlcManager() {
 			// including all the currently pending entries. If the
 			// send was unsuccessful, then abandon the update,
 			// waiting for the revocation window to open up.
-			if !l.updateCommitTxOrFail() {
+			if !l.updateCommitTxOrFail(ctx) {
 				return
 			}
 
@@ -1493,19 +1497,19 @@ func (l *channelLink) htlcManager() {
 		// that the link is an intermediate hop in a multi-hop HTLC
 		// circuit.
 		case pkt := <-l.downstream:
-			l.handleDownstreamPkt(pkt)
+			l.handleDownstreamPkt(ctx, pkt)
 
 		// A message from the connected peer was just received. This
 		// indicates that we have a new incoming HTLC, either directly
 		// for us, or part of a multi-hop HTLC circuit.
 		case msg := <-l.upstream:
-			l.handleUpstreamMsg(msg)
+			l.handleUpstreamMsg(ctx, msg)
 
 		// A htlc resolution is received. This means that we now have a
 		// resolution for a previously accepted htlc.
 		case hodlItem := <-l.hodlQueue.ChanOut():
 			htlcResolution := hodlItem.(invoices.HtlcResolution)
-			err := l.processHodlQueue(htlcResolution)
+			err := l.processHodlQueue(ctx, htlcResolution)
 			switch err {
 			// No error, success.
 			case nil:
@@ -1543,7 +1547,7 @@ func (l *channelLink) htlcManager() {
 				}
 			}
 
-		case <-l.Quit:
+		case <-l.cg.Done():
 			return
 		}
 	}
@@ -1552,7 +1556,7 @@ func (l *channelLink) htlcManager() {
 // processHodlQueue processes a received htlc resolution and continues reading
 // from the hodl queue until no more resolutions remain. When this function
 // returns without an error, the commit tx should be updated.
-func (l *channelLink) processHodlQueue(
+func (l *channelLink) processHodlQueue(ctx context.Context,
 	firstResolution invoices.HtlcResolution) error {
 
 	// Try to read all waiting resolution messages, so that they can all be
@@ -1584,7 +1588,7 @@ loop:
 	}
 
 	// Update the commitment tx.
-	if err := l.updateCommitTx(); err != nil {
+	if err := l.updateCommitTx(ctx); err != nil {
 		return err
 	}
 
@@ -1671,7 +1675,9 @@ func (l *channelLink) randomFeeUpdateTimeout() time.Duration {
 
 // handleDownstreamUpdateAdd processes an UpdateAddHTLC packet sent from the
 // downstream HTLC Switch.
-func (l *channelLink) handleDownstreamUpdateAdd(pkt *htlcPacket) error {
+func (l *channelLink) handleDownstreamUpdateAdd(ctx context.Context,
+	pkt *htlcPacket) error {
+
 	htlc, ok := pkt.htlc.(*lnwire.UpdateAddHTLC)
 	if !ok {
 		return errors.New("not an UpdateAddHTLC packet")
@@ -1775,7 +1781,7 @@ func (l *channelLink) handleDownstreamUpdateAdd(pkt *htlcPacket) error {
 		getEventType(pkt),
 	)
 
-	l.tryBatchUpdateCommitTx()
+	l.tryBatchUpdateCommitTx(ctx)
 
 	return nil
 }
@@ -1786,7 +1792,9 @@ func (l *channelLink) handleDownstreamUpdateAdd(pkt *htlcPacket) error {
 // cleared HTLCs with the upstream peer.
 //
 // TODO(roasbeef): add sync ntfn to ensure switch always has consistent view?
-func (l *channelLink) handleDownstreamPkt(pkt *htlcPacket) {
+func (l *channelLink) handleDownstreamPkt(ctx context.Context,
+	pkt *htlcPacket) {
+
 	if pkt.htlc.MsgType().IsChannelUpdate() &&
 		!l.quiescer.CanSendUpdates() {
 
@@ -1800,7 +1808,7 @@ func (l *channelLink) handleDownstreamPkt(pkt *htlcPacket) {
 	case *lnwire.UpdateAddHTLC:
 		// Handle add message. The returned error can be ignored,
 		// because it is also sent through the mailbox.
-		_ = l.handleDownstreamUpdateAdd(pkt)
+		_ = l.handleDownstreamUpdateAdd(ctx, pkt)
 
 	case *lnwire.UpdateFulfillHTLC:
 		// If hodl.SettleOutgoing mode is active, we exit early to
@@ -1867,7 +1875,7 @@ func (l *channelLink) handleDownstreamPkt(pkt *htlcPacket) {
 		)
 
 		// Immediately update the commitment tx to minimize latency.
-		l.updateCommitTxOrFail()
+		l.updateCommitTxOrFail(ctx)
 
 	case *lnwire.UpdateFailHTLC:
 		// If hodl.FailOutgoing mode is active, we exit early to
@@ -1957,19 +1965,19 @@ func (l *channelLink) handleDownstreamPkt(pkt *htlcPacket) {
 		}
 
 		// Immediately update the commitment tx to minimize latency.
-		l.updateCommitTxOrFail()
+		l.updateCommitTxOrFail(ctx)
 	}
 }
 
 // tryBatchUpdateCommitTx updates the commitment transaction if the batch is
 // full.
-func (l *channelLink) tryBatchUpdateCommitTx() {
+func (l *channelLink) tryBatchUpdateCommitTx(ctx context.Context) {
 	pending := l.channel.NumPendingUpdates(lntypes.Local, lntypes.Remote)
 	if pending < uint64(l.cfg.BatchSize) {
 		return
 	}
 
-	l.updateCommitTxOrFail()
+	l.updateCommitTxOrFail(ctx)
 }
 
 // cleanupSpuriousResponse attempts to ack any AddRef or SettleFailRef
@@ -2039,7 +2047,11 @@ func (l *channelLink) cleanupSpuriousResponse(pkt *htlcPacket) {
 // handleUpstreamMsg processes wire messages related to commitment state
 // updates from the upstream peer. The upstream peer is the peer whom we have a
 // direct channel with, updating our respective commitment chains.
-func (l *channelLink) handleUpstreamMsg(msg lnwire.Message) {
+//
+//nolint:funlen
+func (l *channelLink) handleUpstreamMsg(ctx context.Context,
+	msg lnwire.Message) {
+
 	// First check if the message is an update and we are capable of
 	// receiving updates right now.
 	if msg.MsgType().IsChannelUpdate() && !l.quiescer.CanRecvUpdates() {
@@ -2418,7 +2430,7 @@ func (l *channelLink) handleUpstreamMsg(msg lnwire.Message) {
 		}
 
 		select {
-		case <-l.Quit:
+		case <-l.cg.Done():
 			return
 		default:
 		}
@@ -2430,7 +2442,7 @@ func (l *channelLink) handleUpstreamMsg(msg lnwire.Message) {
 		// reply with a signature as both sides already have a
 		// commitment with the latest accepted.
 		if l.channel.OweCommitment() {
-			if !l.updateCommitTxOrFail() {
+			if !l.updateCommitTxOrFail(ctx) {
 				return
 			}
 		}
@@ -2488,7 +2500,7 @@ func (l *channelLink) handleUpstreamMsg(msg lnwire.Message) {
 		}
 
 		select {
-		case <-l.Quit:
+		case <-l.cg.Done():
 			return
 		default:
 		}
@@ -2542,7 +2554,7 @@ func (l *channelLink) handleUpstreamMsg(msg lnwire.Message) {
 		// but there are still remote updates that are not in the remote
 		// commit tx yet, send out an update.
 		if l.channel.OweCommitment() {
-			if !l.updateCommitTxOrFail() {
+			if !l.updateCommitTxOrFail(ctx) {
 				return
 			}
 		}
@@ -2732,8 +2744,8 @@ func (l *channelLink) ackDownStreamPackets() error {
 
 // updateCommitTxOrFail updates the commitment tx and if that fails, it fails
 // the link.
-func (l *channelLink) updateCommitTxOrFail() bool {
-	err := l.updateCommitTx()
+func (l *channelLink) updateCommitTxOrFail(ctx context.Context) bool {
+	err := l.updateCommitTx(ctx)
 	switch err {
 	// No error encountered, success.
 	case nil:
@@ -2759,7 +2771,7 @@ func (l *channelLink) updateCommitTxOrFail() bool {
 // updateCommitTx signs, then sends an update to the remote peer adding a new
 // commitment to their commitment chain which includes all the latest updates
 // we've received+processed up to this point.
-func (l *channelLink) updateCommitTx() error {
+func (l *channelLink) updateCommitTx(ctx context.Context) error {
 	// Preemptively write all pending keystones to disk, just in case the
 	// HTLCs we have in memory are included in the subsequent attempt to
 	// sign a commitment state.
@@ -2782,7 +2794,7 @@ func (l *channelLink) updateCommitTx() error {
 		return nil
 	}
 
-	ctx, done := l.WithCtxQuitNoTimeout()
+	ctx, done := l.cg.Create(ctx)
 	defer done()
 
 	newCommit, err := l.channel.SignNextCommitment(ctx)
@@ -2822,7 +2834,7 @@ func (l *channelLink) updateCommitTx() error {
 	}
 
 	select {
-	case <-l.Quit:
+	case <-l.cg.Done():
 		return ErrLinkShuttingDown
 	default:
 	}
@@ -3529,7 +3541,7 @@ func (l *channelLink) handleSwitchPacket(pkt *htlcPacket) error {
 // NOTE: Part of the ChannelLink interface.
 func (l *channelLink) HandleChannelUpdate(message lnwire.Message) {
 	select {
-	case <-l.Quit:
+	case <-l.cg.Done():
 		// Return early if the link is already in the process of
 		// quitting. It doesn't make sense to hand the message to the
 		// mailbox here.
@@ -3545,7 +3557,9 @@ func (l *channelLink) HandleChannelUpdate(message lnwire.Message) {
 
 // updateChannelFee updates the commitment fee-per-kw on this channel by
 // committing to an update_fee message.
-func (l *channelLink) updateChannelFee(feePerKw chainfee.SatPerKWeight) error {
+func (l *channelLink) updateChannelFee(ctx context.Context,
+	feePerKw chainfee.SatPerKWeight) error {
+
 	l.log.Infof("updating commit fee to %v", feePerKw)
 
 	// We skip sending the UpdateFee message if the channel is not
@@ -3583,7 +3597,8 @@ func (l *channelLink) updateChannelFee(feePerKw chainfee.SatPerKWeight) error {
 	if err := l.cfg.Peer.SendMessage(false, msg); err != nil {
 		return err
 	}
-	return l.updateCommitTx()
+
+	return l.updateCommitTx(ctx)
 }
 
 // processRemoteSettleFails accepts a batch of settle/fail payment descriptors
@@ -4290,7 +4305,7 @@ func (l *channelLink) forwardBatch(replay bool, packets ...*htlcPacket) {
 		filteredPkts = append(filteredPkts, pkt)
 	}
 
-	err := l.cfg.ForwardPackets(l.Quit, replay, filteredPkts...)
+	err := l.cfg.ForwardPackets(l.cg.Done(), replay, filteredPkts...)
 	if err != nil {
 		log.Errorf("Unhandled error while reforwarding htlc "+
 			"settle/fail over htlcswitch: %v", err)
diff --git a/protofsm/state_machine.go b/protofsm/state_machine.go
index 45267e096..d78563e95 100644
--- a/protofsm/state_machine.go
+++ b/protofsm/state_machine.go
@@ -145,7 +145,7 @@ type StateMachine[Event any, Env Environment] struct {
 	// query the internal state machine state.
 	stateQuery chan stateQuery[Event, Env]
 
-	wg   *fn.GoroutineManager
+	gm   fn.GoroutineManager
 	quit chan struct{}
 
 	startOnce sync.Once
@@ -206,7 +206,7 @@ func NewStateMachine[Event any, Env Environment](
 		),
 		events:         make(chan Event, 1),
 		stateQuery:     make(chan stateQuery[Event, Env]),
-		wg:             fn.NewGoroutineManager(),
+		gm:             *fn.NewGoroutineManager(),
 		newStateEvents: fn.NewEventDistributor[State[Event, Env]](),
 		quit:           make(chan struct{}),
 	}
@@ -216,7 +216,7 @@ func NewStateMachine[Event any, Env Environment](
 // the state machine to completion.
 func (s *StateMachine[Event, Env]) Start(ctx context.Context) {
 	s.startOnce.Do(func() {
-		_ = s.wg.Go(ctx, func(ctx context.Context) {
+		_ = s.gm.Go(ctx, func(ctx context.Context) {
 			s.driveMachine(ctx)
 		})
 	})
@@ -227,7 +227,7 @@ func (s *StateMachine[Event, Env]) Start(ctx context.Context) {
 func (s *StateMachine[Event, Env]) Stop() {
 	s.stopOnce.Do(func() {
 		close(s.quit)
-		s.wg.Stop()
+		s.gm.Stop()
 	})
 }
 
@@ -335,8 +335,6 @@ func (s *StateMachine[Event, Env]) RemoveStateSub(sub StateSubscriber[
 // executeDaemonEvent executes a daemon event, which is a special type of event
 // that can be emitted as part of the state transition function of the state
 // machine. An error is returned if the type of event is unknown.
-//
-//nolint:funlen
 func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context,
 	event DaemonEvent) error {
 
@@ -360,7 +358,7 @@ func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context,
 			// If a post-send event was specified, then we'll funnel
 			// that back into the main state machine now as well.
 			return fn.MapOptionZ(daemonEvent.PostSendEvent, func(event Event) error { //nolint:ll
-				launched := s.wg.Go(
+				launched := s.gm.Go(
 					ctx, func(ctx context.Context) {
 						s.log.DebugS(ctx, "Sending post-send event",
 							"event", lnutils.SpewLogClosure(event))
@@ -386,7 +384,7 @@ func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context,
 		// Otherwise, this has a SendWhen predicate, so we'll need
 		// launch a goroutine to poll the SendWhen, then send only once
 		// the predicate is true.
-		launched := s.wg.Go(ctx, func(ctx context.Context) {
+		launched := s.gm.Go(ctx, func(ctx context.Context) {
 			predicateTicker := time.NewTicker(
 				s.cfg.CustomPollInterval.UnwrapOr(pollInterval),
 			)
@@ -456,7 +454,7 @@ func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context,
 			return fmt.Errorf("unable to register spend: %w", err)
 		}
 
-		launched := s.wg.Go(ctx, func(ctx context.Context) {
+		launched := s.gm.Go(ctx, func(ctx context.Context) {
 			for {
 				select {
 				case spend, ok := <-spendEvent.Spend:
@@ -502,7 +500,7 @@ func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context,
 			return fmt.Errorf("unable to register conf: %w", err)
 		}
 
-		launched := s.wg.Go(ctx, func(ctx context.Context) {
+		launched := s.gm.Go(ctx, func(ctx context.Context) {
 			for {
 				select {
 				case <-confEvent.Confirmed:
@@ -674,7 +672,7 @@ func (s *StateMachine[Event, Env]) driveMachine(ctx context.Context) {
 				return
 			}
 
-		case <-s.wg.Done():
+		case <-s.gm.Done():
 			return
 		}
 	}


Test patch:
diff --git a/htlcswitch/link_test.go b/htlcswitch/link_test.go
index 174710559..cdeae6d81 100644
--- a/htlcswitch/link_test.go
+++ b/htlcswitch/link_test.go
@@ -2257,7 +2257,7 @@ func newSingleLinkTestHarness(t *testing.T, chanAmt,
 			for {
 				select {
 				case <-notifyUpdateChan:
-				case <-chanLink.Quit:
+				case <-chanLink.cg.Done():
 					close(doneChan)
 					return
 				}
@@ -2326,7 +2326,7 @@ func handleStateUpdate(link *channelLink,
 	}
 	link.HandleChannelUpdate(remoteRev)
 
-	ctx, done := link.WithCtxQuitNoTimeout()
+	ctx, done := link.cg.Create(context.Background())
 	defer done()
 
 	remoteSigs, err := remoteChannel.SignNextCommitment(ctx)
@@ -2372,7 +2372,7 @@ func updateState(batchTick chan time.Time, link *channelLink,
 		// Trigger update by ticking the batchTicker.
 		select {
 		case batchTick <- time.Now():
-		case <-link.Quit:
+		case <-link.cg.Done():
 			return fmt.Errorf("link shutting down")
 		}
 		return handleStateUpdate(link, remoteChannel)
@@ -2380,7 +2380,7 @@ func updateState(batchTick chan time.Time, link *channelLink,
 
 	// The remote is triggering the state update, emulate this by
 	// signing and sending CommitSig to the link.
-	ctx, done := link.WithCtxQuitNoTimeout()
+	ctx, done := link.cg.Create(context.Background())
 	defer done()
 
 	remoteSigs, err := remoteChannel.SignNextCommitment(ctx)
@@ -4946,7 +4946,7 @@ func (h *persistentLinkHarness) restartLink(
 			for {
 				select {
 				case <-notifyUpdateChan:
-				case <-chanLink.Quit:
+				case <-chanLink.cg.Done():
 					close(doneChan)
 					return
 				}
@@ -5932,7 +5932,9 @@ func TestChannelLinkFail(t *testing.T) {
 
 				// Sign a commitment that will include
 				// signature for the HTLC just sent.
-				quitCtx, done := c.WithCtxQuitNoTimeout()
+				quitCtx, done := c.cg.Create(
+					context.Background(),
+				)
 				defer done()
 
 				sigs, err := remoteChannel.SignNextCommitment(
@@ -5979,7 +5981,9 @@ func TestChannelLinkFail(t *testing.T) {
 
 				// Sign a commitment that will include
 				// signature for the HTLC just sent.
-				quitCtx, done := c.WithCtxQuitNoTimeout()
+				quitCtx, done := c.cg.Create(
+					context.Background(),
+				)
 				defer done()
 
 				sigs, err := remoteChannel.SignNextCommitment(
diff --git a/htlcswitch/test_utils.go b/htlcswitch/test_utils.go
index 212346588..43902dc33 100644
--- a/htlcswitch/test_utils.go
+++ b/htlcswitch/test_utils.go
@@ -1190,7 +1190,7 @@ func (h *hopNetwork) createChannelLink(server, peer *mockServer,
 			for {
 				select {
 				case <-notifyUpdateChan:
-				case <-chanLink.Quit:
+				case <-chanLink.cg.Done():
 					close(doneChan)
 					return
 				}
diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go
index 1ff0217d9..bda498683 100644
--- a/protofsm/state_machine_test.go
+++ b/protofsm/state_machine_test.go
@@ -1,6 +1,7 @@
 package protofsm
 
 import (
+	"context"
 	"encoding/hex"
 	"fmt"
 	"sync/atomic"
@@ -14,7 +15,6 @@ import (
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/net/context"
 )
 
 type dummyEvents interface {