graph/db: let ForEachNodeCacheable take a context

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)

graph/db/graph_test.go

@@ -1433,6 +1433,7 @@ func TestGraphTraversal(t *testing.T) {
 // working correctly.
 func TestGraphTraversalCacheable(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	graph := MakeTestGraph(t)
 
@@ -1458,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)
@@ -4226,7 +4227,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)

graph/db/interfaces.go

@@ -107,7 +107,7 @@ type V1Store interface { //nolint:interfacebloat
 	// 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

graph/db/kv_store.go

@@ -842,8 +842,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,
+func (c *KVStore) ForEachNodeCacheable(_ context.Context,
-	*lnwire.FeatureVector) error) error {
+	cb func(route.Vertex, *lnwire.FeatureVector) error) error {
 
 	traversal := func(tx kvdb.RTx) error {
 		// First grab the nodes bucket which stores the mapping from

graph/db/sql_store.go

@@ -923,10 +923,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,
+func (s *SQLStore) ForEachNodeCacheable(ctx context.Context,
-	*lnwire.FeatureVector) error) error {
+	cb func(route.Vertex, *lnwire.FeatureVector) error) error {
-
-	ctx := context.TODO()
 
 	err := s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
 		return forEachNodeCacheable(ctx, db, func(nodeID int64,