graph/db: thread context through to UpdateEdgePolicy

graph/builder.go

@@ -915,6 +915,8 @@ func (b *Builder) MarkZombieEdge(chanID uint64) error {
 // ApplyChannelUpdate validates a channel update and if valid, applies it to the
 // database. It returns a bool indicating whether the updates were successful.
 func (b *Builder) ApplyChannelUpdate(msg *lnwire.ChannelUpdate1) bool {
+	ctx := context.TODO()
+
 	ch, _, _, err := b.GetChannelByID(msg.ShortChannelID)
 	if err != nil {
 		log.Errorf("Unable to retrieve channel by id: %v", err)
@@ -959,7 +961,7 @@ func (b *Builder) ApplyChannelUpdate(msg *lnwire.ChannelUpdate1) bool {
 		ExtraOpaqueData:           msg.ExtraOpaqueData,
 	}
 
-	err = b.UpdateEdge(update)
+	err = b.UpdateEdge(ctx, update)
 	if err != nil && !IsError(err, ErrIgnored, ErrOutdated) {
 		log.Errorf("Unable to apply channel update: %v", err)
 		return false
@@ -1118,10 +1120,10 @@ func (b *Builder) addEdge(ctx context.Context, edge *models.ChannelEdgeInfo,
 // considered as not fully constructed.
 //
 // NOTE: This method is part of the ChannelGraphSource interface.
-func (b *Builder) UpdateEdge(update *models.ChannelEdgePolicy,
-	op ...batch.SchedulerOption) error {
+func (b *Builder) UpdateEdge(ctx context.Context,
+	update *models.ChannelEdgePolicy, op ...batch.SchedulerOption) error {
 
-	err := b.updateEdge(update, op...)
+	err := b.updateEdge(ctx, update, op...)
 	if err != nil {
 		logNetworkMsgProcessError(err)
 
@@ -1135,8 +1137,8 @@ func (b *Builder) UpdateEdge(update *models.ChannelEdgePolicy,
 // persisted in the graph, and then applies it to the graph if the update is
 // considered fresh enough and if we actually have a channel persisted for the
 // given update.
-func (b *Builder) updateEdge(policy *models.ChannelEdgePolicy,
-	op ...batch.SchedulerOption) error {
+func (b *Builder) updateEdge(ctx context.Context,
+	policy *models.ChannelEdgePolicy, op ...batch.SchedulerOption) error {
 
 	log.Debugf("Received ChannelEdgePolicy for channel %v",
 		policy.ChannelID)
@@ -1209,7 +1211,7 @@ func (b *Builder) updateEdge(policy *models.ChannelEdgePolicy,
 
 	// Now that we know this isn't a stale update, we'll apply the new edge
 	// policy to the proper directional edge within the channel graph.
-	if err = b.cfg.Graph.UpdateEdgePolicy(policy, op...); err != nil {
+	if err = b.cfg.Graph.UpdateEdgePolicy(ctx, policy, op...); err != nil {
 		err := errors.Errorf("unable to add channel: %v", err)
 		log.Error(err)
 		return err

graph/builder_test.go

@@ -172,7 +172,7 @@ func TestIgnoreChannelEdgePolicyForUnknownChannel(t *testing.T) {
 
 	// Attempt to update the edge. This should be ignored, since the edge
 	// is not yet added to the router.
-	err = ctx.builder.UpdateEdge(edgePolicy)
+	err = ctx.builder.UpdateEdge(ctxb, edgePolicy)
 	if !IsError(err, ErrIgnored) {
 		t.Fatalf("expected to get ErrIgnore, instead got: %v", err)
 	}
@@ -183,7 +183,7 @@ func TestIgnoreChannelEdgePolicyForUnknownChannel(t *testing.T) {
 			"though the vertexes were unknown: %v.", err)
 
 	// Now updating the edge policy should succeed.
-	require.NoError(t, ctx.builder.UpdateEdge(edgePolicy))
+	require.NoError(t, ctx.builder.UpdateEdge(ctxb, edgePolicy))
 }
 
 // TestWakeUpOnStaleBranch tests that upon startup of the ChannelRouter, if the
@@ -1205,7 +1205,7 @@ func TestIsStaleEdgePolicy(t *testing.T) {
 		FeeProportionalMillionths: 10000,
 	}
 	edgePolicy.ChannelFlags = 0
-	if err := ctx.builder.UpdateEdge(edgePolicy); err != nil {
+	if err := ctx.builder.UpdateEdge(ctxb, edgePolicy); err != nil {
 		t.Fatalf("unable to update edge policy: %v", err)
 	}
 
@@ -1219,7 +1219,7 @@ func TestIsStaleEdgePolicy(t *testing.T) {
 		FeeProportionalMillionths: 10000,
 	}
 	edgePolicy.ChannelFlags = 1
-	if err := ctx.builder.UpdateEdge(edgePolicy); err != nil {
+	if err := ctx.builder.UpdateEdge(ctxb, edgePolicy); err != nil {
 		t.Fatalf("unable to update edge policy: %v", err)
 	}
 
@@ -1543,7 +1543,7 @@ func parseTestGraph(t *testing.T, useCache bool, path string) (
 			),
 			ToNode: targetNode,
 		}
-		if err := graph.UpdateEdgePolicy(edgePolicy); err != nil {
+		if err := graph.UpdateEdgePolicy(ctx, edgePolicy); err != nil {
 			return nil, err
 		}
 
@@ -1912,7 +1912,7 @@ func createTestGraphFromChannels(t *testing.T, useCache bool,
 				ToNode:                    node2Vertex,
 				ExtraOpaqueData:           getExtraData(node1),
 			}
-			err := graph.UpdateEdgePolicy(edgePolicy)
+			err := graph.UpdateEdgePolicy(ctx, edgePolicy)
 			if err != nil {
 				return nil, err
 			}
@@ -1943,7 +1943,7 @@ func createTestGraphFromChannels(t *testing.T, useCache bool,
 				ToNode:                    node1Vertex,
 				ExtraOpaqueData:           getExtraData(node2),
 			}
-			err := graph.UpdateEdgePolicy(edgePolicy)
+			err := graph.UpdateEdgePolicy(ctx, edgePolicy)
 			if err != nil {
 				return nil, err
 			}

graph/db/graph.go

@@ -565,10 +565,10 @@ func (c *ChannelGraph) MarkEdgeZombie(chanID uint64,
 // updated, otherwise it's the second node's information. The node ordering is
 // determined by the lexicographical ordering of the identity public keys of the
 // nodes on either side of the channel.
-func (c *ChannelGraph) UpdateEdgePolicy(edge *models.ChannelEdgePolicy,
-	op ...batch.SchedulerOption) error {
+func (c *ChannelGraph) UpdateEdgePolicy(ctx context.Context,
+	edge *models.ChannelEdgePolicy, op ...batch.SchedulerOption) error {
 
-	from, to, err := c.V1Store.UpdateEdgePolicy(edge, op...)
+	from, to, err := c.V1Store.UpdateEdgePolicy(ctx, edge, op...)
 	if err != nil {
 		return err
 	}

graph/db/graph_test.go

@@ -832,7 +832,7 @@ func TestEdgeInfoUpdates(t *testing.T) {
 
 	// Make sure inserting the policy at this point, before the edge info
 	// is added, will fail.
-	err := graph.UpdateEdgePolicy(edge1)
+	err := graph.UpdateEdgePolicy(ctx, edge1)
 	require.ErrorIs(t, err, ErrEdgeNotFound)
 	require.Len(t, graph.graphCache.nodeChannels, 0)
 
@@ -847,11 +847,11 @@ func TestEdgeInfoUpdates(t *testing.T) {
 
 	// Next, insert both edge policies into the database, they should both
 	// be inserted without any issues.
-	if err := graph.UpdateEdgePolicy(edge1); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 	assertEdgeWithPolicyInCache(t, graph, edgeInfo, edge1, true)
-	if err := graph.UpdateEdgePolicy(edge2); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 	assertEdgeWithPolicyInCache(t, graph, edgeInfo, edge2, false)
@@ -1227,14 +1227,14 @@ func TestForEachSourceNodeChannel(t *testing.T) {
 	if !bytes.Equal(abPolicy1.ToNode[:], nodeB.PubKeyBytes[:]) {
 		abPolicyAOutgoing = abPolicy2
 	}
-	require.NoError(t, graph.UpdateEdgePolicy(abPolicyAOutgoing))
+	require.NoError(t, graph.UpdateEdgePolicy(ctx, abPolicyAOutgoing))
 
 	// Now, set the incoming policy for the A-C channel.
 	acPolicyAIncoming := acPolicy1
 	if !bytes.Equal(acPolicy1.ToNode[:], nodeA.PubKeyBytes[:]) {
 		acPolicyAIncoming = acPolicy2
 	}
-	require.NoError(t, graph.UpdateEdgePolicy(acPolicyAIncoming))
+	require.NoError(t, graph.UpdateEdgePolicy(ctx, acPolicyAIncoming))
 
 	type sourceNodeChan struct {
 		otherNode  route.Vertex
@@ -1566,7 +1566,7 @@ func fillTestGraph(t testing.TB, graph *ChannelGraph, numNodes,
 			edge.ChannelFlags = 0
 			edge.ToNode = node2.PubKeyBytes
 			edge.SigBytes = testSig.Serialize()
-			require.NoError(t, graph.UpdateEdgePolicy(edge))
+			require.NoError(t, graph.UpdateEdgePolicy(ctx, edge))
 
 			// Create another random edge that points from
 			// node2 -> node1 this time.
@@ -1574,7 +1574,7 @@ func fillTestGraph(t testing.TB, graph *ChannelGraph, numNodes,
 			edge.ChannelFlags = 1
 			edge.ToNode = node1.PubKeyBytes
 			edge.SigBytes = testSig.Serialize()
-			require.NoError(t, graph.UpdateEdgePolicy(edge))
+			require.NoError(t, graph.UpdateEdgePolicy(ctx, edge))
 
 			chanIndex[chanID] = struct{}{}
 		}
@@ -1764,7 +1764,7 @@ func TestGraphPruning(t *testing.T) {
 		edge.ChannelFlags = 0
 		edge.ToNode = graphNodes[i].PubKeyBytes
 		edge.SigBytes = testSig.Serialize()
-		if err := graph.UpdateEdgePolicy(edge); err != nil {
+		if err := graph.UpdateEdgePolicy(ctx, edge); err != nil {
 			t.Fatalf("unable to update edge: %v", err)
 		}
 
@@ -1774,7 +1774,7 @@ func TestGraphPruning(t *testing.T) {
 		edge.ChannelFlags = 1
 		edge.ToNode = graphNodes[i].PubKeyBytes
 		edge.SigBytes = testSig.Serialize()
-		if err := graph.UpdateEdgePolicy(edge); err != nil {
+		if err := graph.UpdateEdgePolicy(ctx, edge); err != nil {
 			t.Fatalf("unable to update edge: %v", err)
 		}
 	}
@@ -1985,7 +1985,7 @@ func TestChanUpdatesInHorizon(t *testing.T) {
 		edge1.ChannelFlags = 0
 		edge1.ToNode = node2.PubKeyBytes
 		edge1.SigBytes = testSig.Serialize()
-		if err := graph.UpdateEdgePolicy(edge1); err != nil {
+		if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 			t.Fatalf("unable to update edge: %v", err)
 		}
 
@@ -1995,7 +1995,7 @@ func TestChanUpdatesInHorizon(t *testing.T) {
 		edge2.ChannelFlags = 1
 		edge2.ToNode = node1.PubKeyBytes
 		edge2.SigBytes = testSig.Serialize()
-		if err := graph.UpdateEdgePolicy(edge2); err != nil {
+		if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 			t.Fatalf("unable to update edge: %v", err)
 		}
 
@@ -2770,12 +2770,14 @@ func TestFilterChannelRange(t *testing.T) {
 		var updateTime = time.Unix(0, 0)
 		if rand.Int31n(2) == 0 {
 			updateTime = time.Unix(updateTimeSeed, 0)
-			err = graph.UpdateEdgePolicy(&models.ChannelEdgePolicy{
-				ToNode:       node.PubKeyBytes,
-				ChannelFlags: chanFlags,
-				ChannelID:    chanID,
-				LastUpdate:   updateTime,
-			})
+			err = graph.UpdateEdgePolicy(
+				ctx, &models.ChannelEdgePolicy{
+					ToNode:       node.PubKeyBytes,
+					ChannelFlags: chanFlags,
+					ChannelID:    chanID,
+					LastUpdate:   updateTime,
+				},
+			)
 			require.NoError(t, err)
 		}
 		updateTimeSeed++
@@ -2980,7 +2982,7 @@ func TestFetchChanInfos(t *testing.T) {
 		edge1.ChannelFlags = 0
 		edge1.ToNode = node2.PubKeyBytes
 		edge1.SigBytes = testSig.Serialize()
-		if err := graph.UpdateEdgePolicy(edge1); err != nil {
+		if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 			t.Fatalf("unable to update edge: %v", err)
 		}
 
@@ -2988,7 +2990,7 @@ func TestFetchChanInfos(t *testing.T) {
 		edge2.ChannelFlags = 1
 		edge2.ToNode = node1.PubKeyBytes
 		edge2.SigBytes = testSig.Serialize()
-		if err := graph.UpdateEdgePolicy(edge2); err != nil {
+		if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 			t.Fatalf("unable to update edge: %v", err)
 		}
 
@@ -3115,7 +3117,7 @@ func TestIncompleteChannelPolicies(t *testing.T) {
 	edgePolicy.ChannelFlags = 0
 	edgePolicy.ToNode = node2.PubKeyBytes
 	edgePolicy.SigBytes = testSig.Serialize()
-	if err := graph.UpdateEdgePolicy(edgePolicy); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edgePolicy); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 
@@ -3128,7 +3130,7 @@ func TestIncompleteChannelPolicies(t *testing.T) {
 	edgePolicy.ChannelFlags = 1
 	edgePolicy.ToNode = node1.PubKeyBytes
 	edgePolicy.SigBytes = testSig.Serialize()
-	if err := graph.UpdateEdgePolicy(edgePolicy); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edgePolicy); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 
@@ -3178,7 +3180,7 @@ func TestChannelEdgePruningUpdateIndexDeletion(t *testing.T) {
 	edge1.ChannelFlags = 0
 	edge1.ToNode = node1.PubKeyBytes
 	edge1.SigBytes = testSig.Serialize()
-	if err := graph.UpdateEdgePolicy(edge1); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 	edge1 = copyEdgePolicy(edge1) // Avoid read/write race conditions.
@@ -3187,7 +3189,7 @@ func TestChannelEdgePruningUpdateIndexDeletion(t *testing.T) {
 	edge2.ChannelFlags = 1
 	edge2.ToNode = node2.PubKeyBytes
 	edge2.SigBytes = testSig.Serialize()
-	if err := graph.UpdateEdgePolicy(edge2); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 	edge2 = copyEdgePolicy(edge2) // Avoid read/write race conditions.
@@ -3254,12 +3256,12 @@ func TestChannelEdgePruningUpdateIndexDeletion(t *testing.T) {
 	// removed from the update index.
 	edge1.ChannelFlags = 2
 	edge1.LastUpdate = time.Now()
-	if err := graph.UpdateEdgePolicy(edge1); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 	edge2.ChannelFlags = 3
 	edge2.LastUpdate = edge1.LastUpdate.Add(time.Hour)
-	if err := graph.UpdateEdgePolicy(edge2); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 
@@ -3329,7 +3331,7 @@ func TestPruneGraphNodes(t *testing.T) {
 	edge1.ChannelFlags = 0
 	edge1.ToNode = node1.PubKeyBytes
 	edge1.SigBytes = testSig.Serialize()
-	if err := graph.UpdateEdgePolicy(edge1); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 
@@ -3613,7 +3615,7 @@ func TestDisabledChannelIDs(t *testing.T) {
 	// Add one disabled policy and ensure the channel is still not in the
 	// disabled list.
 	edge1.ChannelFlags |= lnwire.ChanUpdateDisabled
-	if err := graph.UpdateEdgePolicy(edge1); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 	disabledChanIds, err = graph.DisabledChannelIDs()
@@ -3626,7 +3628,7 @@ func TestDisabledChannelIDs(t *testing.T) {
 	// Add second disabled policy and ensure the channel is now in the
 	// disabled list.
 	edge2.ChannelFlags |= lnwire.ChanUpdateDisabled
-	if err := graph.UpdateEdgePolicy(edge2); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 	disabledChanIds, err = graph.DisabledChannelIDs()
@@ -3752,7 +3754,7 @@ func TestEdgePolicyMissingMaxHtcl(t *testing.T) {
 	require.NoError(t, err, "error writing db")
 
 	// And add the second, unmodified edge.
-	if err := graph.UpdateEdgePolicy(edge2); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 
@@ -3774,7 +3776,7 @@ func TestEdgePolicyMissingMaxHtcl(t *testing.T) {
 
 	// Now add the original, unmodified edge policy, and make sure the edge
 	// policies then become fully populated.
-	if err := graph.UpdateEdgePolicy(edge1); err != nil {
+	if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
 
@@ -4082,7 +4084,7 @@ func TestBatchedUpdateEdgePolicy(t *testing.T) {
 
 	// Make sure inserting the policy at this point, before the edge info
 	// is added, will fail.
-	require.Error(t, ErrEdgeNotFound, graph.UpdateEdgePolicy(edge1))
+	require.Error(t, ErrEdgeNotFound, graph.UpdateEdgePolicy(ctx, edge1))
 
 	// Add the edge info.
 	require.NoError(t, graph.AddChannelEdge(ctx, edgeInfo))
@@ -4101,7 +4103,7 @@ func TestBatchedUpdateEdgePolicy(t *testing.T) {
 			defer wg.Done()
 
 			select {
-			case errChan <- graph.UpdateEdgePolicy(update):
+			case errChan <- graph.UpdateEdgePolicy(ctx, update):
 			case <-time.After(2 * time.Second):
 				errChan <- errTimeout
 			}
@@ -4226,7 +4228,7 @@ func TestGraphCacheForEachNodeChannel(t *testing.T) {
 		FeeRate: 20,
 	}
 	edge1.InboundFee = fn.Some(inboundFee)
-	require.NoError(t, graph.UpdateEdgePolicy(edge1))
+	require.NoError(t, graph.UpdateEdgePolicy(ctx, edge1))
 	edge1 = copyEdgePolicy(edge1) // Avoid read/write race conditions.
 
 	directedChan := getSingleChannel()
@@ -4241,7 +4243,7 @@ func TestGraphCacheForEachNodeChannel(t *testing.T) {
 	// error when we try to update the edge policy.
 	edge1.LastUpdate = edge1.LastUpdate.Add(time.Second)
 	require.ErrorIs(
-		t, graph.UpdateEdgePolicy(edge1), ErrParsingExtraTLVBytes,
+		t, graph.UpdateEdgePolicy(ctx, edge1), ErrParsingExtraTLVBytes,
 	)
 
 	// Since persistence of the last update failed, we should still bet

graph/db/interfaces.go

@@ -332,7 +332,7 @@ type V1Store interface { //nolint:interfacebloat
 	// node's information. The node ordering is determined by the
 	// lexicographical ordering of the identity public keys of the nodes on
 	// either side of the channel.
-	UpdateEdgePolicy(edge *models.ChannelEdgePolicy,
+	UpdateEdgePolicy(ctx context.Context, edge *models.ChannelEdgePolicy,
 		op ...batch.SchedulerOption) (route.Vertex, route.Vertex, error)
 
 	// SourceNode returns the source node of the graph. The source node is

graph/db/kv_store.go

@@ -2801,11 +2801,11 @@ func makeZombiePubkeys(info *models.ChannelEdgeInfo,
 // updated, otherwise it's the second node's information. The node ordering is
 // determined by the lexicographical ordering of the identity public keys of the
 // nodes on either side of the channel.
-func (c *KVStore) UpdateEdgePolicy(edge *models.ChannelEdgePolicy,
+func (c *KVStore) UpdateEdgePolicy(ctx context.Context,
+	edge *models.ChannelEdgePolicy,
 	opts ...batch.SchedulerOption) (route.Vertex, route.Vertex, error) {
 
 	var (
-		ctx          = context.TODO()
 		isUpdate1    bool
 		edgeNotFound bool
 		from, to     route.Vertex

graph/db/sql_store.go

@@ -577,11 +577,10 @@ func (s *SQLStore) HighestChanID() (uint64, error) {
 // nodes on either side of the channel.
 //
 // NOTE: part of the V1Store interface.
-func (s *SQLStore) UpdateEdgePolicy(edge *models.ChannelEdgePolicy,
+func (s *SQLStore) UpdateEdgePolicy(ctx context.Context,
+	edge *models.ChannelEdgePolicy,
 	opts ...batch.SchedulerOption) (route.Vertex, route.Vertex, error) {
 
-	ctx := context.TODO()
-
 	var (
 		isUpdate1    bool
 		edgeNotFound bool

graph/interfaces.go

@@ -39,7 +39,7 @@ type ChannelGraphSource interface {
 
 	// UpdateEdge is used to update edge information, without this message
 	// edge considered as not fully constructed.
-	UpdateEdge(policy *models.ChannelEdgePolicy,
+	UpdateEdge(ctx context.Context, policy *models.ChannelEdgePolicy,
 		op ...batch.SchedulerOption) error
 
 	// IsStaleNode returns true if the graph source has a node announcement
@@ -231,7 +231,7 @@ type DB interface {
 	// node's information. The node ordering is determined by the
 	// lexicographical ordering of the identity public keys of the nodes on
 	// either side of the channel.
-	UpdateEdgePolicy(edge *models.ChannelEdgePolicy,
+	UpdateEdgePolicy(ctx context.Context, edge *models.ChannelEdgePolicy,
 		op ...batch.SchedulerOption) error
 
 	// HasLightningNode determines if the graph has a vertex identified by

graph/notifications_test.go

@@ -484,10 +484,10 @@ func TestEdgeUpdateNotification(t *testing.T) {
 	require.NoError(t, err, "unable to create a random chan policy")
 	edge2.ChannelFlags = 1
 
-	if err := ctx.builder.UpdateEdge(edge1); err != nil {
+	if err := ctx.builder.UpdateEdge(ctxb, edge1); err != nil {
 		t.Fatalf("unable to add edge update: %v", err)
 	}
-	if err := ctx.builder.UpdateEdge(edge2); err != nil {
+	if err := ctx.builder.UpdateEdge(ctxb, edge2); err != nil {
 		t.Fatalf("unable to add edge update: %v", err)
 	}