graph/db: thread context through ForEachNodeDirectedChannel

graph/db/graph.go

@@ -217,8 +217,9 @@ func (c *ChannelGraph) populateCache(ctx context.Context) error {
 // Unknown policies are passed into the callback as nil values.
 //
 // NOTE: this is part of the graphdb.NodeTraverser interface.
-func (c *ChannelGraph) ForEachNodeDirectedChannel(node route.Vertex,
-	cb func(channel *DirectedChannel) error, reset func()) error {
+func (c *ChannelGraph) ForEachNodeDirectedChannel(ctx context.Context,
+	node route.Vertex, cb func(channel *DirectedChannel) error,
+	reset func()) error {
 
 	if c.graphCache != nil {
 		return c.graphCache.ForEachChannel(node, cb)
@@ -227,7 +228,9 @@ func (c *ChannelGraph) ForEachNodeDirectedChannel(node route.Vertex,
 	// TODO(elle): once the no-cache path needs to support
 	// pathfinding across gossip versions, this should iterate
 	// across all versions rather than defaulting to v1.
-	return c.db.ForEachNodeDirectedChannel(gossipV1, node, cb, reset)
+	return c.db.ForEachNodeDirectedChannel(
+		ctx, gossipV1, node, cb, reset,
+	)
 }
 
 // FetchNodeFeatures returns the features of the given node. If no features are

graph/db/graph_test.go

@@ -1875,7 +1875,7 @@ func testGraphTraversalCacheable(t *testing.T, v lnwire.GossipVersion) {
 		// Query the ChannelGraph which uses the cache to iterate
 		// through the channels for each node.
 		err = graph.ChannelGraph.ForEachNodeDirectedChannel(
-			node, func(d *DirectedChannel) error {
+			ctx, node, func(d *DirectedChannel) error {
 				delete(chanIndex, d.ChannelID)
 				return nil
 			}, func() {},
@@ -1884,7 +1884,7 @@ func testGraphTraversalCacheable(t *testing.T, v lnwire.GossipVersion) {
 
 		// Now skip the cache and query the DB directly.
 		err = graph.db.ForEachNodeDirectedChannel(
-			v, node, func(d *DirectedChannel) error {
+			ctx, v, node, func(d *DirectedChannel) error {
 				delete(chanIndex2, d.ChannelID)
 				return nil
 			}, func() {},
@@ -4932,7 +4932,7 @@ func testGraphCacheForEachNodeChannel(t *testing.T,
 	getSingleChannel := func() *DirectedChannel {
 		var ch *DirectedChannel
 		err := graph.db.ForEachNodeDirectedChannel(
-			v, node1.PubKeyBytes,
+			ctx, v, node1.PubKeyBytes,
 			func(c *DirectedChannel) error {
 				require.Nil(t, ch)
 				ch = c

graph/db/interfaces.go

@@ -21,7 +21,7 @@ import (
 type NodeTraverser interface {
 	// ForEachNodeDirectedChannel calls the callback for every channel of
 	// the given node.
-	ForEachNodeDirectedChannel(nodePub route.Vertex,
+	ForEachNodeDirectedChannel(ctx context.Context, nodePub route.Vertex,
 		cb func(channel *DirectedChannel) error, reset func()) error
 
 	// FetchNodeFeatures returns the features of the given node.
@@ -34,8 +34,9 @@ type NodeTraverser interface {
 type Store interface { //nolint:interfacebloat
 	// ForEachNodeDirectedChannel calls the callback for every channel of
 	// the given node.
-	ForEachNodeDirectedChannel(v lnwire.GossipVersion, nodePub route.Vertex,
-		cb func(channel *DirectedChannel) error, reset func()) error
+	ForEachNodeDirectedChannel(ctx context.Context, v lnwire.GossipVersion,
+		nodePub route.Vertex, cb func(channel *DirectedChannel) error,
+		reset func()) error
 
 	// FetchNodeFeatures returns the features of the given node.
 	FetchNodeFeatures(ctx context.Context, v lnwire.GossipVersion,

graph/db/kv_store.go

@@ -669,9 +669,9 @@ func (c *KVStore) fetchNodeFeatures(tx kvdb.RTx,
 // Unknown policies are passed into the callback as nil values.
 //
 // NOTE: this is part of the graphdb.NodeTraverser interface.
-func (c *KVStore) ForEachNodeDirectedChannel(v lnwire.GossipVersion,
-	nodePub route.Vertex, cb func(channel *DirectedChannel) error,
-	reset func()) error {
+func (c *KVStore) ForEachNodeDirectedChannel(_ context.Context,
+	v lnwire.GossipVersion, nodePub route.Vertex,
+	cb func(channel *DirectedChannel) error, reset func()) error {
 
 	if v != lnwire.GossipVersion1 {
 		return ErrVersionNotSupportedForKVDB
@@ -4501,8 +4501,8 @@ type nodeTraverserSession struct {
 //
 // NOTE: Part of the NodeTraverser interface.
 func (c *nodeTraverserSession) ForEachNodeDirectedChannel(
-	nodePub route.Vertex, cb func(channel *DirectedChannel) error,
-	_ func()) error {
+	_ context.Context, nodePub route.Vertex,
+	cb func(channel *DirectedChannel) error, _ func()) error {
 
 	return c.db.forEachNodeDirectedChannel(c.tx, nodePub, cb, func() {})
 }

graph/db/sql_store.go

@@ -1006,11 +1006,9 @@ func (s *SQLStore) ForEachNode(ctx context.Context,
 // Unknown policies are passed into the callback as nil values.
 //
 // NOTE: this is part of the graphdb.NodeTraverser interface.
-func (s *SQLStore) ForEachNodeDirectedChannel(v lnwire.GossipVersion,
-	nodePub route.Vertex, cb func(channel *DirectedChannel) error,
-	reset func()) error {
-
-	var ctx = context.TODO()
+func (s *SQLStore) ForEachNodeDirectedChannel(ctx context.Context,
+	v lnwire.GossipVersion, nodePub route.Vertex,
+	cb func(channel *DirectedChannel) error, reset func()) error {
 
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
 		return forEachNodeDirectedChannel(ctx, db, v, nodePub, cb)
@@ -3375,10 +3373,8 @@ func newSQLNodeTraverser(db SQLQueries,
 //
 // NOTE: Part of the NodeTraverser interface.
 func (s *sqlNodeTraverser) ForEachNodeDirectedChannel(
-	nodePub route.Vertex, cb func(channel *DirectedChannel) error,
-	_ func()) error {
-
-	ctx := context.TODO()
+	ctx context.Context, nodePub route.Vertex,
+	cb func(channel *DirectedChannel) error, _ func()) error {
 
 	return forEachNodeDirectedChannel(
 		ctx, s.db, lnwire.GossipVersion1, nodePub, cb,

routing/bandwidth.go

@@ -1,6 +1,7 @@
 package routing
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/lightningnetwork/lnd/fn/v2"
@@ -64,7 +65,8 @@ func newBandwidthManager(graph Graph, sourceNode route.Vertex,
 	// First, we'll collect the set of outbound edges from the target
 	// source node and add them to our bandwidth manager's map of channels.
 	err := graph.ForEachNodeDirectedChannel(
-		sourceNode, func(channel *graphdb.DirectedChannel) error {
+		context.TODO(), sourceNode,
+		func(channel *graphdb.DirectedChannel) error {
 			shortID := lnwire.NewShortChanIDFromInt(
 				channel.ChannelID,
 			)

routing/graph.go

@@ -15,7 +15,7 @@ import (
 type Graph interface {
 	// ForEachNodeDirectedChannel calls the callback for every channel of
 	// the given node.
-	ForEachNodeDirectedChannel(nodePub route.Vertex,
+	ForEachNodeDirectedChannel(ctx context.Context, nodePub route.Vertex,
 		cb func(channel *graphdb.DirectedChannel) error,
 		reset func()) error
 

routing/mock_graph_test.go

@@ -166,8 +166,9 @@ func (m *mockGraph) addChannel(id uint64, node1id, node2id byte,
 // forEachNodeChannel calls the callback for every channel of the given node.
 //
 // NOTE: Part of the Graph interface.
-func (m *mockGraph) ForEachNodeDirectedChannel(nodePub route.Vertex,
-	cb func(channel *graphdb.DirectedChannel) error, _ func()) error {
+func (m *mockGraph) ForEachNodeDirectedChannel(_ context.Context,
+	nodePub route.Vertex, cb func(channel *graphdb.DirectedChannel) error,
+	_ func()) error {
 
 	// Look up the mock node.
 	node, ok := m.nodes[nodePub]

routing/pathfind.go

@@ -578,7 +578,7 @@ func getOutgoingBalance(node route.Vertex, outgoingChans map[uint64]struct{},
 
 	// Iterate over all channels of the to node.
 	err := g.ForEachNodeDirectedChannel(
-		node, cb, func() {
+		context.TODO(), node, cb, func() {
 			max = 0
 			total = 0
 		},
@@ -1291,7 +1291,7 @@ func findBlindedPaths(g Graph, target route.Vertex,
 			nextTargetReset   = nextTarget
 		)
 		err := g.ForEachNodeDirectedChannel(
-			nextTarget,
+			context.TODO(), nextTarget,
 			func(channel *graphdb.DirectedChannel) error {
 				// This is not the right channel, continue to
 				// the node's other channels.
@@ -1466,7 +1466,8 @@ func processNodeForBlindedPath(g Graph, node route.Vertex,
 	// Now, iterate over the node's channels in search for paths to this
 	// node that can be used for blinded paths
 	err = g.ForEachNodeDirectedChannel(
-		node, func(channel *graphdb.DirectedChannel) error {
+		context.TODO(), node,
+		func(channel *graphdb.DirectedChannel) error {
 			// Keep track of how many incoming channels this node
 			// has. We only use a node as an introduction node if it
 			// has channels other than the one that lead us to it.

routing/unified_edges.go

@@ -1,6 +1,7 @@
 package routing
 
 import (
+	"context"
 	"math"
 
 	"github.com/btcsuite/btcd/btcutil"
@@ -116,7 +117,7 @@ func (u *nodeEdgeUnifier) addGraphPolicies(g Graph) error {
 
 	// Iterate over all channels of the to node.
 	err := g.ForEachNodeDirectedChannel(
-		u.toNode, cb, func() {
+		context.TODO(), u.toNode, cb, func() {
 			channels = nil
 		},
 	)