Proof of concept for BlockSource multiplexing

jkczyz · jkczyz · commit 6ab4b249f7de · 2020-12-14T16:26:28.000-08:00
diff --git a/lightning-block-sync/src/lib.rs b/lightning-block-sync/src/lib.rs
@@ -357,10 +357,10 @@ where P: Poll<'a, B>,
 	  CL: ChainListener + Sized {
 	chain_tip: (BlockHash, BlockHeaderData),
 	chain_poller: P,
-	backup_block_sources: Vec<B>,
 	header_cache: HeaderCache,
 	chain_notifier: CL,
-	mainnet: bool
+	mainnet: bool,
+	marker: std::marker::PhantomData<B>,
 }
 
 impl<'a, P, B, CL> MicroSPVClient<'a, P, B, CL>
@@ -380,11 +380,12 @@ where P: Poll<'a, B>,
 	/// useful when you have a block source which is more censorship-resistant than others but
 	/// which only provides headers. In this case, we can use such source(s) to learn of a censorship
 	/// attack without giving up privacy by querying a privacy-losing block sources.
-	pub fn init(chain_tip: BlockHeaderData, chain_poller: P, backup_block_sources: Vec<B>, chain_notifier: CL, mainnet: bool) -> Self {
+	pub fn init(chain_tip: BlockHeaderData, chain_poller: P, chain_notifier: CL, mainnet: bool) -> Self {
 		let header_cache = HeaderCache::new();
 		Self {
 			chain_tip: (chain_tip.header.block_hash(), chain_tip),
-			chain_poller, backup_block_sources, header_cache, chain_notifier, mainnet
+			chain_poller, header_cache, chain_notifier, mainnet,
+			marker: std::marker::PhantomData
 		}
 	}
 
@@ -420,15 +421,7 @@ where P: Poll<'a, B>,
 			Ok((ChainTip::Better(new_hash, new_header), block_source)) => {
 				debug_assert_ne!(new_hash, self.chain_tip.0);
 				debug_assert!(new_header.chainwork > self.chain_tip.1.chainwork);
-				let mut blocks_connected = false;
-				let backup_block_sources = self.backup_block_sources.iter_mut().map(|s| &mut **s);
-				for source in std::iter::once(block_source).chain(backup_block_sources) {
-					blocks_connected |= sync_chain_monitor!(new_hash, new_header, source);
-					if self.chain_tip.0 == new_hash {
-						break;
-					}
-				}
-				blocks_connected
+				sync_chain_monitor!(new_hash, new_header, block_source)
 			},
 			Ok((ChainTip::Worse(hash, header), _)) => {
 				debug_assert_ne!(hash, self.chain_tip.0);
@@ -786,8 +779,9 @@ mod tests {
 		let mut source_three = &header_chain;
 		let mut source_four = &backup_chain;
 		let mut client = MicroSPVClient::init((&chain_one).get_header(&block_1a_hash, Some(1)).await.unwrap(),
-			poller::MultipleChainPoller::new(vec![&mut source_one as &mut dyn BlockSource, &mut source_two as &mut dyn BlockSource, &mut source_three as &mut dyn BlockSource]),
-			vec![&mut source_four as &mut dyn BlockSource],
+			poller::ChainMultiplexer::new(
+				vec![&mut source_one as &mut dyn BlockSource, &mut source_two as &mut dyn BlockSource, &mut source_three as &mut dyn BlockSource],
+				vec![&mut source_four as &mut dyn BlockSource]),
 			Arc::clone(&chain_notifier), true);
 
 		// Test that we will reorg onto 2b because chain_one knows about 1b + 2b
diff --git a/lightning-block-sync/src/poller.rs b/lightning-block-sync/src/poller.rs
@@ -1,5 +1,8 @@
 use crate::{AsyncBlockSourceResult, BlockHeaderData, BlockSource, BlockSourceError, ChainTip, Poll};
 
+use bitcoin::blockdata::block::Block;
+use bitcoin::hash_types::BlockHash;
+
 use std::ops::DerefMut;
 
 pub struct ChainPoller<'a, B: DerefMut<Target=dyn BlockSource + 'a> + Sized + Sync + Send> {
@@ -13,7 +16,6 @@ impl<'a, B: DerefMut<Target=dyn BlockSource + 'a> + Sized + Sync + Send> ChainPo
 }
 
 impl<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> Poll<'b, B> for ChainPoller<'b, B> {
-
 	fn poll_chain_tip<'a>(&'a mut self, best_chain_tip: BlockHeaderData) ->
 		AsyncBlockSourceResult<'a, (ChainTip, &'a mut B::Target)>
 	where 'b: 'a {
@@ -59,7 +61,6 @@ impl<'a, B: DerefMut<Target=dyn BlockSource + 'a> + Sized + Sync + Send> Multipl
 }
 
 impl<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> Poll<'b, B> for MultipleChainPoller<'b, B> {
-
 	fn poll_chain_tip<'a>(&'a mut self, best_chain_tip: BlockHeaderData) ->
 		AsyncBlockSourceResult<'a, (ChainTip, &'a mut B::Target)>
 	where 'b: 'a {
@@ -102,6 +103,144 @@ impl<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> Poll<'b
 	}
 }
 
