Merge pull request #10610 from ellemouton/graph-cleanup-todo-contexts

graph/db+refactor: surface all `context.TODO()`s

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

@@ -1,6 +1,7 @@
 package routing
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/btcsuite/btcd/btcutil"
@@ -14,12 +15,13 @@ 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
 
 	// FetchNodeFeatures returns the features of the given node.
-	FetchNodeFeatures(nodePub route.Vertex) (*lnwire.FeatureVector, error)
+	FetchNodeFeatures(ctx context.Context,
+		nodePub route.Vertex) (*lnwire.FeatureVector, error)
 }
 
 // GraphSessionFactory can be used to gain access to a graphdb.NodeTraverser
@@ -30,7 +32,8 @@ type GraphSessionFactory interface {
 	// GraphSession will provide the call-back with access to a
 	// graphdb.NodeTraverser instance which can be used to perform queries
 	// against the channel graph.
-	GraphSession(cb func(graph graphdb.NodeTraverser) error,
+	GraphSession(ctx context.Context,
+		cb func(graph graphdb.NodeTraverser) error,
 		reset func()) error
 }
 

routing/mock_graph_test.go

@@ -2,6 +2,7 @@ package routing
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"testing"
 
@@ -165,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]
@@ -221,8 +223,8 @@ func (m *mockGraph) sourceNode() route.Vertex {
 // fetchNodeFeatures returns the features of the given node.
 //
 // NOTE: Part of the Graph interface.
-func (m *mockGraph) FetchNodeFeatures(nodePub route.Vertex) (
-	*lnwire.FeatureVector, error) {
+func (m *mockGraph) FetchNodeFeatures(_ context.Context,
+	_ route.Vertex) (*lnwire.FeatureVector, error) {
 
 	return lnwire.EmptyFeatureVector(), nil
 }
@@ -232,8 +234,8 @@ func (m *mockGraph) FetchNodeFeatures(nodePub route.Vertex) (
 // the channel graph.
 //
 // NOTE: Part of the GraphSessionFactory interface.
-func (m *mockGraph) GraphSession(cb func(graph graphdb.NodeTraverser) error,
-	_ func()) error {
+func (m *mockGraph) GraphSession(_ context.Context,
+	cb func(graph graphdb.NodeTraverser) error, _ func()) error {
 
 	return cb(m)
 }

routing/pathfind.go

@@ -3,6 +3,7 @@ package routing
 import (
 	"bytes"
 	"container/heap"
+	"context"
 	"errors"
 	"fmt"
 	"math"
@@ -577,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
 		},
@@ -621,7 +622,9 @@ func findPath(g *graphParams, r *RestrictParams, cfg *PathFindingConfig,
 	features := r.DestFeatures
 	if features == nil {
 		var err error
-		features, err = g.graph.FetchNodeFeatures(target)
+		features, err = g.graph.FetchNodeFeatures(
+			context.TODO(), target,
+		)
 		if err != nil {
 			return nil, 0, err
 		}
@@ -1019,7 +1022,9 @@ func findPath(g *graphParams, r *RestrictParams, cfg *PathFindingConfig,
 		}
 
 		// Fetch node features fresh from the graph.
-		fromFeatures, err := g.graph.FetchNodeFeatures(node)
+		fromFeatures, err := g.graph.FetchNodeFeatures(
+			context.TODO(), node,
+		)
 		if err != nil {
 			return nil, err
 		}
@@ -1286,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.
@@ -1349,7 +1354,7 @@ func findBlindedPaths(g Graph, target route.Vertex,
 			return true, nil
 		}
 
-		features, err := g.FetchNodeFeatures(node)
+		features, err := g.FetchNodeFeatures(context.TODO(), node)
 		if err != nil {
 			return false, err
 		}
@@ -1461,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/pathfind_test.go

@@ -2199,7 +2199,9 @@ func runRouteFailMaxHTLC(t *testing.T, useCache bool) {
 	// Next, update the middle edge policy to only allow payments up to 100k
 	// msat.
 	graph := ctx.testGraphInstance.graph
-	_, midEdge, _, err := graph.FetchChannelEdgesByID(firstToSecondID)
+	_, midEdge, _, err := graph.FetchChannelEdgesByID(
+		t.Context(), firstToSecondID,
+	)
 	require.NoError(t, err, "unable to fetch channel edges by ID")
 	midEdge.MessageFlags = 1
 	midEdge.MaxHTLC = payAmt - 1
@@ -2243,7 +2245,9 @@ func runRouteFailDisabledEdge(t *testing.T, useCache bool) {
 	// path finding, as we don't consider the disable flag for local
 	// channels (and roasbeef is the source).
 	roasToPham := uint64(999991)
-	_, e1, e2, err := graph.graph.FetchChannelEdgesByID(roasToPham)
+	_, e1, e2, err := graph.graph.FetchChannelEdgesByID(
+		t.Context(), roasToPham,
+	)
 	require.NoError(t, err, "unable to fetch edge")
 	e1.ChannelFlags |= lnwire.ChanUpdateDisabled
 	e1.LastUpdate = e1.LastUpdate.Add(time.Second)
@@ -2266,7 +2270,9 @@ func runRouteFailDisabledEdge(t *testing.T, useCache bool) {
 	// Now, we'll modify the edge from phamnuwen -> sophon, to read that
 	// it's disabled.
 	phamToSophon := uint64(99999)
-	_, e, _, err := graph.graph.FetchChannelEdgesByID(phamToSophon)
+	_, e, _, err := graph.graph.FetchChannelEdgesByID(
+		t.Context(), phamToSophon,
+	)
 	require.NoError(t, err, "unable to fetch edge")
 	e.ChannelFlags |= lnwire.ChanUpdateDisabled
 	e.LastUpdate = e.LastUpdate.Add(time.Second)
@@ -2349,7 +2355,9 @@ func runPathSourceEdgesBandwidth(t *testing.T, useCache bool) {
 	// Finally, set the roasbeef->songoku bandwidth, but also set its
 	// disable flag.
 	bandwidths.hints[roasToSongoku] = 2 * payAmt
-	_, e1, e2, err := graph.graph.FetchChannelEdgesByID(roasToSongoku)
+	_, e1, e2, err := graph.graph.FetchChannelEdgesByID(
+		t.Context(), roasToSongoku,
+	)
 	require.NoError(t, err, "unable to fetch edge")
 	e1.ChannelFlags |= lnwire.ChanUpdateDisabled
 	e1.LastUpdate = e1.LastUpdate.Add(time.Second)
@@ -3318,7 +3326,7 @@ func dbFindPath(graph *graphdb.VersionedGraph,
 	}
 
 	var route []*unifiedEdge
-	err = graph.GraphSession(func(graph graphdb.NodeTraverser) error {
+	err = graph.GraphSession(ctx, func(graph graphdb.NodeTraverser) error {
 		route, _, err = findPath(
 			&graphParams{
 				additionalEdges: additionalEdges,

routing/payment_session.go

@@ -1,6 +1,7 @@
 package routing
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -344,6 +345,7 @@ func (p *paymentSession) RequestRoute(maxAmt, feeLimit lnwire.MilliSatoshi,
 
 	for {
 		err := p.graphSessFactory.GraphSession(
+			context.TODO(),
 			findPath, func() {
 				path = nil
 			},

routing/payment_session_test.go

@@ -1,6 +1,7 @@
 package routing
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -260,8 +261,8 @@ func (g *sessionGraph) sourceNode() route.Vertex {
 	return route.Vertex{}
 }
 
-func (g *sessionGraph) GraphSession(cb func(graph graphdb.NodeTraverser) error,
-	_ func()) error {
+func (g *sessionGraph) GraphSession(_ context.Context,
+	cb func(graph graphdb.NodeTraverser) error, _ func()) error {
 
 	return cb(g)
 }

routing/router_test.go

@@ -458,7 +458,7 @@ func TestChannelUpdateValidation(t *testing.T) {
 
 	// Assert that the initially configured fee is retrieved correctly.
 	_, e1, e2, err := ctx.graph.FetchChannelEdgesByID(
-		lnwire.NewShortChanIDFromInt(1).ToUint64(),
+		t.Context(), lnwire.NewShortChanIDFromInt(1).ToUint64(),
 	)
 	require.NoError(t, err, "cannot retrieve channel")
 
@@ -530,7 +530,7 @@ func TestChannelUpdateValidation(t *testing.T) {
 	require.Error(t, err, "expected route to fail with channel update")
 
 	_, e1, e2, err = ctx.graph.FetchChannelEdgesByID(
-		lnwire.NewShortChanIDFromInt(1).ToUint64(),
+		t.Context(), lnwire.NewShortChanIDFromInt(1).ToUint64(),
 	)
 	require.NoError(t, err, "cannot retrieve channel")
 
@@ -552,7 +552,7 @@ func TestChannelUpdateValidation(t *testing.T) {
 	// This time a valid signature was supplied and the policy change should
 	// have been applied to the graph.
 	_, e1, e2, err = ctx.graph.FetchChannelEdgesByID(
-		lnwire.NewShortChanIDFromInt(1).ToUint64(),
+		t.Context(), lnwire.NewShortChanIDFromInt(1).ToUint64(),
 	)
 	require.NoError(t, err, "cannot retrieve channel")
 
@@ -593,7 +593,7 @@ func TestSendPaymentErrorRepeatedFeeInsufficient(t *testing.T) {
 	// to sophon. We'll obtain this as we'll need to to generate the
 	// FeeInsufficient error that we'll send back.
 	_, _, edgeUpdateToFail, err := ctx.graph.FetchChannelEdgesByID(
-		songokuSophonChanID,
+		t.Context(), songokuSophonChanID,
 	)
 	require.NoError(t, err, "unable to fetch chan id")
 
@@ -946,7 +946,9 @@ func TestSendPaymentErrorNonFinalTimeLockErrors(t *testing.T) {
 	chanID := ctx.getChannelIDFromAlias(t, "roasbeef", "songoku")
 	roasbeefSongoku := lnwire.NewShortChanIDFromInt(chanID)
 
-	_, _, edgeUpdateToFail, err := ctx.graph.FetchChannelEdgesByID(chanID)
+	_, _, edgeUpdateToFail, err := ctx.graph.FetchChannelEdgesByID(
+		t.Context(), chanID,
+	)
 	require.NoError(t, err, "unable to fetch chan id")
 
 	errChanUpdate := lnwire.ChannelUpdate1{

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
 		},
 	)