autopilot: update AttachmentHeuristics with context

Continue threading context through the autopilot system and remove the
remaining context.TODOs.

autopilot/agent.go

@@ -648,7 +648,7 @@ func (a *Agent) openChans(ctx context.Context, availableFunds btcutil.Amount,
 	// graph.
 	log.Debugf("Scoring %d nodes for chan_size=%v", len(nodes), chanSize)
 	scores, err := a.cfg.Heuristic.NodeScores(
-		a.cfg.Graph, totalChans, chanSize, nodes,
+		ctx, a.cfg.Graph, totalChans, chanSize, nodes,
 	)
 	if err != nil {
 		return fmt.Errorf("unable to calculate node scores : %w", err)

autopilot/agent_test.go

@@ -86,9 +86,9 @@ func (m *mockHeuristic) Name() string {
 	return "mock"
 }
 
-func (m *mockHeuristic) NodeScores(g ChannelGraph, chans []LocalChannel,
+func (m *mockHeuristic) NodeScores(_ context.Context, g ChannelGraph,
-	chanSize btcutil.Amount, nodes map[NodeID]struct{}) (
+	chans []LocalChannel, chanSize btcutil.Amount,
-	map[NodeID]*NodeScore, error) {
+	nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) {
 
 	if m.nodeScoresArgs != nil {
 		directive := directiveArg{

autopilot/combinedattach.go

@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/btcsuite/btcd/btcutil"
@@ -70,9 +71,9 @@ func (c *WeightedCombAttachment) Name() string {
 // is the maximum possible improvement in connectivity.
 //
 // NOTE: This is a part of the AttachmentHeuristic interface.
-func (c *WeightedCombAttachment) NodeScores(g ChannelGraph, chans []LocalChannel,
+func (c *WeightedCombAttachment) NodeScores(ctx context.Context, g ChannelGraph,
-	chanSize btcutil.Amount, nodes map[NodeID]struct{}) (
+	chans []LocalChannel, chanSize btcutil.Amount,
-	map[NodeID]*NodeScore, error) {
+	nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) {
 
 	// We now query each heuristic to determine the score they give to the
 	// nodes for the given channel size.
@@ -81,7 +82,7 @@ func (c *WeightedCombAttachment) NodeScores(g ChannelGraph, chans []LocalChannel
 		log.Tracef("Getting scores from sub heuristic %v", h.Name())
 
 		s, err := h.NodeScores(
-			g, chans, chanSize, nodes,
+			ctx, g, chans, chanSize, nodes,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("unable to get sub score: %w",

autopilot/externalscoreattach.go

@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
@@ -80,9 +81,9 @@ func (s *ExternalScoreAttachment) SetNodeScores(targetHeuristic string,
 // not known will get a score of 0.
 //
 // NOTE: This is a part of the AttachmentHeuristic interface.
-func (s *ExternalScoreAttachment) NodeScores(g ChannelGraph, chans []LocalChannel,
+func (s *ExternalScoreAttachment) NodeScores(_ context.Context, g ChannelGraph,
-	chanSize btcutil.Amount, nodes map[NodeID]struct{}) (
+	chans []LocalChannel, chanSize btcutil.Amount,
-	map[NodeID]*NodeScore, error) {
+	nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) {
 
 	existingPeers := make(map[NodeID]struct{})
 	for _, c := range chans {

autopilot/externalscoreattach_test.go

@@ -1,6 +1,7 @@
 package autopilot_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -22,6 +23,7 @@ func randKey() (*btcec.PublicKey, error) {
 // ExternalScoreAttachment correctly reflects the scores we set last.
 func TestSetNodeScores(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	const name = "externalscore"
 
@@ -62,7 +64,7 @@ func TestSetNodeScores(t *testing.T) {
 		q[nID] = struct{}{}
 	}
 	resp, err := h.NodeScores(
-		nil, nil, btcutil.Amount(btcutil.SatoshiPerBitcoin), q,
+		ctx, nil, nil, btcutil.Amount(btcutil.SatoshiPerBitcoin), q,
 	)
 	if err != nil {
 		t.Fatal(err)

autopilot/interface.go

@@ -144,7 +144,7 @@ type AttachmentHeuristic interface {
 	//
 	// NOTE: A NodeID not found in the returned map is implicitly given a
 	// score of 0.
-	NodeScores(g ChannelGraph, chans []LocalChannel,
+	NodeScores(ctx context.Context, g ChannelGraph, chans []LocalChannel,
 		chanSize btcutil.Amount, nodes map[NodeID]struct{}) (
 		map[NodeID]*NodeScore, error)
 }

autopilot/manager.go

@@ -276,8 +276,8 @@ func (m *Manager) StopAgent() error {
 }
 
 // QueryHeuristics queries the available autopilot heuristics for node scores.
-func (m *Manager) QueryHeuristics(nodes []NodeID, localState bool) (
+func (m *Manager) QueryHeuristics(ctx context.Context, nodes []NodeID,
-	HeuristicScores, error) {
+	localState bool) (HeuristicScores, error) {
 
 	m.Lock()
 	defer m.Unlock()
@@ -288,7 +288,8 @@ func (m *Manager) QueryHeuristics(nodes []NodeID, localState bool) (
 	}
 
 	log.Debugf("Querying heuristics for %d nodes", len(n))
-	return m.queryHeuristics(n, localState)
+
+	return m.queryHeuristics(ctx, n, localState)
 }
 
 // HeuristicScores is an alias for a map that maps heuristic names to a map of
@@ -299,8 +300,8 @@ type HeuristicScores map[string]map[NodeID]float64
 // the agent's current active heuristic.
 //
 // NOTE: Must be called with the manager's lock.
-func (m *Manager) queryHeuristics(nodes map[NodeID]struct{}, localState bool) (
+func (m *Manager) queryHeuristics(ctx context.Context,
-	HeuristicScores, error) {
+	nodes map[NodeID]struct{}, localState bool) (HeuristicScores, error) {
 
 	// If we want to take the local state into action when querying the
 	// heuristics, we fetch it. If not we'll just pass an empty slice to
@@ -348,7 +349,7 @@ func (m *Manager) queryHeuristics(nodes map[NodeID]struct{}, localState bool) (
 		}
 
 		s, err := h.NodeScores(
-			m.cfg.PilotCfg.Graph, totalChans, chanSize, nodes,
+			ctx, m.cfg.PilotCfg.Graph, totalChans, chanSize, nodes,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("unable to get sub score: %w",

autopilot/prefattach.go

@@ -79,11 +79,9 @@ func (p *PrefAttachment) Name() string {
 // given to nodes already having high connectivity in the graph.
 //
 // NOTE: This is a part of the AttachmentHeuristic interface.
-func (p *PrefAttachment) NodeScores(g ChannelGraph, chans []LocalChannel,
+func (p *PrefAttachment) NodeScores(ctx context.Context, g ChannelGraph,
-	chanSize btcutil.Amount, nodes map[NodeID]struct{}) (
+	chans []LocalChannel, chanSize btcutil.Amount,
-	map[NodeID]*NodeScore, error) {
+	nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) {
-
-	ctx := context.TODO()
 
 	// We first run though the graph once in order to find the median
 	// channel size.

autopilot/prefattach_test.go

@@ -88,6 +88,8 @@ var chanGraphs = []struct {
 // TestPrefAttachmentSelectEmptyGraph ensures that when passed an
 // empty graph, the NodeSores function always returns a score of 0.
 func TestPrefAttachmentSelectEmptyGraph(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
 	prefAttach := NewPrefAttachment()
 
 	// Create a random public key, which we will query to get a score for.
@@ -108,7 +110,7 @@ func TestPrefAttachmentSelectEmptyGraph(t *testing.T) {
 			// attempt to get the score for this one node.
 			const walletFunds = btcutil.SatoshiPerBitcoin
 			scores, err := prefAttach.NodeScores(
-				graph, nil, walletFunds, nodes,
+				ctx, graph, nil, walletFunds, nodes,
 			)
 			require.NoError(t1, err)
 
@@ -172,7 +174,7 @@ func TestPrefAttachmentSelectTwoVertexes(t *testing.T) {
 			// attempt to get our candidates channel score given
 			// the current state of the graph.
 			candidates, err := prefAttach.NodeScores(
-				graph, nil, maxChanSize, nodes,
+				ctx, graph, nil, maxChanSize, nodes,
 			)
 			require.NoError(t1, err)
 
@@ -280,7 +282,7 @@ func TestPrefAttachmentSelectGreedyAllocation(t *testing.T) {
 			// result, the heuristic should try to greedily
 			// allocate funds to channels.
 			scores, err := prefAttach.NodeScores(
-				graph, nil, maxChanSize, nodes,
+				ctx, graph, nil, maxChanSize, nodes,
 			)
 			require.NoError(t1, err)
 
@@ -298,7 +300,7 @@ func TestPrefAttachmentSelectGreedyAllocation(t *testing.T) {
 			// candidates of that size.
 			const remBalance = btcutil.SatoshiPerBitcoin * 0.5
 			scores, err = prefAttach.NodeScores(
-				graph, nil, remBalance, nodes,
+				ctx, graph, nil, remBalance, nodes,
 			)
 			require.NoError(t1, err)
 
@@ -358,7 +360,7 @@ func TestPrefAttachmentSelectSkipNodes(t *testing.T) {
 			// With our graph created, we'll now get the scores for
 			// all nodes in the graph.
 			scores, err := prefAttach.NodeScores(
-				graph, nil, maxChanSize, nodes,
+				ctx, graph, nil, maxChanSize, nodes,
 			)
 			require.NoError(t1, err)
 
@@ -386,7 +388,7 @@ func TestPrefAttachmentSelectSkipNodes(t *testing.T) {
 			// then all nodes should have a score of zero, since we
 			// already got channels to them.
 			scores, err = prefAttach.NodeScores(
-				graph, chans, maxChanSize, nodes,
+				ctx, graph, chans, maxChanSize, nodes,
 			)
 			require.NoError(t1, err)
 

autopilot/top_centrality.go

@@ -51,11 +51,9 @@ func (g *TopCentrality) Name() string {
 // As our current implementation of betweenness centrality is non-incremental,
 // NodeScores will recalculate the centrality values on every call, which is
 // slow for large graphs.
-func (g *TopCentrality) NodeScores(graph ChannelGraph, chans []LocalChannel,
+func (g *TopCentrality) NodeScores(ctx context.Context, graph ChannelGraph,
-	chanSize btcutil.Amount, nodes map[NodeID]struct{}) (
+	chans []LocalChannel, chanSize btcutil.Amount,
-	map[NodeID]*NodeScore, error) {
+	nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) {
-
-	ctx := context.TODO()
 
 	// Calculate betweenness centrality for the whole graph.
 	if err := g.centralityMetric.Refresh(ctx, graph); err != nil {

autopilot/top_centrality_test.go

@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"testing"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -58,7 +59,7 @@ func testTopCentrality(t *testing.T, graph testGraph,
 		// Attempt to get centrality scores and expect
 		// that the result equals with the expected set.
 		scores, err := topCentrality.NodeScores(
-			graph, channels, chanSize, nodes,
+			context.Background(), graph, channels, chanSize, nodes,
 		)
 
 		require.NoError(t, err)

lnrpc/autopilotrpc/autopilot_server.go

@@ -236,7 +236,7 @@ func (s *Server) QueryScores(ctx context.Context, in *QueryScoresRequest) (
 
 	// Query the heuristics.
 	heuristicScores, err := s.manager.QueryHeuristics(
-		nodes, !in.IgnoreLocalState,
+		ctx, nodes, !in.IgnoreLocalState,
 	)
 	if err != nil {
 		return nil, err