From 46f09ba1ff701940820d78dddf85c20f0cd6ce89 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Wed, 9 Apr 2025 06:31:05 +0200
Subject: [PATCH] autopilot: continue threading context

Remove one context.TODO and add one more.
---
 autopilot/betweenness_centrality.go      |  4 ++--
 autopilot/betweenness_centrality_test.go | 11 +++++++++--
 autopilot/interface.go                   |  2 +-
 autopilot/top_centrality.go              |  5 ++++-
 rpcserver.go                             |  2 +-
 5 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/autopilot/betweenness_centrality.go b/autopilot/betweenness_centrality.go
index 2a45fe1f1..c3f0a3d02 100644
--- a/autopilot/betweenness_centrality.go
+++ b/autopilot/betweenness_centrality.go
@@ -169,8 +169,8 @@ func betweennessCentrality(g *SimpleGraph, s int, centrality []float64) {
 }
 
 // Refresh recalculates and stores centrality values.
-func (bc *BetweennessCentrality) Refresh(graph ChannelGraph) error {
-	ctx := context.TODO()
+func (bc *BetweennessCentrality) Refresh(ctx context.Context,
+	graph ChannelGraph) error {
 
 	cache, err := NewSimpleGraph(ctx, graph)
 	if err != nil {
diff --git a/autopilot/betweenness_centrality_test.go b/autopilot/betweenness_centrality_test.go
index 8e3e07ce2..b1de27834 100644
--- a/autopilot/betweenness_centrality_test.go
+++ b/autopilot/betweenness_centrality_test.go
@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -30,6 +31,9 @@ func TestBetweennessCentralityMetricConstruction(t *testing.T) {
 
 // Tests that empty graph results in empty centrality result.
 func TestBetweennessCentralityEmptyGraph(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
 	centralityMetric, err := NewBetweennessCentralityMetric(1)
 	require.NoError(
 		t, err,
@@ -42,7 +46,7 @@ func TestBetweennessCentralityEmptyGraph(t *testing.T) {
 		require.NoError(t, err, "unable to create graph")
 
 		success := t.Run(chanGraph.name, func(t1 *testing.T) {
-			err = centralityMetric.Refresh(graph)
+			err = centralityMetric.Refresh(ctx, graph)
 			require.NoError(t1, err)
 
 			centrality := centralityMetric.GetMetric(false)
@@ -59,6 +63,9 @@ func TestBetweennessCentralityEmptyGraph(t *testing.T) {
 
 // Test betweenness centrality calculating using an example graph.
 func TestBetweennessCentralityWithNonEmptyGraph(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
 	workers := []int{1, 3, 9, 100}
 
 	tests := []struct {
@@ -100,7 +107,7 @@ func TestBetweennessCentralityWithNonEmptyGraph(t *testing.T) {
 					t1, graph, centralityTestGraph,
 				)
 
-				err = metric.Refresh(graph)
+				err = metric.Refresh(ctx, graph)
 				require.NoError(t1, err)
 
 				for _, expected := range tests {
diff --git a/autopilot/interface.go b/autopilot/interface.go
index 0991d9864..ae803632f 100644
--- a/autopilot/interface.go
+++ b/autopilot/interface.go
@@ -157,7 +157,7 @@ type NodeMetric interface {
 	Name() string
 
 	// Refresh refreshes the metric values based on the current graph.
-	Refresh(graph ChannelGraph) error
+	Refresh(ctx context.Context, graph ChannelGraph) error
 
 	// GetMetric returns the latest value of this metric. Values in the
 	// map are per node and can be in arbitrary domain. If normalize is
diff --git a/autopilot/top_centrality.go b/autopilot/top_centrality.go
index 65157c621..90bf66ae8 100644
--- a/autopilot/top_centrality.go
+++ b/autopilot/top_centrality.go
@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"runtime"
 
 	"github.com/btcsuite/btcd/btcutil"
@@ -54,8 +55,10 @@ func (g *TopCentrality) NodeScores(graph ChannelGraph, chans []LocalChannel,
 	chanSize btcutil.Amount, nodes map[NodeID]struct{}) (
 	map[NodeID]*NodeScore, error) {
 
+	ctx := context.TODO()
+
 	// Calculate betweenness centrality for the whole graph.
-	if err := g.centralityMetric.Refresh(graph); err != nil {
+	if err := g.centralityMetric.Refresh(ctx, graph); err != nil {
 		return nil, err
 	}
 
diff --git a/rpcserver.go b/rpcserver.go
index d51cf1cab..72cf953bf 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -6857,7 +6857,7 @@ func (r *rpcServer) GetNodeMetrics(ctx context.Context,
 	if err != nil {
 		return nil, err
 	}
-	if err := centralityMetric.Refresh(channelGraph); err != nil {
+	if err := centralityMetric.Refresh(ctx, channelGraph); err != nil {
 		return nil, err
 	}