diff --git a/autopilot/agent.go b/autopilot/agent.go index 7b951a66d..95e95c55f 100644 --- a/autopilot/agent.go +++ b/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) diff --git a/autopilot/agent_test.go b/autopilot/agent_test.go index 9b3c30bb5..47305a2eb 100644 --- a/autopilot/agent_test.go +++ b/autopilot/agent_test.go @@ -86,9 +86,9 @@ func (m *mockHeuristic) Name() string { return "mock" } -func (m *mockHeuristic) NodeScores(g ChannelGraph, chans []LocalChannel, - chanSize btcutil.Amount, nodes map[NodeID]struct{}) ( - map[NodeID]*NodeScore, error) { +func (m *mockHeuristic) NodeScores(_ context.Context, g ChannelGraph, + chans []LocalChannel, chanSize btcutil.Amount, + nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) { if m.nodeScoresArgs != nil { directive := directiveArg{ diff --git a/autopilot/combinedattach.go b/autopilot/combinedattach.go index b43856d24..9064a8229 100644 --- a/autopilot/combinedattach.go +++ b/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, - chanSize btcutil.Amount, nodes map[NodeID]struct{}) ( - map[NodeID]*NodeScore, error) { +func (c *WeightedCombAttachment) NodeScores(ctx context.Context, g ChannelGraph, + chans []LocalChannel, chanSize btcutil.Amount, + 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", diff --git a/autopilot/externalscoreattach.go b/autopilot/externalscoreattach.go index 1144d7d4a..a979e15ee 100644 --- a/autopilot/externalscoreattach.go +++ b/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, - chanSize btcutil.Amount, nodes map[NodeID]struct{}) ( - map[NodeID]*NodeScore, error) { +func (s *ExternalScoreAttachment) NodeScores(_ context.Context, g ChannelGraph, + chans []LocalChannel, chanSize btcutil.Amount, + nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) { existingPeers := make(map[NodeID]struct{}) for _, c := range chans { diff --git a/autopilot/externalscoreattach_test.go b/autopilot/externalscoreattach_test.go index aa0f87fb2..bef50a674 100644 --- a/autopilot/externalscoreattach_test.go +++ b/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) diff --git a/autopilot/interface.go b/autopilot/interface.go index ae803632f..db25d9bb2 100644 --- a/autopilot/interface.go +++ b/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) } diff --git a/autopilot/manager.go b/autopilot/manager.go index 036bf3a31..c0ca40b55 100644 --- a/autopilot/manager.go +++ b/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) ( - HeuristicScores, error) { +func (m *Manager) QueryHeuristics(ctx context.Context, nodes []NodeID, + 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) ( - HeuristicScores, error) { +func (m *Manager) queryHeuristics(ctx context.Context, + 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", diff --git a/autopilot/prefattach.go b/autopilot/prefattach.go index 4f55e87ea..76d814a46 100644 --- a/autopilot/prefattach.go +++ b/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, - chanSize btcutil.Amount, nodes map[NodeID]struct{}) ( - map[NodeID]*NodeScore, error) { - - ctx := context.TODO() +func (p *PrefAttachment) NodeScores(ctx context.Context, g ChannelGraph, + chans []LocalChannel, chanSize btcutil.Amount, + nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) { // We first run though the graph once in order to find the median // channel size. diff --git a/autopilot/prefattach_test.go b/autopilot/prefattach_test.go index 7dec5f49f..b79a396b1 100644 --- a/autopilot/prefattach_test.go +++ b/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) diff --git a/autopilot/top_centrality.go b/autopilot/top_centrality.go index 90bf66ae8..e96b34097 100644 --- a/autopilot/top_centrality.go +++ b/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, - chanSize btcutil.Amount, nodes map[NodeID]struct{}) ( - map[NodeID]*NodeScore, error) { - - ctx := context.TODO() +func (g *TopCentrality) NodeScores(ctx context.Context, graph ChannelGraph, + chans []LocalChannel, chanSize btcutil.Amount, + nodes map[NodeID]struct{}) (map[NodeID]*NodeScore, error) { // Calculate betweenness centrality for the whole graph. if err := g.centralityMetric.Refresh(ctx, graph); err != nil { diff --git a/autopilot/top_centrality_test.go b/autopilot/top_centrality_test.go index 9e8242761..e8ca3616d 100644 --- a/autopilot/top_centrality_test.go +++ b/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) diff --git a/lnrpc/autopilotrpc/autopilot_server.go b/lnrpc/autopilotrpc/autopilot_server.go index 3e3c6f8f7..67e6c283a 100644 --- a/lnrpc/autopilotrpc/autopilot_server.go +++ b/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