Support async results in TestChainSource, count get_utxo calls

Author: TheBlueMatt

lightning/src/routing/gossip.rs

@@ -1927,7 +1927,7 @@ mod tests {
 	#[cfg(feature = "std")]
 	use crate::ln::features::InitFeatures;
 	use crate::routing::gossip::{P2PGossipSync, NetworkGraph, NetworkUpdate, NodeAlias, MAX_EXCESS_BYTES_FOR_RELAY, NodeId, RoutingFees, ChannelUpdateInfo, ChannelInfo, NodeAnnouncementInfo, NodeInfo};
-	use crate::routing::utxo::UtxoLookupError;
+	use crate::routing::utxo::{UtxoLookupError, UtxoResult};
 	use crate::ln::msgs::{RoutingMessageHandler, UnsignedNodeAnnouncement, NodeAnnouncement,
 		UnsignedChannelAnnouncement, ChannelAnnouncement, UnsignedChannelUpdate, ChannelUpdate,
 		ReplyChannelRange, QueryChannelRange, QueryShortChannelIds, MAX_VALUE_MSAT};
@@ -2159,7 +2159,7 @@ mod tests {
 
 		// Test if an associated transaction were not on-chain (or not confirmed).
 		let chain_source = test_utils::TestChainSource::new(Network::Testnet);
-		*chain_source.utxo_ret.lock().unwrap() = Err(UtxoLookupError::UnknownTx);
+		*chain_source.utxo_ret.lock().unwrap() = UtxoResult::Sync(Err(UtxoLookupError::UnknownTx));
 		let network_graph = NetworkGraph::new(genesis_hash, &logger);
 		gossip_sync = P2PGossipSync::new(&network_graph, Some(&chain_source), &logger);
 
@@ -2172,7 +2172,8 @@ mod tests {
 		};
 
 		// Now test if the transaction is found in the UTXO set and the script is correct.
-		*chain_source.utxo_ret.lock().unwrap() = Ok(TxOut { value: 0, script_pubkey: good_script.clone() });
+		*chain_source.utxo_ret.lock().unwrap() =
+			UtxoResult::Sync(Ok(TxOut { value: 0, script_pubkey: good_script.clone() }));
 		let valid_announcement = get_signed_channel_announcement(|unsigned_announcement| {
 			unsigned_announcement.short_channel_id += 2;
 		}, node_1_privkey, node_2_privkey, &secp_ctx);
@@ -2190,7 +2191,8 @@ mod tests {
 
 		// If we receive announcement for the same channel, once we've validated it against the
 		// chain, we simply ignore all new (duplicate) announcements.
-		*chain_source.utxo_ret.lock().unwrap() = Ok(TxOut { value: 0, script_pubkey: good_script });
+		*chain_source.utxo_ret.lock().unwrap() =
+			UtxoResult::Sync(Ok(TxOut { value: 0, script_pubkey: good_script }));
 		match gossip_sync.handle_channel_announcement(&valid_announcement) {
 			Ok(_) => panic!(),
 			Err(e) => assert_eq!(e.err, "Already have chain-validated channel")
@@ -2264,7 +2266,8 @@ mod tests {
 		{
 			// Announce a channel we will update
 			let good_script = get_channel_script(&secp_ctx);
-			*chain_source.utxo_ret.lock().unwrap() = Ok(TxOut { value: amount_sats, script_pubkey: good_script.clone() });
+			*chain_source.utxo_ret.lock().unwrap() =
+				UtxoResult::Sync(Ok(TxOut { value: amount_sats, script_pubkey: good_script.clone() }));
 
 			let valid_channel_announcement = get_signed_channel_announcement(|_| {}, node_1_privkey, node_2_privkey, &secp_ctx);
 			short_channel_id = valid_channel_announcement.contents.short_channel_id;

lightning/src/routing/router.rs

@@ -2112,6 +2112,7 @@ fn build_route_from_hops_internal<L: Deref>(
 #[cfg(test)]
 mod tests {
 	use crate::routing::gossip::{NetworkGraph, P2PGossipSync, NodeId, EffectiveCapacity};
+	use crate::routing::utxo::UtxoResult;
 	use crate::routing::router::{get_route, build_route_from_hops_internal, add_random_cltv_offset, default_node_features,
 		PaymentParameters, Route, RouteHint, RouteHintHop, RouteHop, RoutingFees,
 		DEFAULT_MAX_TOTAL_CLTV_EXPIRY_DELTA, MAX_PATH_LENGTH_ESTIMATE};
@@ -3526,7 +3527,8 @@ mod tests {
 		.push_opcode(opcodes::all::OP_PUSHNUM_2)
 		.push_opcode(opcodes::all::OP_CHECKMULTISIG).into_script().to_v0_p2wsh();
 
-		*chain_monitor.utxo_ret.lock().unwrap() = Ok(TxOut { value: 15, script_pubkey: good_script.clone() });
+		*chain_monitor.utxo_ret.lock().unwrap() =
+			UtxoResult::Sync(Ok(TxOut { value: 15, script_pubkey: good_script.clone() }));
 		gossip_sync.add_utxo_lookup(Some(chain_monitor));
 
 		add_channel(&gossip_sync, &secp_ctx, &privkeys[0], &privkeys[2], ChannelFeatures::from_le_bytes(id_to_feature_flags(3)), 333);

lightning/src/routing/utxo.rs

@@ -42,6 +42,7 @@ pub enum UtxoLookupError {
 /// The result of a [`UtxoLookup::get_utxo`] call. A call may resolve either synchronously,
 /// returning the `Sync` variant, or asynchronously, returning an [`UtxoFuture`] in the `Async`
 /// variant.
+#[derive(Clone)]
 pub enum UtxoResult {
 	/// A result which was resolved synchronously. It either includes a [`TxOut`] for the output
 	/// requested or a [`UtxoLookupError`].

lightning/src/util/test_utils.rs

@@ -841,7 +841,8 @@ impl core::fmt::Debug for OnGetShutdownScriptpubkey {
 
 pub struct TestChainSource {
 	pub genesis_hash: BlockHash,
-	pub utxo_ret: Mutex<Result<TxOut, UtxoLookupError>>,
+	pub utxo_ret: Mutex<UtxoResult>,
+	pub get_utxo_call_count: AtomicUsize,
 	pub watched_txn: Mutex<HashSet<(Txid, Script)>>,
 	pub watched_outputs: Mutex<HashSet<(OutPoint, Script)>>,
 }
@@ -851,7 +852,8 @@ impl TestChainSource {
 		let script_pubkey = Builder::new().push_opcode(opcodes::OP_TRUE).into_script();
 		Self {
 			genesis_hash: genesis_block(network).block_hash(),
-			utxo_ret: Mutex::new(Ok(TxOut { value: u64::max_value(), script_pubkey })),
+			utxo_ret: Mutex::new(UtxoResult::Sync(Ok(TxOut { value: u64::max_value(), script_pubkey }))),
+			get_utxo_call_count: AtomicUsize::new(0),
 			watched_txn: Mutex::new(HashSet::new()),
 			watched_outputs: Mutex::new(HashSet::new()),
 		}
@@ -860,11 +862,12 @@ impl TestChainSource {
 
 impl UtxoLookup for TestChainSource {
 	fn get_utxo(&self, genesis_hash: &BlockHash, _short_channel_id: u64) -> UtxoResult {
+		self.get_utxo_call_count.fetch_add(1, Ordering::Relaxed);
 		if self.genesis_hash != *genesis_hash {
 			return UtxoResult::Sync(Err(UtxoLookupError::UnknownChain));
 		}
 
-		UtxoResult::Sync(self.utxo_ret.lock().unwrap().clone())
+		self.utxo_ret.lock().unwrap().clone()
 	}
 }