+pub struct ChainMultiplexer<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> {
+	block_sources: Vec<(B, BlockSourceError)>,
+	backup_block_sources: Vec<(B, BlockSourceError)>,
+	best_block_source: usize,
+}
+
+impl<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> ChainMultiplexer<'b, B> {
+	pub fn new(mut block_sources: Vec<B>, mut backup_block_sources: Vec<B>) -> Self {
+		assert!(!block_sources.is_empty());
+		let block_sources = block_sources.drain(..).map(|block_source| {
+			(block_source, BlockSourceError::Transient)
+		}).collect();
+
+		let backup_block_sources = backup_block_sources.drain(..).map(|block_source| {
+			(block_source, BlockSourceError::Transient)
+		}).collect();
+
+		Self { block_sources, backup_block_sources, best_block_source: 0 }
+	}
+
+	fn best_and_backup_block_sources(&mut self) -> Vec<&mut (B, BlockSourceError)> {
+		let best_block_source = self.block_sources.get_mut(self.best_block_source).unwrap();
+		let backup_block_sources = self.backup_block_sources.iter_mut();
+		std::iter::once(best_block_source)
+			.chain(backup_block_sources)
+			.filter(|(_, e)| e == &BlockSourceError::Transient)
+			.collect()
+	}
+}
+
+impl<'b, B: 'b + DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> Poll<'b, B> for ChainMultiplexer<'b, B> {
+	fn poll_chain_tip<'a>(&'a mut self, best_chain_tip: BlockHeaderData) ->
+		AsyncBlockSourceResult<'a, (ChainTip, &'a mut B::Target)>
+	where 'b: 'a {
+		Box::pin(async move {
+			let mut heaviest_chain_tip = best_chain_tip;
+			let mut best_result = Err(BlockSourceError::Persistent);
+			for (i, (block_source, error)) in self.block_sources.iter_mut().enumerate() {
+				if let BlockSourceError::Persistent = error {
+					continue;
+				}
+
+				let result = match block_source.get_best_block().await {
+					Err(e) => Err(e),
+					Ok((block_hash, height)) => {
+						if block_hash == heaviest_chain_tip.header.block_hash() {
+							Ok(ChainTip::Common)
+						} else {
+							match block_source.get_header(&block_hash, height).await {
+								Err(e) => Err(e),
+								Ok(chain_tip) => {
+									crate::stateless_check_header(&chain_tip.header)?;
+									if chain_tip.header.block_hash() != block_hash {
+										Err(BlockSourceError::Persistent)
+									} else if chain_tip.chainwork <= heaviest_chain_tip.chainwork {
+										Ok(ChainTip::Worse(block_hash, chain_tip))
+									} else {
+										Ok(ChainTip::Better(block_hash, chain_tip))
+									}
+								},
+							}
+						}
+					},
+				};
+
+				match result {
+					Err(BlockSourceError::Persistent) => {
+						*error = BlockSourceError::Persistent;
+					},
+					Err(BlockSourceError::Transient) => {
+						if best_result.is_err() {
+							best_result = result;
+						}
+					},
+					Ok(ChainTip::Common) => {
+						if let Ok(ChainTip::Better(_, _)) = best_result {} else {
+							best_result = result;
+						}
+					},
+					Ok(ChainTip::Better(_, header)) => {
+						self.best_block_source = i;
+						best_result = result;
+						heaviest_chain_tip = header;
+					},
+					Ok(ChainTip::Worse(_, _)) => {
+						if best_result.is_err() {
+							best_result = result;
+						}
+					},
+				}
+			}
+
+			best_result.map(move |chain_tip| (chain_tip, self as &mut dyn BlockSource))
+		})
+	}
+}
+
+impl<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> BlockSource for ChainMultiplexer<'b, B> {
+	fn get_header<'a>(&'a mut self, header_hash: &'a BlockHash, height: Option<u32>) -> AsyncBlockSourceResult<'a, BlockHeaderData> {
+		Box::pin(async move {
+			for (block_source, error) in self.best_and_backup_block_sources() {
+				let result = block_source.get_header(header_hash, height).await;
+				match result {
+					Err(e) => *error = e,
+					Ok(_) => return result,
+				}
+			}
+			Err(BlockSourceError::Persistent)
+		})
+	}
+
+	fn get_block<'a>(&'a mut self, header_hash: &'a BlockHash) -> AsyncBlockSourceResult<'a, Block> {
+		Box::pin(async move {
+			for (block_source, error) in self.best_and_backup_block_sources() {
+				let result = block_source.get_block(header_hash).await;
+				match result {
+					Err(e) => *error = e,
+					Ok(_) => return result,
+				}
+			}
+			Err(BlockSourceError::Persistent)
+		})
+	}
+
+	fn get_best_block<'a>(&'a mut self) -> AsyncBlockSourceResult<'a, (BlockHash, Option<u32>)> {
+		Box::pin(async move {
+			for (block_source, error) in self.best_and_backup_block_sources() {
+				let result = block_source.get_best_block().await;
+				match result {
+					Err(e) => *error = e,
+					Ok(_) => return result,
+				}
+			}
+			Err(BlockSourceError::Persistent)
+		})
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use crate::*;
