Merge pull request #10043 from ellemouton/ctx3

multi: add context.Context param to more graphdb.V1Store methods

graph/db/graph.go

@@ -75,7 +75,7 @@ func (c *ChannelGraph) Start() error {
 	defer log.Debug("ChannelGraph started")
 
 	if c.graphCache != nil {
-		if err := c.populateCache(); err != nil {
+		if err := c.populateCache(context.TODO()); err != nil {
 			return fmt.Errorf("could not populate the graph "+
 				"cache: %w", err)
 		}
@@ -159,12 +159,12 @@ func (c *ChannelGraph) handleTopologySubscriptions() {
 // populateCache loads the entire channel graph into the in-memory graph cache.
 //
 // NOTE: This should only be called if the graphCache has been constructed.
-func (c *ChannelGraph) populateCache() error {
+func (c *ChannelGraph) populateCache(ctx context.Context) error {
 	startTime := time.Now()
 	log.Info("Populating in-memory channel graph, this might take a " +
 		"while...")
 
-	err := c.V1Store.ForEachNodeCacheable(func(node route.Vertex,
+	err := c.V1Store.ForEachNodeCacheable(ctx, func(node route.Vertex,
 		features *lnwire.FeatureVector) error {
 
 		c.graphCache.AddNodeFeatures(node, features)
@@ -245,14 +245,15 @@ func (c *ChannelGraph) GraphSession(cb func(graph NodeTraverser) error) error {
 // graph, executing the passed callback with each node encountered.
 //
 // NOTE: The callback contents MUST not be modified.
-func (c *ChannelGraph) ForEachNodeCached(cb func(node route.Vertex,
-	chans map[uint64]*DirectedChannel) error) error {
+func (c *ChannelGraph) ForEachNodeCached(ctx context.Context,
+	cb func(node route.Vertex,
+		chans map[uint64]*DirectedChannel) error) error {
 
 	if c.graphCache != nil {
 		return c.graphCache.ForEachNode(cb)
 	}
 
-	return c.V1Store.ForEachNodeCached(cb)
+	return c.V1Store.ForEachNodeCached(ctx, cb)
 }
 
 // AddLightningNode adds a vertex/node to the graph database. If the node is not

graph/db/graph_test.go

@@ -922,15 +922,20 @@ func TestEdgePolicyCRUD(t *testing.T) {
 		// Use the ForEachChannel method to fetch the policies and
 		// assert that the deserialized policies match the original
 		// ones.
-		err := graph.ForEachChannel(func(info *models.ChannelEdgeInfo,
-			policy1 *models.ChannelEdgePolicy,
-			policy2 *models.ChannelEdgePolicy) error {
+		err := graph.ForEachChannel(
+			ctx, func(info *models.ChannelEdgeInfo,
+				policy1 *models.ChannelEdgePolicy,
+				policy2 *models.ChannelEdgePolicy) error {
 
-			require.NoError(t, compareEdgePolicies(edge1, policy1))
-			require.NoError(t, compareEdgePolicies(edge2, policy2))
+				require.NoError(
+					t, compareEdgePolicies(edge1, policy1),
+				)
+				require.NoError(
+					t, compareEdgePolicies(edge2, policy2),
+				)
 
-			return nil
-		})
+				return nil
+			})
 		require.NoError(t, err)
 	}
 
@@ -1311,7 +1316,7 @@ func TestForEachSourceNodeChannel(t *testing.T) {
 
 	// Now, we'll use the ForEachSourceNodeChannel and assert that it
 	// returns the expected data in the call-back.
-	err := graph.ForEachSourceNodeChannel(func(chanPoint wire.OutPoint,
+	err := graph.ForEachSourceNodeChannel(ctx, func(chanPoint wire.OutPoint,
 		havePolicy bool, otherNode *models.LightningNode) error {
 
 		require.Contains(t, expectedSrcChans, chanPoint)
@@ -1332,6 +1337,7 @@ func TestForEachSourceNodeChannel(t *testing.T) {
 
 func TestGraphTraversal(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	graph := MakeTestGraph(t)
 
@@ -1352,7 +1358,7 @@ func TestGraphTraversal(t *testing.T) {
 	// set of channels (to force the fall back), we should find all the
 	// channel as well as the nodes included.
 	graph.graphCache = nil
-	err := graph.ForEachNodeCached(func(node route.Vertex,
+	err := graph.ForEachNodeCached(ctx, func(node route.Vertex,
 		chans map[uint64]*DirectedChannel) error {
 
 		if _, ok := nodeIndex[node]; !ok {
@@ -1373,7 +1379,7 @@ func TestGraphTraversal(t *testing.T) {
 	// Iterate through all the known channels within the graph DB, once
 	// again if the map is empty that indicates that all edges have
 	// properly been reached.
-	err = graph.ForEachChannel(func(ei *models.ChannelEdgeInfo,
+	err = graph.ForEachChannel(ctx, func(ei *models.ChannelEdgeInfo,
 		_ *models.ChannelEdgePolicy,
 		_ *models.ChannelEdgePolicy) error {
 
@@ -1387,7 +1393,7 @@ func TestGraphTraversal(t *testing.T) {
 	// outgoing channels for a particular node.
 	numNodeChans := 0
 	firstNode, secondNode := nodeList[0], nodeList[1]
-	err = graph.ForEachNodeChannel(firstNode.PubKeyBytes,
+	err = graph.ForEachNodeChannel(ctx, firstNode.PubKeyBytes,
 		func(_ *models.ChannelEdgeInfo, outEdge,
 			inEdge *models.ChannelEdgePolicy) error {
 
@@ -1427,6 +1433,7 @@ func TestGraphTraversal(t *testing.T) {
 // working correctly.
 func TestGraphTraversalCacheable(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	graph := MakeTestGraph(t)
 
@@ -1440,7 +1447,7 @@ func TestGraphTraversalCacheable(t *testing.T) {
 	// Create a map of all nodes with the iteration we know works (because
 	// it is tested in another test).
 	nodeMap := make(map[route.Vertex]struct{})
-	err := graph.ForEachNode(func(tx NodeRTx) error {
+	err := graph.ForEachNode(ctx, func(tx NodeRTx) error {
 		nodeMap[tx.Node().PubKeyBytes] = struct{}{}
 
 		return nil
@@ -1452,7 +1459,7 @@ func TestGraphTraversalCacheable(t *testing.T) {
 	// iterating over each node, once again if the map is empty that
 	// indicates that all edges have properly been reached.
 	var nodes []route.Vertex
-	err = graph.ForEachNodeCacheable(func(node route.Vertex,
+	err = graph.ForEachNodeCacheable(ctx, func(node route.Vertex,
 		features *lnwire.FeatureVector) error {
 
 		delete(nodeMap, node)
@@ -1570,7 +1577,7 @@ func fillTestGraph(t testing.TB, graph *ChannelGraph, numNodes,
 
 	// Iterate over each node as returned by the graph, if all nodes are
 	// reached, then the map created above should be empty.
-	err := graph.ForEachNode(func(tx NodeRTx) error {
+	err := graph.ForEachNode(ctx, func(tx NodeRTx) error {
 		delete(nodeIndex, tx.Node().Alias)
 		return nil
 	})
@@ -1662,13 +1669,14 @@ func assertPruneTip(t *testing.T, graph *ChannelGraph,
 
 func assertNumChans(t *testing.T, graph *ChannelGraph, n int) {
 	numChans := 0
-	if err := graph.ForEachChannel(func(*models.ChannelEdgeInfo,
-		*models.ChannelEdgePolicy,
-		*models.ChannelEdgePolicy) error {
+	if err := graph.ForEachChannel(context.Background(),
+		func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
+			*models.ChannelEdgePolicy) error {
 
-		numChans++
-		return nil
-	}); err != nil {
+			numChans++
+			return nil
+		},
+	); err != nil {
 		_, _, line, _ := runtime.Caller(1)
 		t.Fatalf("line %v: unable to scan channels: %v", line, err)
 	}
@@ -1681,7 +1689,7 @@ func assertNumChans(t *testing.T, graph *ChannelGraph, n int) {
 
 func assertNumNodes(t *testing.T, graph *ChannelGraph, n int) {
 	numNodes := 0
-	err := graph.ForEachNode(func(tx NodeRTx) error {
+	err := graph.ForEachNode(context.Background(), func(tx NodeRTx) error {
 		numNodes++
 
 		return nil
@@ -3138,7 +3146,7 @@ func TestIncompleteChannelPolicies(t *testing.T) {
 		expectedOut bool) {
 
 		calls := 0
-		err := graph.ForEachNodeChannel(node.PubKeyBytes,
+		err := graph.ForEachNodeChannel(ctx, node.PubKeyBytes,
 			func(_ *models.ChannelEdgeInfo, outEdge,
 				inEdge *models.ChannelEdgePolicy) error {
 
@@ -4201,6 +4209,7 @@ func TestBatchedUpdateEdgePolicy(t *testing.T) {
 // allocations and the total memory consumed by the full graph traversal.
 func BenchmarkForEachChannel(b *testing.B) {
 	graph := MakeTestGraph(b)
+	ctx := context.Background()
 
 	const numNodes = 100
 	const numChannels = 4
@@ -4215,7 +4224,7 @@ func BenchmarkForEachChannel(b *testing.B) {
 		)
 
 		var nodes []route.Vertex
-		err := graph.ForEachNodeCacheable(func(node route.Vertex,
+		err := graph.ForEachNodeCacheable(ctx, func(node route.Vertex,
 			vector *lnwire.FeatureVector) error {
 
 			nodes = append(nodes, node)
@@ -4241,7 +4250,7 @@ func BenchmarkForEachChannel(b *testing.B) {
 				return nil
 			}
 
-			err := graph.ForEachNodeChannel(n, cb)
+			err := graph.ForEachNodeChannel(ctx, n, cb)
 			require.NoError(b, err)
 		}
 	}

graph/db/interfaces.go

@@ -70,8 +70,9 @@ type V1Store interface { //nolint:interfacebloat
 	// node, executing the passed callback on each. The call-back is
 	// provided with the channel's outpoint, whether we have a policy for
 	// the channel and the channel peer's node information.
-	ForEachSourceNodeChannel(cb func(chanPoint wire.OutPoint,
-		havePolicy bool, otherNode *models.LightningNode) error) error
+	ForEachSourceNodeChannel(ctx context.Context,
+		cb func(chanPoint wire.OutPoint, havePolicy bool,
+			otherNode *models.LightningNode) error) error
 
 	// ForEachNodeChannel iterates through all channels of the given node,
 	// executing the passed callback with an edge info structure and the
@@ -82,7 +83,7 @@ type V1Store interface { //nolint:interfacebloat
 	// to the caller.
 	//
 	// Unknown policies are passed into the callback as nil values.
-	ForEachNodeChannel(nodePub route.Vertex,
+	ForEachNodeChannel(ctx context.Context, nodePub route.Vertex,
 		cb func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
 			*models.ChannelEdgePolicy) error) error
 
@@ -90,7 +91,7 @@ type V1Store interface { //nolint:interfacebloat
 	// DirectedChannel data to the call-back.
 	//
 	// NOTE: The callback contents MUST not be modified.
-	ForEachNodeCached(cb func(node route.Vertex,
+	ForEachNodeCached(ctx context.Context, cb func(node route.Vertex,
 		chans map[uint64]*DirectedChannel) error) error
 
 	// ForEachNode iterates through all the stored vertices/nodes in the
@@ -100,13 +101,13 @@ type V1Store interface { //nolint:interfacebloat
 	// passed to the call-back are executed under the same read transaction
 	// and so, methods on the NodeTx object _MUST_ only be called from
 	// within the call-back.
-	ForEachNode(cb func(tx NodeRTx) error) error
+	ForEachNode(ctx context.Context, cb func(tx NodeRTx) error) error
 
 	// ForEachNodeCacheable iterates through all the stored vertices/nodes
 	// in the graph, executing the passed callback with each node
 	// encountered. If the callback returns an error, then the transaction
 	// is aborted and the iteration stops early.
-	ForEachNodeCacheable(cb func(route.Vertex,
+	ForEachNodeCacheable(ctx context.Context, cb func(route.Vertex,
 		*lnwire.FeatureVector) error) error
 
 	// LookupAlias attempts to return the alias as advertised by the target
@@ -157,7 +158,7 @@ type V1Store interface { //nolint:interfacebloat
 	// NOTE: If an edge can't be found, or wasn't advertised, then a nil
 	// pointer for that particular channel edge routing policy will be
 	// passed into the callback.
-	ForEachChannel(cb func(*models.ChannelEdgeInfo,
+	ForEachChannel(ctx context.Context, cb func(*models.ChannelEdgeInfo,
 		*models.ChannelEdgePolicy,
 		*models.ChannelEdgePolicy) error) error
 

graph/db/kv_store.go

@@ -404,8 +404,9 @@ func (c *KVStore) AddrsForNode(ctx context.Context,
 // NOTE: If an edge can't be found, or wasn't advertised, then a nil pointer
 // for that particular channel edge routing policy will be passed into the
 // callback.
-func (c *KVStore) ForEachChannel(cb func(*models.ChannelEdgeInfo,
-	*models.ChannelEdgePolicy, *models.ChannelEdgePolicy) error) error {
+func (c *KVStore) ForEachChannel(_ context.Context,
+	cb func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
+		*models.ChannelEdgePolicy) error) error {
 
 	return forEachChannel(c.db, cb)
 }
@@ -666,8 +667,9 @@ func (c *KVStore) FetchNodeFeatures(nodePub route.Vertex) (
 // data to the call-back.
 //
 // NOTE: The callback contents MUST not be modified.
-func (c *KVStore) ForEachNodeCached(cb func(node route.Vertex,
-	chans map[uint64]*DirectedChannel) error) error {
+func (c *KVStore) ForEachNodeCached(_ context.Context,
+	cb func(node route.Vertex,
+		chans map[uint64]*DirectedChannel) error) error {
 
 	// Otherwise call back to a version that uses the database directly.
 	// We'll iterate over each node, then the set of channels for each
@@ -788,7 +790,9 @@ func (c *KVStore) DisabledChannelIDs() ([]uint64, error) {
 // early. Any operations performed on the NodeTx passed to the call-back are
 // executed under the same read transaction and so, methods on the NodeTx object
 // _MUST_ only be called from within the call-back.
-func (c *KVStore) ForEachNode(cb func(tx NodeRTx) error) error {
+func (c *KVStore) ForEachNode(_ context.Context,
+	cb func(tx NodeRTx) error) error {
+
 	return forEachNode(c.db, func(tx kvdb.RTx,
 		node *models.LightningNode) error {
 
@@ -841,8 +845,8 @@ func forEachNode(db kvdb.Backend,
 // graph, executing the passed callback with each node encountered. If the
 // callback returns an error, then the transaction is aborted and the iteration
 // stops early.
-func (c *KVStore) ForEachNodeCacheable(cb func(route.Vertex,
-	*lnwire.FeatureVector) error) error {
+func (c *KVStore) ForEachNodeCacheable(_ context.Context,
+	cb func(route.Vertex, *lnwire.FeatureVector) error) error {
 
 	traversal := func(tx kvdb.RTx) error {
 		// First grab the nodes bucket which stores the mapping from
@@ -3242,7 +3246,7 @@ func nodeTraversal(tx kvdb.RTx, nodePub []byte, db kvdb.Backend,
 // halted with the error propagated back up to the caller.
 //
 // Unknown policies are passed into the callback as nil values.
-func (c *KVStore) ForEachNodeChannel(nodePub route.Vertex,
+func (c *KVStore) ForEachNodeChannel(_ context.Context, nodePub route.Vertex,
 	cb func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
 		*models.ChannelEdgePolicy) error) error {
 
@@ -3258,8 +3262,9 @@ func (c *KVStore) ForEachNodeChannel(nodePub route.Vertex,
 // executing the passed callback on each. The callback is provided with the
 // channel's outpoint, whether we have a policy for the channel and the channel
 // peer's node information.
-func (c *KVStore) ForEachSourceNodeChannel(cb func(chanPoint wire.OutPoint,
-	havePolicy bool, otherNode *models.LightningNode) error) error {
+func (c *KVStore) ForEachSourceNodeChannel(_ context.Context,
+	cb func(chanPoint wire.OutPoint, havePolicy bool,
+		otherNode *models.LightningNode) error) error {
 
 	return kvdb.View(c.db, func(tx kvdb.RTx) error {
 		nodes := tx.ReadBucket(nodeBucket)

graph/db/sql_migration_test.go

@@ -283,7 +283,7 @@ func assertInSync(t *testing.T, kvDB *KVStore, sqlDB *SQLStore,
 func fetchAllNodes(t *testing.T, store V1Store) []*models.LightningNode {
 	nodes := make([]*models.LightningNode, 0)
 
-	err := store.ForEachNode(func(tx NodeRTx) error {
+	err := store.ForEachNode(context.Background(), func(tx NodeRTx) error {
 		node := tx.Node()
 
 		// Call PubKey to ensure the objects cached pubkey is set so that

graph/db/sql_store.go

@@ -724,10 +724,9 @@ func (s *SQLStore) updateEdgeCache(e *models.ChannelEdgePolicy,
 // peer's node information.
 //
 // NOTE: part of the V1Store interface.
-func (s *SQLStore) ForEachSourceNodeChannel(cb func(chanPoint wire.OutPoint,
-	havePolicy bool, otherNode *models.LightningNode) error) error {
-
-	var ctx = context.TODO()
+func (s *SQLStore) ForEachSourceNodeChannel(ctx context.Context,
+	cb func(chanPoint wire.OutPoint, havePolicy bool,
+		otherNode *models.LightningNode) error) error {
 
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
 		nodeID, nodePub, err := s.getSourceNode(ctx, db, ProtocolV1)
@@ -784,12 +783,10 @@ func (s *SQLStore) ForEachSourceNodeChannel(cb func(chanPoint wire.OutPoint,
 // _MUST_ only be called from within the call-back.
 //
 // NOTE: part of the V1Store interface.
-func (s *SQLStore) ForEachNode(cb func(tx NodeRTx) error) error {
-	var (
-		ctx          = context.TODO()
-		lastID int64 = 0
-	)
+func (s *SQLStore) ForEachNode(ctx context.Context,
+	cb func(tx NodeRTx) error) error {
 
+	var lastID int64 = 0
 	handleNode := func(db SQLQueries, dbNode sqlc.Node) error {
 		node, err := buildNode(ctx, db, &dbNode)
 		if err != nil {
@@ -924,10 +921,8 @@ func (s *SQLStore) ForEachNodeDirectedChannel(nodePub route.Vertex,
 // stops early.
 //
 // NOTE: This is a part of the V1Store interface.
-func (s *SQLStore) ForEachNodeCacheable(cb func(route.Vertex,
-	*lnwire.FeatureVector) error) error {
-
-	ctx := context.TODO()
+func (s *SQLStore) ForEachNodeCacheable(ctx context.Context,
+	cb func(route.Vertex, *lnwire.FeatureVector) error) error {
 
 	err := s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
 		return forEachNodeCacheable(ctx, db, func(nodeID int64,
@@ -959,12 +954,10 @@ func (s *SQLStore) ForEachNodeCacheable(cb func(route.Vertex,
 // Unknown policies are passed into the callback as nil values.
 //
 // NOTE: part of the V1Store interface.
-func (s *SQLStore) ForEachNodeChannel(nodePub route.Vertex,
+func (s *SQLStore) ForEachNodeChannel(ctx context.Context, nodePub route.Vertex,
 	cb func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
 		*models.ChannelEdgePolicy) error) error {
 
-	var ctx = context.TODO()
-
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
 		dbNode, err := db.GetNodeByPubKey(
 			ctx, sqlc.GetNodeByPubKeyParams{
@@ -1109,10 +1102,9 @@ func (s *SQLStore) ChanUpdatesInHorizon(startTime,
 // NOTE: The callback contents MUST not be modified.
 //
 // NOTE: part of the V1Store interface.
-func (s *SQLStore) ForEachNodeCached(cb func(node route.Vertex,
-	chans map[uint64]*DirectedChannel) error) error {
-
-	var ctx = context.TODO()
+func (s *SQLStore) ForEachNodeCached(ctx context.Context,
+	cb func(node route.Vertex,
+		chans map[uint64]*DirectedChannel) error) error {
 
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
 		return forEachNodeCacheable(ctx, db, func(nodeID int64,
@@ -1336,10 +1328,9 @@ func (s *SQLStore) ForEachChannelCacheable(cb func(*models.CachedEdgeInfo,
 // callback.
 //
 // NOTE: part of the V1Store interface.
-func (s *SQLStore) ForEachChannel(cb func(*models.ChannelEdgeInfo,
-	*models.ChannelEdgePolicy, *models.ChannelEdgePolicy) error) error {
-
-	ctx := context.TODO()
+func (s *SQLStore) ForEachChannel(ctx context.Context,
+	cb func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
+		*models.ChannelEdgePolicy) error) error {
 
 	handleChannel := func(db SQLQueries,
 		row sqlc.ListChannelsWithPoliciesPaginatedRow) error {