From 3d5e852c8cb5d5f622a44db11bbddd2576f846b1 Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Mon, 7 Jul 2025 09:48:39 +0200 Subject: [PATCH] graph/db: let ForEachNodeCached take a context --- autopilot/graph.go | 2 +- autopilot/interface.go | 2 +- graph/db/graph.go | 7 ++++--- graph/db/graph_test.go | 2 +- graph/db/interfaces.go | 2 +- graph/db/kv_store.go | 5 +++-- graph/db/sql_store.go | 7 +++---- rpcserver.go | 2 +- 8 files changed, 15 insertions(+), 14 deletions(-) diff --git a/autopilot/graph.go b/autopilot/graph.go index d50871591..3aa3ad055 100644 --- a/autopilot/graph.go +++ b/autopilot/graph.go @@ -219,7 +219,7 @@ func (nc dbNodeCached) ForEachChannel(ctx context.Context, func (dc *databaseChannelGraphCached) ForEachNode(ctx context.Context, cb func(context.Context, Node) error) error { - return dc.db.ForEachNodeCached(func(n route.Vertex, + return dc.db.ForEachNodeCached(ctx, func(n route.Vertex, channels map[uint64]*graphdb.DirectedChannel) error { if len(channels) > 0 { diff --git a/autopilot/interface.go b/autopilot/interface.go index 594f2ece5..703f49020 100644 --- a/autopilot/interface.go +++ b/autopilot/interface.go @@ -234,6 +234,6 @@ type GraphSource interface { // channel graph cache if one is available. It is less consistent than // ForEachNode since any further calls are made across multiple // transactions. - ForEachNodeCached(cb func(node route.Vertex, + ForEachNodeCached(ctx context.Context, cb func(node route.Vertex, chans map[uint64]*graphdb.DirectedChannel) error) error } diff --git a/graph/db/graph.go b/graph/db/graph.go index 86c14a6cf..3168a6618 100644 --- a/graph/db/graph.go +++ b/graph/db/graph.go @@ -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 diff --git a/graph/db/graph_test.go b/graph/db/graph_test.go index 828d49697..6edce2fb1 100644 --- a/graph/db/graph_test.go +++ b/graph/db/graph_test.go @@ -1358,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 { diff --git a/graph/db/interfaces.go b/graph/db/interfaces.go index fb5b88c99..ecb07ec6d 100644 --- a/graph/db/interfaces.go +++ b/graph/db/interfaces.go @@ -91,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 diff --git a/graph/db/kv_store.go b/graph/db/kv_store.go index e19d80256..9d1fc7cb7 100644 --- a/graph/db/kv_store.go +++ b/graph/db/kv_store.go @@ -667,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 diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go index cf707579f..a77709b8e 100644 --- a/graph/db/sql_store.go +++ b/graph/db/sql_store.go @@ -1102,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, diff --git a/rpcserver.go b/rpcserver.go index 9674d689f..f974ea38b 100644 --- a/rpcserver.go +++ b/rpcserver.go @@ -7130,7 +7130,7 @@ func (r *rpcServer) GetNetworkInfo(ctx context.Context, // network, tallying up the total number of nodes, and also gathering // each node so we can measure the graph diameter and degree stats // below. - err := graph.ForEachNodeCached(func(node route.Vertex, + err := graph.ForEachNodeCached(ctx, func(node route.Vertex, edges map[uint64]*graphdb.DirectedChannel) error { // Increment the total number of nodes with each iteration.