Merge pull request #1592 from tnull/2022-07-manual-penalty

Allow to set manual node penalties

Author: TheBlueMatt

lightning/src/routing/router.rs

@@ -1882,7 +1882,7 @@ fn build_route_from_hops_internal<L: Deref>(
 
 #[cfg(test)]
 mod tests {
-	use routing::gossip::{NetworkGraph, P2PGossipSync, NodeId};
+	use routing::gossip::{NetworkGraph, P2PGossipSync, NodeId, EffectiveCapacity};
 	use routing::router::{get_route, build_route_from_hops_internal, add_random_cltv_offset, default_node_features,
 		PaymentParameters, Route, RouteHint, RouteHintHop, RouteHop, RoutingFees,
 		DEFAULT_MAX_TOTAL_CLTV_EXPIRY_DELTA, MAX_PATH_LENGTH_ESTIMATE};
@@ -5727,7 +5727,7 @@ mod tests {
 	}
 
 	#[test]
-	fn avoids_banned_nodes() {
+	fn honors_manual_penalties() {
 		let (secp_ctx, network_graph, _, _, logger) = build_line_graph();
 		let (_, our_id, _, nodes) = get_nodes(&secp_ctx);
 
@@ -5737,7 +5737,17 @@ mod tests {
 		let scorer_params = ProbabilisticScoringParameters::default();
 		let mut scorer = ProbabilisticScorer::new(scorer_params, Arc::clone(&network_graph), Arc::clone(&logger));
 
-		// First check we can get a route.
+		// First check set manual penalties are returned by the scorer.
+		let usage = ChannelUsage {
+			amount_msat: 0,
+			inflight_htlc_msat: 0,
+			effective_capacity: EffectiveCapacity::Total { capacity_msat: 1_024_000, htlc_maximum_msat: Some(1_000) },
+		};
+		scorer.set_manual_penalty(&NodeId::from_pubkey(&nodes[3]), 123);
+		scorer.set_manual_penalty(&NodeId::from_pubkey(&nodes[4]), 456);
+		assert_eq!(scorer.channel_penalty_msat(42, &NodeId::from_pubkey(&nodes[3]), &NodeId::from_pubkey(&nodes[4]), usage), 456);
+
+		// Then check we can get a normal route
 		let payment_params = PaymentParameters::from_node_id(nodes[10]);
 		let route = get_route(&our_id, &payment_params, &network_graph.read_only(), None, 100, 42, Arc::clone(&logger), &scorer, &random_seed_bytes);
 		assert!(route.is_ok());

lightning/src/routing/scoring.rs

@@ -380,10 +380,12 @@ pub struct ProbabilisticScoringParameters {
 	/// Default value: 256 msat
 	pub amount_penalty_multiplier_msat: u64,
 
-	/// A list of nodes that won't be considered during path finding.
+	/// Manual penalties used for the given nodes. Allows to set a particular penalty for a given
+	/// node. Note that a manual penalty of `u64::max_value()` means the node would not ever be
+	/// considered during path finding.
 	///
 	/// (C-not exported)
-	pub banned_nodes: HashSet<NodeId>,
+	pub manual_node_penalties: HashMap<NodeId, u64>,
 
 	/// This penalty is applied when `htlc_maximum_msat` is equal to or larger than half of the
 	/// channel's capacity, which makes us prefer nodes with a smaller `htlc_maximum_msat`. We
@@ -486,17 +488,27 @@ impl<G: Deref<Target = NetworkGraph<L>>, L: Deref, T: Time> ProbabilisticScorerU
 	/// Marks the node with the given `node_id` as banned, i.e.,
 	/// it will be avoided during path finding.
 	pub fn add_banned(&mut self, node_id: &NodeId) {
-		self.params.banned_nodes.insert(*node_id);
+		self.params.manual_node_penalties.insert(*node_id, u64::max_value());
 	}
 
 	/// Removes the node with the given `node_id` from the list of nodes to avoid.
 	pub fn remove_banned(&mut self, node_id: &NodeId) {
-		self.params.banned_nodes.remove(node_id);
+		self.params.manual_node_penalties.remove(node_id);
 	}
 
-	/// Clears the list of nodes that are avoided during path finding.
-	pub fn clear_banned(&mut self) {
-		self.params.banned_nodes = HashSet::new();
+	/// Sets a manual penalty for the given node.
+	pub fn set_manual_penalty(&mut self, node_id: &NodeId, penalty: u64) {
+		self.params.manual_node_penalties.insert(*node_id, penalty);
+	}
+
+	/// Removes the node with the given `node_id` from the list of manual penalties.
+	pub fn remove_manual_penalty(&mut self, node_id: &NodeId) {
+		self.params.manual_node_penalties.remove(node_id);
+	}
+
+	/// Clears the list of manual penalties that are applied during path finding.
+	pub fn clear_manual_penalties(&mut self) {
+		self.params.manual_node_penalties = HashMap::new();
 	}
 }
 
@@ -508,7 +520,7 @@ impl ProbabilisticScoringParameters {
 			liquidity_penalty_multiplier_msat: 0,
 			liquidity_offset_half_life: Duration::from_secs(3600),
 			amount_penalty_multiplier_msat: 0,
-			banned_nodes: HashSet::new(),
+			manual_node_penalties: HashMap::new(),
 			anti_probing_penalty_msat: 0,
 		}
 	}
@@ -517,7 +529,7 @@ impl ProbabilisticScoringParameters {
 	/// they will be avoided during path finding.
 	pub fn add_banned_from_list(&mut self, node_ids: Vec<NodeId>) {
 		for id in node_ids {
-			self.banned_nodes.insert(id);
+			self.manual_node_penalties.insert(id, u64::max_value());
 		}
 	}
 }
@@ -529,7 +541,7 @@ impl Default for ProbabilisticScoringParameters {
 			liquidity_penalty_multiplier_msat: 40_000,
 			liquidity_offset_half_life: Duration::from_secs(3600),
 			amount_penalty_multiplier_msat: 256,
-			banned_nodes: HashSet::new(),
+			manual_node_penalties: HashMap::new(),
 			anti_probing_penalty_msat: 250,
 		}
 	}
@@ -731,8 +743,8 @@ impl<G: Deref<Target = NetworkGraph<L>>, L: Deref, T: Time> Score for Probabilis
 	fn channel_penalty_msat(
 		&self, short_channel_id: u64, source: &NodeId, target: &NodeId, usage: ChannelUsage
 	) -> u64 {
-		if self.params.banned_nodes.contains(source) || self.params.banned_nodes.contains(target) {
-			return u64::max_value();
+		if let Some(penalty) = self.params.manual_node_penalties.get(target) {
+			return *penalty;
 		}
 
 		let mut anti_probing_penalty_msat = 0;