Fix flaky reorg tests

tests/common/mod.rs

@@ -21,17 +21,17 @@ pub(crate) mod lnd;
 use std::collections::{HashMap, HashSet};
 use std::env;
 use std::future::Future;
+use std::net::TcpListener;
 use std::path::PathBuf;
 use std::str::FromStr;
-use std::sync::atomic::{AtomicU16, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
 use bitcoin::hashes::hex::FromHex;
 use bitcoin::hashes::sha256::Hash as Sha256;
 use bitcoin::hashes::Hash;
 use bitcoin::{
-	Address, Amount, Network, OutPoint, ScriptBuf, Sequence, Transaction, Txid, Witness,
+	Address, Amount, BlockHash, Network, OutPoint, ScriptBuf, Sequence, Transaction, Txid, Witness,
 };
 use electrsd::corepc_node::{Client as BitcoindClient, Node as BitcoinD};
 use electrsd::electrum_client::ElectrumApi;
@@ -358,13 +358,14 @@ pub(crate) fn random_storage_path() -> PathBuf {
 	temp_path
 }
 
-static NEXT_PORT: AtomicU16 = AtomicU16::new(20000);
-
 pub(crate) fn generate_listening_addresses() -> Vec<SocketAddress> {
-	let port = NEXT_PORT.fetch_add(2, Ordering::Relaxed);
+	let listener_a = TcpListener::bind(("127.0.0.1", 0)).expect("available listener port");
+	let listener_b = TcpListener::bind(("127.0.0.1", 0)).expect("available listener port");
+	let port_a = listener_a.local_addr().expect("listener address").port();
+	let port_b = listener_b.local_addr().expect("listener address").port();
 	vec![
-		SocketAddress::TcpIpV4 { addr: [127, 0, 0, 1], port },
-		SocketAddress::TcpIpV4 { addr: [127, 0, 0, 1], port: port + 1 },
+		SocketAddress::TcpIpV4 { addr: [127, 0, 0, 1], port: port_a },
+		SocketAddress::TcpIpV4 { addr: [127, 0, 0, 1], port: port_b },
 	]
 }
 
@@ -634,8 +635,12 @@ pub(crate) async fn generate_blocks_and_wait<E: ElectrumApi>(
 	let cur_height = blockchain_info.blocks;
 	let address = bitcoind.new_address().expect("failed to get new address");
 	// TODO: expect this Result once the WouldBlock issue is resolved upstream.
-	let _block_hashes_res = bitcoind.generate_to_address(num, &address);
-	wait_for_block(electrs, cur_height as usize + num).await;
+	let block_hashes_res = bitcoind.generate_to_address(num, &address);
+	let min_height = cur_height as usize + num;
+	let expected_block_hash = block_hashes_res.ok().and_then(|block_hashes| {
+		block_hashes.0.last().map(|hash| hash.parse().expect("block hash should be valid"))
+	});
+	wait_for_bitcoind_block_in_electrs(bitcoind, electrs, min_height, expected_block_hash).await;
 	print!(" Done!");
 	println!("\n");
 }
@@ -655,26 +660,65 @@ pub(crate) fn invalidate_blocks(bitcoind: &BitcoindClient, num_blocks: usize) {
 	assert!(new_cur_height + num_blocks == cur_height);
 }
 
+async fn wait_for_bitcoind_block_in_electrs<E: ElectrumApi>(
+	bitcoind: &BitcoindClient, electrs: &E, min_height: usize,
+	mut expected_block_hash: Option<BlockHash>,
+) {
+	let mut delay = Duration::from_millis(64);
+	let mut tries = 0;
+	loop {
+		let mut height = None;
+		if expected_block_hash.is_none() {
+			let bitcoind_height =
+				bitcoind.get_blockchain_info().expect("failed to get blockchain info").blocks
+					as usize;
+			height = Some(bitcoind_height);
+			if bitcoind_height >= min_height {
+				expected_block_hash = Some(
+					bitcoind
+						.get_block_hash(min_height as u64)
+						.expect("failed to get block hash")
+						.block_hash()
+						.expect("block hash should be present"),
+				);
+			}
+		}
+
+		if let Some(expected_block_hash) = expected_block_hash.as_ref() {
+			if let Ok(header) = electrs.block_header(min_height) {
+				if header.block_hash() == *expected_block_hash {
+					return;
+				}
+			}
+		}
+		assert!(
+			tries < 120,
+			"electrs did not serve bitcoind block header {} within 60 seconds; bitcoind height {:?}, expected hash {:?}",
+			min_height,
+			height,
+			expected_block_hash
+		);
+		tries += 1;
+		tokio::time::sleep(delay).await;
+		if delay.as_millis() < 512 {
+			delay = delay.mul_f32(2.0);
+		}
+	}
+}
+
 pub(crate) async fn wait_for_block<E: ElectrumApi>(electrs: &E, min_height: usize) {
-	let mut header = match electrs.block_headers_subscribe() {
-		Ok(header) => header,
-		Err(_) => {
-			// While subscribing should succeed the first time around, we ran into some cases where
-			// it didn't. Since we can't proceed without subscribing, we try again after a delay
-			// and panic if it still fails.
-			tokio::time::sleep(Duration::from_secs(3)).await;
-			electrs.block_headers_subscribe().expect("failed to subscribe to block headers")
-		},
-	};
+	let mut delay = Duration::from_millis(64);
+	let mut tries = 0;
 	loop {
-		if header.height >= min_height {
-			break;
+		if electrs.block_header(min_height).is_ok() {
+			return;
+		}
+		assert!(tries < 120, "electrs did not serve block header {} within 60 seconds", min_height);
+		tries += 1;
+		tokio::time::sleep(delay).await;
+		if delay.as_millis() < 512 {
+			delay = delay.mul_f32(2.0);
 		}
-		header = exponential_backoff_poll(|| {
-			electrs.ping().expect("failed to ping electrs");
-			electrs.block_headers_pop().expect("failed to pop block header")
-		})
-		.await;
 	}
 }
 
@@ -729,6 +773,23 @@ where
 	}
 }
 
