graphdb: pass context to PruneGraph

graph/builder.go

@@ -185,7 +185,8 @@ func (b *Builder) Start() error {
 			// the prune height to the current best height of the
 			// chain backend.
 			_, err = b.cfg.Graph.PruneGraph(
-				nil, bestHash, uint32(bestHeight),
+				context.TODO(), nil, bestHash,
+				uint32(bestHeight),
 			)
 			if err != nil {
 				return err
@@ -443,7 +444,7 @@ func (b *Builder) syncGraphWithChain() error {
 	// With the spent outputs gathered, attempt to prune the channel graph,
 	// also passing in the best hash+height so the prune tip can be updated.
 	closedChans, err := b.cfg.Graph.PruneGraph(
-		spentOutputs, bestHash, uint32(bestHeight),
+		context.TODO(), spentOutputs, bestHash, uint32(bestHeight),
 	)
 	if err != nil {
 		return err
@@ -857,8 +858,10 @@ func (b *Builder) updateGraphWithClosedChannels(
 	// With the spent outputs gathered, attempt to prune the channel graph,
 	// also passing in the hash+height of the block being pruned so the
 	// prune tip can be updated.
-	chansClosed, err := b.cfg.Graph.PruneGraph(spentOutputs,
-		&chainUpdate.Hash, chainUpdate.Height)
+	chansClosed, err := b.cfg.Graph.PruneGraph(
+		context.TODO(), spentOutputs, &chainUpdate.Hash,
+		chainUpdate.Height,
+	)
 	if err != nil {
 		log.Errorf("unable to prune routing table: %v", err)
 		return err

graph/db/graph.go

@@ -446,12 +446,13 @@ func (c *ChannelGraph) DisconnectBlockAtHeight(height uint32) (
 // prune the graph is stored so callers can ensure the graph is fully in sync
 // with the current UTXO state. A slice of channels that have been closed by
 // the target block are returned if the function succeeds without error.
-func (c *ChannelGraph) PruneGraph(spentOutputs []*wire.OutPoint,
+func (c *ChannelGraph) PruneGraph(ctx context.Context,
+	spentOutputs []*wire.OutPoint,
 	blockHash *chainhash.Hash, blockHeight uint32) (
 	[]*models.ChannelEdgeInfo, error) {
 
 	edges, nodes, err := c.db.PruneGraph(
-		spentOutputs, blockHash, blockHeight,
+		ctx, spentOutputs, blockHash, blockHeight,
 	)
 	if err != nil {
 		return nil, err

graph/db/graph_test.go

@@ -863,12 +863,12 @@ func TestDisconnectBlockAtHeight(t *testing.T) {
 
 	// Prune the graph a few times to make sure we have entries in the
 	// prune log.
-	_, err := graph.PruneGraph(spendOutputs, &blockHash, 155)
+	_, err := graph.PruneGraph(ctx, spendOutputs, &blockHash, 155)
 	require.NoError(t, err, "unable to prune graph")
 	var blockHash2 chainhash.Hash
 	copy(blockHash2[:], bytes.Repeat([]byte{2}, 32))
 
-	_, err = graph.PruneGraph(spendOutputs, &blockHash2, 156)
+	_, err = graph.PruneGraph(ctx, spendOutputs, &blockHash2, 156)
 	require.NoError(t, err, "unable to prune graph")
 
 	// We'll create 3 almost identical edges, so first create a helper
@@ -2287,7 +2287,9 @@ func TestGraphPruning(t *testing.T) {
 	copy(blockHash[:], bytes.Repeat([]byte{1}, 32))
 	blockHeight := uint32(1)
 	block := channelPoints[:2]
-	prunedChans, err := graph.PruneGraph(block, &blockHash, blockHeight)
+	prunedChans, err := graph.PruneGraph(
+		ctx, block, &blockHash, blockHeight,
+	)
 	require.NoError(t, err, "unable to prune graph")
 	require.Len(t, prunedChans, 2)
 
@@ -2313,7 +2315,7 @@ func TestGraphPruning(t *testing.T) {
 	blockHash = sha256.Sum256(blockHash[:])
 	blockHeight = 2
 	prunedChans, err = graph.PruneGraph(
-		[]*wire.OutPoint{nonChannel}, &blockHash, blockHeight,
+		ctx, []*wire.OutPoint{nonChannel}, &blockHash, blockHeight,
 	)
 	require.NoError(t, err, "unable to prune graph")
 
@@ -2331,7 +2333,7 @@ func TestGraphPruning(t *testing.T) {
 	blockHash = sha256.Sum256(blockHash[:])
 	blockHeight = 3
 	prunedChans, err = graph.PruneGraph(
-		channelPoints[2:], &blockHash, blockHeight,
+		ctx, channelPoints[2:], &blockHash, blockHeight,
 	)
 	require.NoError(t, err, "unable to prune graph")
 
@@ -3383,7 +3385,7 @@ func TestStressTestChannelGraphAPI(t *testing.T) {
 				}
 
 				_, err := graph.PruneGraph(
-					spentOutpoints, &blockHash, 100,
+					ctx, spentOutpoints, &blockHash, 100,
 				)
 
 				return err
@@ -3995,7 +3997,8 @@ func TestChannelEdgePruningUpdateIndexDeletion(t *testing.T) {
 	var blockHash chainhash.Hash
 	copy(blockHash[:], bytes.Repeat([]byte{2}, 32))
 	_, err := graph.PruneGraph(
-		[]*wire.OutPoint{&edgeInfo.ChannelPoint}, &blockHash, 101,
+		ctx, []*wire.OutPoint{&edgeInfo.ChannelPoint}, &blockHash,
+		101,
 	)
 	require.NoError(t, err, "unable to prune graph")
 
@@ -4762,12 +4765,12 @@ func testBatchedAddChannelEdge(t *testing.T, v lnwire.GossipVersion) {
 
 	// Prune the graph a few times to make sure we have entries in the
 	// prune log.
-	_, err := graph.PruneGraph(spendOutputs, &blockHash, 155)
+	_, err := graph.PruneGraph(ctx, spendOutputs, &blockHash, 155)
 	require.Nil(t, err)
 	var blockHash2 chainhash.Hash
 	copy(blockHash2[:], bytes.Repeat([]byte{2}, 32))
 
-	_, err = graph.PruneGraph(spendOutputs, &blockHash2, 156)
+	_, err = graph.PruneGraph(ctx, spendOutputs, &blockHash2, 156)
 	require.Nil(t, err)
 
 	// We'll create 3 almost identical edges, so first create a helper

graph/db/interfaces.go

@@ -401,7 +401,7 @@ type Store interface { //nolint:interfacebloat
 	// slice of channels that have been closed by the target block along
 	// with any pruned nodes are returned if the function succeeds without
 	// error.
-	PruneGraph(spentOutputs []*wire.OutPoint,
+	PruneGraph(ctx context.Context, spentOutputs []*wire.OutPoint,
 		blockHash *chainhash.Hash, blockHeight uint32) (
 		[]*models.ChannelEdgeInfo, []route.Vertex, error)
 

graph/db/kv_store.go

@@ -1492,7 +1492,7 @@ const (
 // with the current UTXO state. A slice of channels that have been closed by
 // the target block along with any pruned nodes are returned if the function
 // succeeds without error.
-func (c *KVStore) PruneGraph(spentOutputs []*wire.OutPoint,
+func (c *KVStore) PruneGraph(_ context.Context, spentOutputs []*wire.OutPoint,
 	blockHash *chainhash.Hash, blockHeight uint32) (
 	[]*models.ChannelEdgeInfo, []route.Vertex, error) {
 

graph/db/sql_store.go

@@ -2831,11 +2831,10 @@ func (s *SQLStore) PruneGraphNodes(ctx context.Context) (
 // succeeds without error.
 //
 // NOTE: part of the Store interface.
-func (s *SQLStore) PruneGraph(spentOutputs []*wire.OutPoint,
-	blockHash *chainhash.Hash, blockHeight uint32) (
-	[]*models.ChannelEdgeInfo, []route.Vertex, error) {
-
-	ctx := context.TODO()
+func (s *SQLStore) PruneGraph(ctx context.Context,
+	spentOutputs []*wire.OutPoint, blockHash *chainhash.Hash,
+	blockHeight uint32) ([]*models.ChannelEdgeInfo, []route.Vertex,
+	error) {
 
 	s.cacheMu.Lock()
 	defer s.cacheMu.Unlock()