WIP: New interface for syncing multiple listeners

Author: jkczyz

lightning-block-sync/src/init.rs

@@ -0,0 +1,243 @@
+use crate::{BlockSource, BlockSourceResult, Cache, ChainListener, ChainNotifier};
+use crate::poll::{ChainPoller, Validate, ValidatedBlockHeader};
+
+use bitcoin::blockdata::block::{Block, BlockHeader};
+use bitcoin::hash_types::BlockHash;
+use bitcoin::network::constants::Network;
+
+/// Unbounded cache of block headers keyed by block hash.
+pub type UnboundedCache = std::collections::HashMap<BlockHash, ValidatedBlockHeader>;
+
+impl Cache for UnboundedCache {
+	fn look_up(&self, block_hash: &BlockHash) -> Option<&ValidatedBlockHeader> {
+		self.get(block_hash)
+	}
+
+	fn block_connected(&mut self, block_hash: BlockHash, block_header: ValidatedBlockHeader) {
+		self.insert(block_hash, block_header);
+	}
+
+	fn block_disconnected(&mut self, _block_hash: &BlockHash) -> Option<ValidatedBlockHeader> {
+		None
+	}
+}
+
+/// Performs a one-time sync of chain listeners from a single *trusted* block source, bringing their
+/// view of the latest chain tip to `new_block`.
+///
+/// Useful upon startup to bring each `ChannelMonitor` and the `ChannelManager` in sync, before
+/// switching to `SpvClient`.
+pub async fn sync_listeners<B: BlockSource, C: Cache>(
+	new_block: BlockHash,
+	block_source: &mut B,
+	network: Network,
+	chain_notifier: &mut ChainNotifier<C>,
+	mut chain_listeners: Vec<(BlockHash, &mut dyn ChainListener)>,
+) -> BlockSourceResult<()> {
+	let new_header = match chain_notifier.header_cache.look_up(&new_block) {
+		Some(header) => *header,
+		None => block_source
+			.get_header(&new_block, None).await?
+			.validate(new_block)?,
+	};
+
+	// Find differences and disconnect blocks for each listener individually.
+	let mut chain_listeners_with_common_ancestors = Vec::new();
+	let mut most_common_ancestor = None;
+	let mut most_connected_blocks = Vec::new();
+	for (old_block, chain_listener) in chain_listeners.drain(..) {
+		let old_header = match chain_notifier.header_cache.look_up(&old_block) {
+			Some(header) => *header,
+			None => block_source
+				.get_header(&old_block, None).await?
+				.validate(old_block)?
+		};
+
+		// Disconnect any stale blocks.
+		let mut chain_poller = ChainPoller::new(block_source as &mut dyn BlockSource, network);
+		let difference =
+			chain_notifier.find_difference(new_header, &old_header, &mut chain_poller).await?;
+		chain_notifier.disconnect_blocks(
+			difference.disconnected_blocks,
+			&mut DynamicChainListener(chain_listener),
+		);
+
+		// Keep track of the most common ancestor and all blocks connected across all listeners.
+		chain_listeners_with_common_ancestors.push((difference.common_ancestor, chain_listener));
+		if difference.connected_blocks.len() > most_connected_blocks.len() {
+			most_common_ancestor = Some(difference.common_ancestor);
+			most_connected_blocks = difference.connected_blocks;
+		}
+	}
+
+	// Connect new blocks for all listeners at once to avoid re-fetching blocks.
+	if let Some(common_ancestor) = most_common_ancestor {
+		let mut chain_poller = ChainPoller::new(block_source as &mut dyn BlockSource, network);
+		let mut chain_listener = ChainListenerSet(chain_listeners_with_common_ancestors);
+		chain_notifier.connect_blocks(
+			common_ancestor,
+			most_connected_blocks,
+			&mut chain_poller,
+			&mut chain_listener,
+		).await.unwrap();
+	}
+	Ok(())
+}
+
+struct DynamicChainListener<'a>(&'a mut dyn ChainListener);
+
+impl<'a> ChainListener for DynamicChainListener<'a> {
+	fn block_connected(&mut self, _block: &Block, _height: u32) {
+		unreachable!()
+	}
+
+	fn block_disconnected(&mut self, header: &BlockHeader, height: u32) {
+		self.0.block_disconnected(header, height)
+	}
+}
+
+struct ChainListenerSet<'a>(Vec<(ValidatedBlockHeader, &'a mut dyn ChainListener)>);
+
+impl<'a> ChainListener for ChainListenerSet<'a> {
+	fn block_connected(&mut self, block: &Block, height: u32) {
+		for (header, chain_listener) in self.0.iter_mut() {
+			if height > header.height {
+				chain_listener.block_connected(block, height);
+			}
+		}
+	}
+
+	fn block_disconnected(&mut self, _header: &BlockHeader, _height: u32) {
+		unreachable!()
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use crate::test_utils::{Blockchain, MockChainListener};
+	use super::*;
+
+	use bitcoin::network::constants::Network;
+
+	#[tokio::test]
+	async fn sync_from_same_chain() {
+		let mut chain = Blockchain::default().with_height(4);
+		let new_tip = chain.tip();
+		let old_tip_1 = chain.at_height(1);
+		let old_tip_2 = chain.at_height(2);
+		let old_tip_3 = chain.at_height(3);
+
+		let mut listener_1 = MockChainListener::new()
+			.expect_block_connected(*old_tip_2)
+			.expect_block_connected(*old_tip_3)
+			.expect_block_connected(*new_tip);
+		let mut listener_2 = MockChainListener::new()
+			.expect_block_connected(*old_tip_3)
+			.expect_block_connected(*new_tip);
+		let mut listener_3 = MockChainListener::new()
+			.expect_block_connected(*new_tip);
+
+		let listeners = vec![
+			(old_tip_1.block_hash, &mut listener_1 as &mut dyn ChainListener),
+			(old_tip_2.block_hash, &mut listener_2 as &mut dyn ChainListener),
+			(old_tip_3.block_hash, &mut listener_3 as &mut dyn ChainListener),
+		];
+		let mut notifier = ChainNotifier { header_cache: chain.header_cache(0..=4) };
+		match sync_listeners(new_tip.block_hash, &mut chain, Network::Bitcoin, &mut notifier, listeners).await {
+			Ok(()) => {},
+			Err(e) => panic!("Unexpected error: {:?}", e),
+		}
+	}
+
+	#[tokio::test]
+	async fn sync_from_different_chains() {
+		let mut main_chain = Blockchain::default().with_height(4);
+		let fork_chain_1 = main_chain.fork_at_height(1);
+		let fork_chain_2 = main_chain.fork_at_height(2);
+		let fork_chain_3 = main_chain.fork_at_height(3);
+
+		let new_tip = main_chain.tip();
+		let old_tip_1 = fork_chain_1.tip();
+		let old_tip_2 = fork_chain_2.tip();
+		let old_tip_3 = fork_chain_3.tip();
+
+		let mut listener_1 = MockChainListener::new()
+			.expect_block_disconnected(*fork_chain_1.at_height(4))
+			.expect_block_disconnected(*fork_chain_1.at_height(3))
+			.expect_block_disconnected(*fork_chain_1.at_height(2))
+			.expect_block_connected(*main_chain.at_height(2))
+			.expect_block_connected(*main_chain.at_height(3))
+			.expect_block_connected(*main_chain.at_height(4));
+		let mut listener_2 = MockChainListener::new()
+			.expect_block_disconnected(*fork_chain_2.at_height(4))
+			.expect_block_disconnected(*fork_chain_2.at_height(3))
+			.expect_block_connected(*main_chain.at_height(3))
+			.expect_block_connected(*main_chain.at_height(4));
+		let mut listener_3 = MockChainListener::new()
+			.expect_block_disconnected(*fork_chain_3.at_height(4))
+			.expect_block_connected(*main_chain.at_height(4));
+
+		let listeners = vec![
+			(old_tip_1.block_hash, &mut listener_1 as &mut dyn ChainListener),
+			(old_tip_2.block_hash, &mut listener_2 as &mut dyn ChainListener),
+			(old_tip_3.block_hash, &mut listener_3 as &mut dyn ChainListener),
+		];
+		let mut header_cache = fork_chain_1.header_cache(2..=4);
+		header_cache.extend(fork_chain_2.header_cache(3..=4));
+		header_cache.extend(fork_chain_3.header_cache(4..=4));
+		let mut notifier = ChainNotifier { header_cache };
+		match sync_listeners(new_tip.block_hash, &mut main_chain, Network::Bitcoin, &mut notifier, listeners).await {
+			Ok(()) => {},
+			Err(e) => panic!("Unexpected error: {:?}", e),
+		}
+	}
+
+	#[tokio::test]
+	async fn sync_from_overlapping_chains() {
+		let mut main_chain = Blockchain::default().with_height(4);
+		let fork_chain_1 = main_chain.fork_at_height(1);
+		let fork_chain_2 = fork_chain_1.fork_at_height(2);
+		let fork_chain_3 = fork_chain_2.fork_at_height(3);
+
+		let new_tip = main_chain.tip();
+		let old_tip_1 = fork_chain_1.tip();
+		let old_tip_2 = fork_chain_2.tip();
+		let old_tip_3 = fork_chain_3.tip();
+
+		let mut listener_1 = MockChainListener::new()
+			.expect_block_disconnected(*fork_chain_1.at_height(4))
+			.expect_block_disconnected(*fork_chain_1.at_height(3))
+			.expect_block_disconnected(*fork_chain_1.at_height(2))
+			.expect_block_connected(*main_chain.at_height(2))
+			.expect_block_connected(*main_chain.at_height(3))
+			.expect_block_connected(*main_chain.at_height(4));
+		let mut listener_2 = MockChainListener::new()
+			.expect_block_disconnected(*fork_chain_2.at_height(4))
+			.expect_block_disconnected(*fork_chain_2.at_height(3))
+			.expect_block_disconnected(*fork_chain_2.at_height(2))
+			.expect_block_connected(*main_chain.at_height(2))
+			.expect_block_connected(*main_chain.at_height(3))
+			.expect_block_connected(*main_chain.at_height(4));
+		let mut listener_3 = MockChainListener::new()
+			.expect_block_disconnected(*fork_chain_3.at_height(4))
+			.expect_block_disconnected(*fork_chain_3.at_height(3))
+			.expect_block_disconnected(*fork_chain_3.at_height(2))
+			.expect_block_connected(*main_chain.at_height(2))
+			.expect_block_connected(*main_chain.at_height(3))
+			.expect_block_connected(*main_chain.at_height(4));
+
+		let listeners = vec![
+			(old_tip_1.block_hash, &mut listener_1 as &mut dyn ChainListener),
+			(old_tip_2.block_hash, &mut listener_2 as &mut dyn ChainListener),
+			(old_tip_3.block_hash, &mut listener_3 as &mut dyn ChainListener),
+		];
+		let mut header_cache = fork_chain_1.header_cache(2..=4);
+		header_cache.extend(fork_chain_2.header_cache(3..=4));
+		header_cache.extend(fork_chain_3.header_cache(4..=4));
+		let mut notifier = ChainNotifier { header_cache };
+		match sync_listeners(new_tip.block_hash, &mut main_chain, Network::Bitcoin, &mut notifier, listeners).await {
+			Ok(()) => {},
+			Err(e) => panic!("Unexpected error: {:?}", e),
+		}
+	}
+}