graphdb: thread topology update context

graph/db/graph.go

@@ -73,15 +73,17 @@ func (c *ChannelGraph) Start() error {
 	log.Debugf("ChannelGraph starting")
 	defer log.Debug("ChannelGraph started")
 
+	ctx := context.TODO()
+
 	if c.graphCache != nil {
-		if err := c.populateCache(context.TODO()); err != nil {
+		if err := c.populateCache(ctx); err != nil {
 			return fmt.Errorf("could not populate the graph "+
 				"cache: %w", err)
 		}
 	}
 
 	c.wg.Add(1)
-	go c.handleTopologySubscriptions()
+	go c.handleTopologySubscriptions(ctx)
 
 	return nil
 }
@@ -106,7 +108,7 @@ func (c *ChannelGraph) Stop() error {
 // synchronously.
 //
 // NOTE: this MUST be run in a goroutine.
-func (c *ChannelGraph) handleTopologySubscriptions() {
+func (c *ChannelGraph) handleTopologySubscriptions(ctx context.Context) {
 	defer c.wg.Done()
 
 	for {
@@ -118,7 +120,7 @@ func (c *ChannelGraph) handleTopologySubscriptions() {
 			// synchronously so that we can guarantee the order of
 			// notification delivery.
 			c.wg.Add(1)
-			go c.handleTopologyUpdate(update)
+			go c.handleTopologyUpdate(ctx, update)
 
 			// TODO(roasbeef): remove all unconnected vertexes
 			// after N blocks pass with no corresponding
@@ -149,6 +151,9 @@ func (c *ChannelGraph) handleTopologySubscriptions() {
 				exit:     make(chan struct{}),
 			})
 
+		case <-ctx.Done():
+			return
+
 		case <-c.quit:
 			return
 		}

graph/db/notifications.go

@@ -202,11 +202,11 @@ func (c *ChannelGraph) notifyTopologyChange(topologyDiff *TopologyChange) {
 //
 // NOTE: must be run inside goroutine and must only ever be called from within
 // handleTopologySubscriptions.
-func (c *ChannelGraph) handleTopologyUpdate(update any) {
+func (c *ChannelGraph) handleTopologyUpdate(ctx context.Context, update any) {
 	defer c.wg.Done()
 
 	topChange := &TopologyChange{}
-	err := c.addToTopologyChange(topChange, update)
+	err := c.addToTopologyChange(ctx, topChange, update)
 	if err != nil {
 		log.Errorf("unable to update topology change notification: %v",
 			err)
@@ -376,8 +376,8 @@ type ChannelEdgeUpdate struct {
 // constitutes. This function will also fetch any required auxiliary
 // information required to create the topology change update from the graph
 // database.
-func (c *ChannelGraph) addToTopologyChange(update *TopologyChange,
-	msg any) error {
+func (c *ChannelGraph) addToTopologyChange(ctx context.Context,
+	update *TopologyChange, msg any) error {
 
 	switch m := msg.(type) {
 
@@ -414,7 +414,7 @@ func (c *ChannelGraph) addToTopologyChange(update *TopologyChange,
 		// in order to get the information concerning which nodes are
 		// being connected.
 		edgeInfo, _, _, err := c.FetchChannelEdgesByID(
-			context.TODO(), m.ChannelID,
+			ctx, m.ChannelID,
 		)
 		if err != nil {
 			return fmt.Errorf("unable fetch channel edge: %w", err)