Fix probing tests

Added an explicit start of probing, which eliminated previously occured
race condition

Author: randomlogin

tests/probing_tests.rs

@@ -19,6 +19,7 @@
 //      runs payment rounds and prints probing perfomance tables.
 
 mod common;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 use lightning::routing::gossip::NodeAlias;
 use lightning_invoice::{Bolt11InvoiceDescription, Description};
@@ -49,31 +50,27 @@ const PROBING_DIVERSITY_PENALTY: u64 = 50_000;
 struct FixedDestStrategy {
 	destination: PublicKey,
 	amount_msat: u64,
+	ready_to_probe: AtomicBool,
 }
 
 impl FixedDestStrategy {
 	fn new(destination: PublicKey, amount_msat: u64) -> Arc<Self> {
-		Arc::new(Self { destination, amount_msat })
+		Arc::new(Self { destination, amount_msat, ready_to_probe: AtomicBool::new(false) })
 	}
-}
 
-impl ProbingStrategy for FixedDestStrategy {
-	fn next_probe(&self) -> Option<Probe> {
-		Some(Probe::Destination { final_node: self.destination, amount_msat: self.amount_msat })
+	fn start_probing(&self) {
+		self.ready_to_probe.store(true, Ordering::Relaxed);
 	}
 }
 
-async fn wait_until(timeout: Duration, predicate: impl Fn() -> bool) -> bool {
-	tokio::time::timeout(timeout, async {
-		loop {
-			if predicate() {
-				return;
-			}
-			tokio::time::sleep(Duration::from_millis(1)).await;
+impl ProbingStrategy for FixedDestStrategy {
+	fn next_probe(&self) -> Option<Probe> {
+		if self.ready_to_probe.load(Ordering::Relaxed) {
+			Some(Probe::Destination { final_node: self.destination, amount_msat: self.amount_msat })
+		} else {
+			None
 		}
-	})
-	.await
-	.is_ok()
+	}
 }
 
 fn config_with_label(label: &str) -> common::TestConfig {
@@ -85,20 +82,6 @@ fn config_with_label(label: &str) -> common::TestConfig {
 	config
 }
 
-fn build_node_fixed_dest_probing(
-	chain_source: &TestChainSource<'_>, destination_node_id: PublicKey,
-) -> TestNode {
-	let mut config = random_config(false);
-	let strategy = FixedDestStrategy::new(destination_node_id, PROBE_AMOUNT_MSAT);
-	config.probing = Some(
-		ProbingConfig::custom(strategy)
-			.interval(Duration::from_millis(PROBING_INTERVAL_MILLISECONDS))
-			.max_locked_msat(PROBE_AMOUNT_MSAT)
-			.build(),
-	);
-	setup_node(chain_source, config)
-}
-
 fn build_node_random_probing(chain_source: &TestChainSource<'_>, max_hops: usize) -> TestNode {
 	let mut config = config_with_label("Random");
 	config.probing = Some(
@@ -259,14 +242,23 @@ fn print_probing_perfomance(observers: &[&TestNode], all_nodes: &[&TestNode]) {
 
 /// Verifies that `locked_msat` increases when a probe is dispatched and returns
 /// to zero once the probe resolves (succeeds or fails).
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn probe_budget_increments_and_decrements() {
 	let (bitcoind, electrsd) = setup_bitcoind_and_electrsd();
 	let chain_source = TestChainSource::Electrum(&electrsd);
 
 	let node_b = setup_node(&chain_source, random_config(false));
 	let node_c = setup_node(&chain_source, random_config(false));
-	let node_a = build_node_fixed_dest_probing(&chain_source, node_c.node_id());
+
+	let mut config_a = random_config(false);
+	let strategy = FixedDestStrategy::new(node_c.node_id(), PROBE_AMOUNT_MSAT);
+	config_a.probing = Some(
+		ProbingConfig::custom(strategy.clone())
+			.interval(Duration::from_millis(PROBING_INTERVAL_MILLISECONDS))
+			.max_locked_msat(PROBE_AMOUNT_MSAT)
+			.build(),
+	);
+	let node_a = setup_node(&chain_source, config_a);
 
 	let addr_a = node_a.onchain_payment().new_address().unwrap();
 	let addr_b = node_b.onchain_payment().new_address().unwrap();
@@ -295,18 +287,33 @@ async fn probe_budget_increments_and_decrements() {
 	expect_event!(node_b, ChannelReady);
 	expect_event!(node_c, ChannelReady);
 
-	// Give gossip time to propagate to A.
+	// Give gossip time to propagate to A, then enable probing.
 	tokio::time::sleep(Duration::from_secs(3)).await;
+	strategy.start_probing();
 
-	let went_up =
-		wait_until(Duration::from_secs(10), || node_a.prober().map_or(0, |p| p.locked_msat()) > 0)
-			.await;
+	let went_up = tokio::time::timeout(Duration::from_secs(10), async {
+		loop {
+			if node_a.prober().map_or(0, |p| p.locked_msat()) > 0 {
+				break;
+			}
+			tokio::time::sleep(Duration::from_millis(1)).await;
+		}
+	})
+	.await
+	.is_ok();
 	assert!(went_up, "locked_msat never increased — no probe was dispatched");
 	println!("First probe dispatched; locked_msat = {}", node_a.prober().unwrap().locked_msat());
 
-	let cleared =
-		wait_until(Duration::from_secs(20), || node_a.prober().map_or(1, |p| p.locked_msat()) == 0)
-			.await;
+	let cleared = tokio::time::timeout(Duration::from_secs(20), async {
+		loop {
+			if node_a.prober().map_or(1, |p| p.locked_msat()) == 0 {
+				break;
+			}
+			tokio::time::sleep(Duration::from_millis(1)).await;
+		}
+	})
+	.await
+	.is_ok();
 	assert!(cleared, "locked_msat never returned to zero after probe resolved");
 	println!("Probe resolved; locked_msat = 0");
 
@@ -320,7 +327,7 @@ async fn probe_budget_increments_and_decrements() {
 /// Exhaustion is triggered by stopping the intermediate node (B) while a probe HTLC
 /// is in-flight, preventing resolution and keeping the budget locked. After B restarts
 /// the HTLC fails, the budget clears, and probing resumes.
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn exhausted_probe_budget_blocks_new_probes() {
 	let (bitcoind, electrsd) = setup_bitcoind_and_electrsd();
 	let chain_source = TestChainSource::Electrum(&electrsd);
@@ -332,8 +339,8 @@ async fn exhausted_probe_budget_blocks_new_probes() {
 	let mut config_a = random_config(false);
 	let strategy = FixedDestStrategy::new(node_c.node_id(), PROBE_AMOUNT_MSAT);
 	config_a.probing = Some(
-		ProbingConfig::custom(strategy)
-			.interval(Duration::from_secs(3))
+		ProbingConfig::custom(strategy.clone())
+			.interval(Duration::from_secs(10))
 			.max_locked_msat(PROBE_AMOUNT_MSAT)
 			.build(),
 	);
@@ -374,10 +381,20 @@ async fn exhausted_probe_budget_blocks_new_probes() {
 		.map(|ch| ch.outbound_capacity_msat)
 		.expect("A→B channel not found");
 
+	assert_eq!(node_a.prober().map_or(1, |p| p.locked_msat()), 0, "initial locked_msat is nonzero");
+	strategy.start_probing();
+
 	// Give gossip time to propagate to A, then wait for the first probe.
-	let locked =
-		wait_until(Duration::from_secs(15), || node_a.prober().map_or(0, |p| p.locked_msat()) > 0)
-			.await;
+	let locked = tokio::time::timeout(Duration::from_secs(15), async {
+		loop {
+			if node_a.prober().map_or(0, |p| p.locked_msat()) > 0 {
+				break;
+			}
+			tokio::time::sleep(Duration::from_millis(1)).await;
+		}
+	})
+	.await
+	.is_ok();
 	assert!(locked, "no probe dispatched within 15 s");
 
 	// Capacity should have decreased due to the in-flight probe HTLC.
@@ -421,15 +438,29 @@ async fn exhausted_probe_budget_blocks_new_probes() {
 	node_b.connect(node_a.node_id(), node_a_addr, false).unwrap();
 	node_b.connect(node_c.node_id(), node_c_addr, false).unwrap();
 
-	let cleared =
-		wait_until(Duration::from_secs(15), || node_a.prober().map_or(1, |p| p.locked_msat()) == 0)
-			.await;
+	let cleared = tokio::time::timeout(Duration::from_secs(15), async {
+		loop {
+			if node_a.prober().map_or(1, |p| p.locked_msat()) == 0 {
+				break;
+			}
+			tokio::time::sleep(Duration::from_millis(1)).await;
+		}
+	})
+	.await
+	.is_ok();
 	assert!(cleared, "locked_msat never cleared after B came back online");
 
 	// Once the budget is freed, a new probe should be dispatched within a few ticks.
-	let new_probe =
-		wait_until(Duration::from_secs(10), || node_a.prober().map_or(0, |p| p.locked_msat()) > 0)
-			.await;
+	let new_probe = tokio::time::timeout(Duration::from_secs(10), async {
+		loop {
+			if node_a.prober().map_or(0, |p| p.locked_msat()) > 0 {
+				break;
+			}
+			tokio::time::sleep(Duration::from_millis(1)).await;
+		}
+	})
+	.await
+	.is_ok();
 	assert!(new_probe, "no new probe dispatched after budget was freed");
 
 	node_a.stop().unwrap();