Fix race condition in probing tests

Changed the exhaust test to be statistical (locked amount never exceeds
the cap) instead of trying to turn intermediary routing node offline
when it hasn't yet forwarded the probe htlc.

Author: randomlogin

tests/probing_tests.rs

@@ -7,9 +7,8 @@
 //      to zero once the probe resolves.
 //
 //   exhausted_probe_budget_blocks_new_probes
-//      Stops B mid-flight so the HTLC cannot resolve; confirms the budget
-//      stays exhausted and no further probes are sent. After B restarts
-//      the probe fails, the budget clears, and new probes resume.
+//      Samples locked_msat across multiple probe cycles and asserts it never
+//      exceeds the configured max_locked_msat budget cap.
 
 mod common;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -94,7 +93,7 @@ fn build_probe_path(
 				short_channel_id: ch_ab.short_channel_id.unwrap(),
 				channel_features: ChannelFeatures::empty(),
 				fee_msat: 1000,
-				cltv_expiry_delta: 40,
+				cltv_expiry_delta: 144,
 				maybe_announced_channel: true,
 			},
 			RouteHop {
@@ -103,7 +102,7 @@ fn build_probe_path(
 				short_channel_id: ch_bc.short_channel_id.unwrap(),
 				channel_features: ChannelFeatures::empty(),
 				fee_msat: amount_msat,
-				cltv_expiry_delta: 0,
+				cltv_expiry_delta: 18,
 				maybe_announced_channel: true,
 			},
 		],
@@ -194,11 +193,7 @@ async fn probe_budget_increments_and_decrements() {
 	node_c.stop().unwrap();
 }
 
-/// Verifies that no new probes are dispatched once the in-flight budget is exhausted.
-///
-/// 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.
+/// Verifies that `locked_msat` never exceeds `max_locked_msat` across multiple probe cycles.
 #[tokio::test(flavor = "multi_thread")]
 async fn exhausted_probe_budget_blocks_new_probes() {
 	let (bitcoind, electrsd) = setup_bitcoind_and_electrsd();
@@ -209,10 +204,11 @@ async fn exhausted_probe_budget_blocks_new_probes() {
 
 	let mut config_a = random_config(false);
 	let strategy = FixedPathStrategy::new();
+	let max_locked_msat = 2 * PROBE_AMOUNT_MSAT;
 	config_a.probing = Some(
 		ProbingConfigBuilder::custom(strategy.clone())
 			.interval(Duration::from_millis(PROBING_INTERVAL_MILLISECONDS))
-			.max_locked_msat(2 * PROBE_AMOUNT_MSAT)
+			.max_locked_msat(max_locked_msat)
 			.build(),
 	);
 	let node_a = setup_node(&chain_source, config_a);
@@ -244,100 +240,28 @@ async fn exhausted_probe_budget_blocks_new_probes() {
 	expect_event!(node_b, ChannelReady);
 	expect_event!(node_c, ChannelReady);
 
-	let capacity_at_open = node_a
-		.list_channels()
-		.iter()
-		.find(|ch| ch.counterparty_node_id == node_b.node_id())
-		.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.set_path(build_probe_path(&node_a, &node_b, &node_c, PROBE_AMOUNT_MSAT));
 	tokio::time::sleep(Duration::from_secs(3)).await;
 	strategy.start_probing();
 
-	// Wait for the first probe to be in-flight.
-	let locked = tokio::time::timeout(Duration::from_secs(30), async {
-		loop {
-			if node_a.prober().map_or(0, |p| p.locked_msat()) > 0 {
-				break;
-			}
-			tokio::time::sleep(Duration::from_millis(100)).await;
+	// Sample locked_msat across multiple probe cycles and assert the budget cap is never exceeded
+	let mut observed_locked = false;
+	let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
+	while tokio::time::Instant::now() < deadline {
+		let msat = node_a.prober().map_or(0, |p| p.locked_msat());
+		if msat > 0 {
+			observed_locked = true;
 		}
-	})
-	.await
-	.is_ok();
-	assert!(locked, "no probe dispatched within 30 s");
-
-	// Capacity should have decreased due to the in-flight probe HTLC.
-	let capacity_with_probe = node_a
-		.list_channels()
-		.iter()
-		.find(|ch| ch.counterparty_node_id == node_b.node_id())
-		.map(|ch| ch.outbound_capacity_msat)
-		.expect("A→B channel not found");
-	assert!(
-		capacity_with_probe < capacity_at_open,
-		"HTLC not visible in channel state: capacity unchanged ({capacity_at_open} msat)"
-	);
-
-	// Stop B while the probe HTLC is in-flight.
-	node_b.stop().unwrap();
-	// Pause probing so the budget can clear without a new probe re-locking it.
-	strategy.stop_probing();
-
-	tokio::time::sleep(Duration::from_secs(5)).await;
-	assert!(
-		node_a.prober().map_or(0, |p| p.locked_msat()) > 0,
-		"probe resolved unexpectedly while B was offline"
-	);
-	let capacity_after_wait = node_a
-		.list_channels()
-		.iter()
-		.find(|ch| ch.counterparty_node_id == node_b.node_id())
-		.map(|ch| ch.outbound_capacity_msat)
-		.unwrap_or(u64::MAX);
-	assert!(
-		capacity_after_wait >= capacity_with_probe,
-		"a new probe HTLC was sent despite budget being exhausted"
-	);
-
-	// strategy.stop_probing();
-
-	// Bring B back and explicitly reconnect to A and C so the stuck HTLC resolves
-	// without waiting for the background reconnection backoff.
-	node_b.start().unwrap();
-	let node_a_addr = node_a.listening_addresses().unwrap().first().unwrap().clone();
-	let node_c_addr = node_c.listening_addresses().unwrap().first().unwrap().clone();
-	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 = tokio::time::timeout(Duration::from_secs(60), async {
-		loop {
-			if node_a.prober().map_or(1, |p| p.locked_msat()) == 0 {
-				break;
-			}
-			tokio::time::sleep(Duration::from_millis(100)).await;
-		}
-	})
-	.await
-	.is_ok();
-	assert!(cleared, "locked_msat never cleared after B came back online");
+		assert!(
+			msat <= max_locked_msat,
+			"locked_msat {msat} exceeded budget cap {max_locked_msat}"
+		);
+		tokio::time::sleep(Duration::from_millis(25)).await;
+	}
 
-	// Re-enable probing; a new probe should be dispatched within a few ticks.
-	strategy.start_probing();
-	let new_probe = tokio::time::timeout(Duration::from_secs(60), async {
-		loop {
-			if node_a.prober().map_or(0, |p| p.locked_msat()) > 0 {
-				break;
-			}
-			tokio::time::sleep(Duration::from_millis(100)).await;
-		}
-	})
-	.await
-	.is_ok();
-	assert!(new_probe, "no new probe dispatched after budget was freed");
+	assert!(observed_locked, "no probe was dispatched during the observation window");
 
 	node_a.stop().unwrap();
 	node_b.stop().unwrap();