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,