[graph-work-side-branch]: temp side branch for graph work

autopilot/agent.go

@@ -2,6 +2,7 @@ package autopilot
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"math/rand"
 	"net"
@@ -11,6 +12,7 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/davecgh/go-spew/spew"
+	"github.com/lightningnetwork/lnd/fn/v2"
 	"github.com/lightningnetwork/lnd/lnwire"
 )
 
@@ -166,8 +168,9 @@ type Agent struct {
 	pendingOpens map[NodeID]LocalChannel
 	pendingMtx   sync.Mutex
 
-	quit chan struct{}
-	wg   sync.WaitGroup
+	quit   chan struct{}
+	wg     sync.WaitGroup
+	cancel fn.Option[context.CancelFunc]
 }
 
 // New creates a new instance of the Agent instantiated using the passed
@@ -202,17 +205,20 @@ func New(cfg Config, initialState []LocalChannel) (*Agent, error) {
 func (a *Agent) Start() error {
 	var err error
 	a.started.Do(func() {
-		err = a.start()
+		ctx, cancel := context.WithCancel(context.Background())
+		a.cancel = fn.Some(cancel)
+
+		err = a.start(ctx)
 	})
 	return err
 }
 
-func (a *Agent) start() error {
+func (a *Agent) start(ctx context.Context) error {
 	rand.Seed(time.Now().Unix())
 	log.Infof("Autopilot Agent starting")
 
 	a.wg.Add(1)
-	go a.controller()
+	go a.controller(ctx)
 
 	return nil
 }
@@ -230,6 +236,7 @@ func (a *Agent) Stop() error {
 func (a *Agent) stop() error {
 	log.Infof("Autopilot Agent stopping")
 
+	a.cancel.WhenSome(func(fn context.CancelFunc) { fn() })
 	close(a.quit)
 	a.wg.Wait()
 
@@ -401,7 +408,7 @@ func mergeChanState(pendingChans map[NodeID]LocalChannel,
 // and external state changes as a result of decisions it makes w.r.t channel
 // allocation, or attributes affecting its control loop being updated by the
 // backing Lightning Node.
-func (a *Agent) controller() {
+func (a *Agent) controller(ctx context.Context) {
 	defer a.wg.Done()
 
 	// We'll start off by assigning our starting balance, and injecting
@@ -502,6 +509,9 @@ func (a *Agent) controller() {
 		// immediately.
 		case <-a.quit:
 			return
+
+		case <-ctx.Done():
+			return
 		}
 
 		a.pendingMtx.Lock()
@@ -539,7 +549,7 @@ func (a *Agent) controller() {
 		log.Infof("Triggering attachment directive dispatch, "+
 			"total_funds=%v", a.totalBalance)
 
-		err := a.openChans(availableFunds, numChans, totalChans)
+		err := a.openChans(ctx, availableFunds, numChans, totalChans)
 		if err != nil {
 			log.Errorf("Unable to open channels: %v", err)
 		}
@@ -548,8 +558,8 @@ func (a *Agent) controller() {
 
 // openChans queries the agent's heuristic for a set of channel candidates, and
 // attempts to open channels to them.
-func (a *Agent) openChans(availableFunds btcutil.Amount, numChans uint32,
-	totalChans []LocalChannel) error {
+func (a *Agent) openChans(ctx context.Context, availableFunds btcutil.Amount,
+	numChans uint32, totalChans []LocalChannel) error {
 
 	// As channel size we'll use the maximum channel size available.
 	chanSize := a.cfg.Constraints.MaxChanSize()
@@ -598,7 +608,9 @@ func (a *Agent) openChans(availableFunds btcutil.Amount, numChans uint32,
 	selfPubBytes := a.cfg.Self.SerializeCompressed()
 	nodes := make(map[NodeID]struct{})
 	addresses := make(map[NodeID][]net.Addr)
-	if err := a.cfg.Graph.ForEachNode(func(node Node) error {
+	if err := a.cfg.Graph.ForEachNode(ctx, func(_ context.Context,
+		node Node) error {
+
 		nID := NodeID(node.PubKey())
 
 		// If we come across ourselves, them we'll continue in
@@ -636,7 +648,7 @@ func (a *Agent) openChans(availableFunds btcutil.Amount, numChans uint32,
 	// 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

@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
@@ -85,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{

autopilot/betweenness_centrality.go

@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"fmt"
 	"sync"
 )
@@ -168,8 +169,10 @@ func betweennessCentrality(g *SimpleGraph, s int, centrality []float64) {
 }
 
 // Refresh recalculates and stores centrality values.
-func (bc *BetweennessCentrality) Refresh(graph ChannelGraph) error {
-	cache, err := NewSimpleGraph(graph)
+func (bc *BetweennessCentrality) Refresh(ctx context.Context,
+	graph ChannelGraph) error {
+
+	cache, err := NewSimpleGraph(ctx, graph)
 	if err != nil {
 		return err
 	}

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 {

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",

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 {

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/graph.go

@@ -1,6 +1,7 @@
 package autopilot
 
 import (
+	"context"
 	"encoding/hex"
 	"net"
 	"sort"
@@ -80,7 +81,9 @@ func (d *dbNode) Addrs() []net.Addr {
 // describes the active channel.
 //
 // NOTE: Part of the autopilot.Node interface.
-func (d *dbNode) ForEachChannel(cb func(ChannelEdge) error) error {
+func (d *dbNode) ForEachChannel(ctx context.Context,
+	cb func(context.Context, ChannelEdge) error) error {
+
 	return d.tx.ForEachChannel(func(ei *models.ChannelEdgeInfo, ep,
 		_ *models.ChannelEdgePolicy) error {
 
@@ -108,7 +111,7 @@ func (d *dbNode) ForEachChannel(cb func(ChannelEdge) error) error {
 			},
 		}
 
-		return cb(edge)
+		return cb(ctx, edge)
 	})
 }
 
@@ -117,7 +120,9 @@ func (d *dbNode) ForEachChannel(cb func(ChannelEdge) error) error {
 // error, then execution should be terminated.
 //
 // NOTE: Part of the autopilot.ChannelGraph interface.
-func (d *databaseChannelGraph) ForEachNode(cb func(Node) error) error {
+func (d *databaseChannelGraph) ForEachNode(ctx context.Context,
+	cb func(context.Context, Node) error) error {
+
 	return d.db.ForEachNode(func(nodeTx graphdb.NodeRTx) error {
 		// We'll skip over any node that doesn't have any advertised
 		// addresses. As we won't be able to reach them to actually
@@ -129,7 +134,8 @@ func (d *databaseChannelGraph) ForEachNode(cb func(Node) error) error {
 		node := &dbNode{
 			tx: nodeTx,
 		}
-		return cb(node)
+
+		return cb(ctx, node)
 	})
 }
 
@@ -185,7 +191,9 @@ func (nc dbNodeCached) Addrs() []net.Addr {
 // describes the active channel.
 //
 // NOTE: Part of the autopilot.Node interface.
-func (nc dbNodeCached) ForEachChannel(cb func(ChannelEdge) error) error {
+func (nc dbNodeCached) ForEachChannel(ctx context.Context,
+	cb func(context.Context, ChannelEdge) error) error {
+
 	for cid, channel := range nc.channels {
 		edge := ChannelEdge{
 			ChanID:   lnwire.NewShortChanIDFromInt(cid),
@@ -195,7 +203,7 @@ func (nc dbNodeCached) ForEachChannel(cb func(ChannelEdge) error) error {
 			},
 		}
 
-		if err := cb(edge); err != nil {
+		if err := cb(ctx, edge); err != nil {
 			return err
 		}
 	}
@@ -208,7 +216,9 @@ func (nc dbNodeCached) ForEachChannel(cb func(ChannelEdge) error) error {
 // error, then execution should be terminated.
 //
 // NOTE: Part of the autopilot.ChannelGraph interface.
-func (dc *databaseChannelGraphCached) ForEachNode(cb func(Node) error) error {
+func (dc *databaseChannelGraphCached) ForEachNode(ctx context.Context,
+	cb func(context.Context, Node) error) error {
+
 	return dc.db.ForEachNodeCached(func(n route.Vertex,
 		channels map[uint64]*graphdb.DirectedChannel) error {
 
@@ -217,7 +227,8 @@ func (dc *databaseChannelGraphCached) ForEachNode(cb func(Node) error) error {
 				node:     n,
 				channels: channels,
 			}
-			return cb(node)
+
+			return cb(ctx, node)
 		}
 		return nil
 	})
@@ -262,9 +273,11 @@ func (m memNode) Addrs() []net.Addr {
 // describes the active channel.
 //
 // NOTE: Part of the autopilot.Node interface.
-func (m memNode) ForEachChannel(cb func(ChannelEdge) error) error {
+func (m memNode) ForEachChannel(ctx context.Context,
+	cb func(context.Context, ChannelEdge) error) error {
+
 	for _, channel := range m.chans {
-		if err := cb(channel); err != nil {
+		if err := cb(ctx, channel); err != nil {
 			return err
 		}
 	}