diff --git a/graph/db/graph.go b/graph/db/graph.go index a1eb2e32d..d60ca24cf 100644 --- a/graph/db/graph.go +++ b/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 } diff --git a/graph/db/notifications.go b/graph/db/notifications.go index 6676cbebe..1c8a88986 100644 --- a/graph/db/notifications.go +++ b/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)