+pub(crate) async fn stop_nodes_concurrently(nodes: Vec<TestNode>) {
+	let stop_receivers = nodes
+		.into_iter()
+		.map(|node| {
+			let (stop_sender, stop_receiver) = tokio::sync::oneshot::channel();
+			std::thread::spawn(move || {
+				let _ = stop_sender.send(node.stop());
+			});
+			stop_receiver
+		})
+		.collect::<Vec<_>>();
+
+	for stop_receiver in stop_receivers {
+		stop_receiver.await.expect("node stop thread panicked").unwrap();
+	}
+}
+
 pub(crate) async fn premine_and_distribute_funds<E: ElectrumApi>(
 	bitcoind: &BitcoindClient, electrs: &E, addrs: Vec<Address>, amount: Amount,
 ) {

tests/reorg_test.rs

@@ -1,5 +1,6 @@
 mod common;
 use std::collections::HashMap;
+use std::time::Duration;
 
 use bitcoin::Amount;
 use ldk_node::payment::{PaymentDirection, PaymentKind};
@@ -10,9 +11,34 @@ use proptest::proptest;
 use crate::common::{
 	expect_event, generate_blocks_and_wait, invalidate_blocks, open_channel,
 	premine_and_distribute_funds, random_chain_source, random_config, setup_bitcoind_and_electrsd,
-	setup_node, wait_for_outpoint_spend,
+	setup_node, stop_nodes_concurrently, wait_for_outpoint_spend,
 };
 
+async fn wait_for_pending_sweep_balance<F>(node: &ldk_node::Node, mut matches_pending_balance: F)
+where
+	F: FnMut(&PendingSweepBalance) -> bool,
+{
+	let mut delay = Duration::from_millis(64);
+	let mut tries = 0;
+	loop {
+		node.sync_wallets().unwrap();
+		let balances = node.list_balances();
+		if balances
+			.pending_balances_from_channel_closures
+			.iter()
+			.any(|balance| matches_pending_balance(balance))
+		{
+			return;
+		}
+		assert!(tries < 20, "Unexpected balance state: {:?}", balances);
+		tries += 1;
+		tokio::time::sleep(delay).await;
+		if delay.as_millis() < 512 {
+			delay = delay.mul_f32(2.0);
+		}
+	}
+}
+
 proptest! {
 	#![proptest_config(proptest::test_runner::Config::with_cases(5))]
 	#[test]
@@ -143,50 +169,62 @@ proptest! {
 			generate_blocks_and_wait(bitcoind, electrs, 1).await;
 			sync_wallets!();
 
-			if force_close {
-				for node in &nodes {
-					node.sync_wallets().unwrap();
-					// If there is no more balance, there is nothing to process here.
-					if node.list_balances().lightning_balances.len() < 1 {
-						return;
-					}
-					match node.list_balances().lightning_balances[0] {
-						LightningBalance::ClaimableAwaitingConfirmations {
-							confirmation_height,
-							..
-						} => {
-							let cur_height = node.status().current_best_block.height;
-							let blocks_to_go = confirmation_height - cur_height;
+				if force_close {
+					let mut found_claimable_balance = false;
+					for node in &nodes {
+						node.sync_wallets().unwrap();
+						let balances = node.list_balances();
+						let confirmation_height = balances.lightning_balances.iter().find_map(|b| {
+							match b {
+								LightningBalance::ClaimableAwaitingConfirmations {
+									confirmation_height,
+									..
+								} => Some(*confirmation_height),
+								_ => None,
+							}
+						});
+						let Some(confirmation_height) = confirmation_height else {
+							continue;
+						};
+						found_claimable_balance = true;
+
+						let cur_height = node.status().current_best_block.height;
+						let blocks_to_go = confirmation_height.saturating_sub(cur_height);
+						if blocks_to_go > 0 {
 							generate_blocks_and_wait(bitcoind, electrs, blocks_to_go as usize).await;
 							node.sync_wallets().unwrap();
-						},
-						_ => panic!("Unexpected balance state for node_hub!"),
-					}
+						}
 
-					assert!(node.list_balances().lightning_balances.len() < 2);
-					assert!(node.list_balances().pending_balances_from_channel_closures.len() > 0);
-					match node.list_balances().pending_balances_from_channel_closures[0] {
-						PendingSweepBalance::BroadcastAwaitingConfirmation { .. } => {},
-						_ => panic!("Unexpected balance state!"),
-					}
+						wait_for_pending_sweep_balance(node, |balance| {
+							matches!(
+								balance,
+								PendingSweepBalance::BroadcastAwaitingConfirmation { .. }
+									| PendingSweepBalance::AwaitingThresholdConfirmations { .. }
+							)
+						})
+						.await;
 
-					generate_blocks_and_wait(&bitcoind, electrs, 1).await;
-					node.sync_wallets().unwrap();
-					assert!(node.list_balances().lightning_balances.len() < 2);
-					assert!(node.list_balances().pending_balances_from_channel_closures.len() > 0);
-					match node.list_balances().pending_balances_from_channel_closures[0] {
-						PendingSweepBalance::AwaitingThresholdConfirmations { .. } => {},
-						_ => panic!("Unexpected balance state!"),
+						generate_blocks_and_wait(bitcoind, electrs, 1).await;
+						node.sync_wallets().unwrap();
+						wait_for_pending_sweep_balance(node, |balance| {
+							matches!(balance, PendingSweepBalance::AwaitingThresholdConfirmations { .. })
+						})
+						.await;
 					}
+					assert!(found_claimable_balance);
 				}
-			}
 
 			generate_blocks_and_wait(bitcoind, electrs, 6).await;
 			sync_wallets!();
 
 			reorg!(reorg_depth);
 			sync_wallets!();
 
+			// The final reorg can leave close or sweep transactions below the wallet's
+			// trusted spendable depth even after they are confirmed on the replacement chain.
+			generate_blocks_and_wait(bitcoind, electrs, 6).await;
+			sync_wallets!();
+
 			let fee_sat = 7000;
 			// Check balance after close channel
 			nodes.iter().for_each(|node| {
@@ -198,6 +236,8 @@ proptest! {
 
 				assert_eq!(node.next_event(), None);
 			});
+
+			stop_nodes_concurrently(nodes).await;
 		})
 	}
 }