Use MockInstant for tests

Author: dunxen

lightning/src/routing/gossip.rs

@@ -45,6 +45,10 @@ use bitcoin::hashes::hex::ToHex;
 
 #[cfg(feature = "std")]
 use std::time::{SystemTime, UNIX_EPOCH};
+#[cfg(all(not(test), feature = "std"))]
+use std::time::Instant;
+#[cfg(all(test, feature = "std"))]
+use util::time::tests::MockInstant as Instant;
 
 /// We remove stale channel directional info two weeks after the last update, per BOLT 7's
 /// suggestion.
@@ -130,20 +134,18 @@ pub struct NetworkGraph<L: Deref> where L::Target: Logger {
 	// Lock order: channels -> nodes
 	channels: RwLock<BTreeMap<u64, ChannelInfo>>,
 	nodes: RwLock<BTreeMap<NodeId, NodeInfo>>,
-	/// Keeps track of [`NodeId`]s for we have explicitly
-	/// removed due to permanent failure so that we don't resync them
-	/// from gossip. Each [`NodeId`] is mapped to the UNIX timestamp of when
-	/// we removed it so that we can forget we removed it after some
-	/// time (by removing it from this map) and it can be resynced from
-	/// gossip if it appears again.
-	removed_nodes: Mutex<HashMap<NodeId, u64>>,
-	/// Keeps track of short channel IDs for channels we have explicitly
-	/// removed due to permanent failure so that we don't resync them
-	/// from gossip. Each ID is mapped to the UNIX timestamp of when
-	/// we removed it so that we can forget we removed them after some
-	/// time (by removing them from this map) and they can be resynced from
-	/// gossip if they appear again.
-	removed_channels: Mutex<HashMap<u64, u64>>,
+	#[cfg(feature = "std")]
+	/// Keeps track of [`NodeId`]s we have explicitly removed due to permanent failure so that we
+	/// don't resync them from gossip. Each [`NodeId`] is mapped to the isntant when we removed it
+	/// from the network graph so that we can forget we removed it after some time (by removing it
+	/// from this map) and it can be resynced from gossip if it appears again.
+	removed_nodes: Mutex<HashMap<NodeId, Instant>>,
+	#[cfg(feature = "std")]
+	/// Keeps track of short channel IDs for channels we have explicitly removed due to permanent
+	/// failure so that we don't resync them from gossip. Each SCID is mapped to the instant when we
+	/// removed it from the network graph so that we can forget we removed it after some time (by
+	/// removing it from this map) and it can be resynced from gossip if it appears again.
+	removed_channels: Mutex<HashMap<u64, Instant>>,
 }
 
 /// A read-only view of [`NetworkGraph`].
@@ -1206,7 +1208,9 @@ impl<L: Deref> ReadableArgs<L> for NetworkGraph<L> where L::Target: Logger {
 			channels: RwLock::new(channels),
 			nodes: RwLock::new(nodes),
 			last_rapid_gossip_sync_timestamp: Mutex::new(last_rapid_gossip_sync_timestamp),
+			#[cfg(feature = "std")]
 			removed_nodes: Mutex::new(HashMap::new()),
+			#[cfg(feature = "std")]
 			removed_channels: Mutex::new(HashMap::new()),
 		})
 	}
@@ -1244,7 +1248,9 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 			channels: RwLock::new(BTreeMap::new()),
 			nodes: RwLock::new(BTreeMap::new()),
 			last_rapid_gossip_sync_timestamp: Mutex::new(None),
+			#[cfg(feature = "std")]
 			removed_nodes: Mutex::new(HashMap::new()),
+			#[cfg(feature = "std")]
 			removed_channels: Mutex::new(HashMap::new()),
 		}
 	}
@@ -1301,6 +1307,7 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 	}
 
 	fn update_node_from_announcement_intern(&self, msg: &msgs::UnsignedNodeAnnouncement, full_msg: Option<&msgs::NodeAnnouncement>) -> Result<(), LightningError> {
+		#[cfg(feature = "std")]
 		{
 			if self.removed_nodes.lock().unwrap().contains_key(&NodeId::from_pubkey(&msg.node_id)) {
 				return Err(LightningError{err: "Node was removed from our network graph recently".to_owned(), action: ErrorAction::IgnoreAndLog(Level::Gossip)});
@@ -1462,6 +1469,7 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 			return Err(LightningError{err: "Channel announcement node had a channel with itself".to_owned(), action: ErrorAction::IgnoreError});
 		}
 
+		#[cfg(feature = "std")]
 		{
 			let channels = self.channels.read().unwrap();
 
@@ -1495,7 +1503,8 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 			}
 		}
 
-		{
+		#[cfg(feature = "std")]
+        {
 			let removed_nodes = self.removed_nodes.lock().unwrap();
 			let removed_channels = self.removed_channels.lock().unwrap();
 			if removed_channels.contains_key(&msg.short_channel_id) ||
@@ -1586,7 +1595,7 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 		let mut nodes = self.nodes.write().unwrap();
 		let mut removed_nodes = self.removed_nodes.lock().unwrap();
 		let mut removed_channels = self.removed_channels.lock().unwrap();
-		let time = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time must be > 1970").as_secs();
+		let time = Instant::now();
 
 		if let Some(node) = nodes.remove(&node_id) {
 			for scid in node.channels.iter() {
@@ -1625,10 +1634,12 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 	/// can be resynced from gossip in the future.
 	///
 	/// This method is only available with the `std` feature. See
-	/// [`NetworkGraph::remove_stale_channels_and_tracking_with_time`] for `no-std` use.
+	/// [`NetworkGraph::remove_stale_channels_with_time`] for `no-std` use.
 	pub fn remove_stale_channels_and_tracking(&self) {
 		let time = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time must be > 1970").as_secs();
-		self.remove_stale_channels_and_tracking_with_time(time);
+		self.remove_stale_channels_with_time(time);
+		self.removed_nodes.lock().unwrap().retain(|_, time| time.elapsed().as_secs() < STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
+		self.removed_channels.lock().unwrap().retain(|_, time| time.elapsed().as_secs() < STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
 	}
 
 	/// Removes information about channels that we haven't heard any updates about in some time.
@@ -1640,12 +1651,8 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 	/// pruning occur for updates which are at least two weeks old, which we implement here.
 	///
 	/// This function takes the current unix time as an argument. For users with the `std` feature
-	/// enabled, [`NetworkGraph::remove_stale_channels`] may be preferable.
-	///
-	/// When the `std` feature is enabled then this method will also remove entries from
-	/// [`NetworkGraph::removed_nodes`] and [`NetworkGraph::removed_channels`] if they have been in
-	/// the map for a while so that these can be resynced from gossip in the future.
-	pub fn remove_stale_channels_and_tracking_with_time(&self, current_time_unix: u64) {
+	/// enabled, [`NetworkGraph::remove_stale_channels_and_tracking`] may be preferable.
+	pub fn remove_stale_channels_with_time(&self, current_time_unix: u64) {
 		let mut channels = self.channels.write().unwrap();
 		// Time out if we haven't received an update in at least 14 days.
 		if current_time_unix > u32::max_value() as u64 { return; } // Remove by 2106
@@ -1677,14 +1684,6 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 				self.remove_channel_in_nodes(&mut nodes, &info, scid);
 			}
 		}
-
-		// We only bother to remove nodes and channels from tracking in std as they're only ever
-		// added to tracking with std.
-		#[cfg(feature = "std")]
-		{
-			self.removed_nodes.lock().unwrap().retain(|_, timestamp| current_time_unix.saturating_sub(*timestamp) < STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
-			self.removed_channels.lock().unwrap().retain(|_, timestamp| current_time_unix.saturating_sub(*timestamp) < STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
-		}
 	}
 
 	/// For an already known (from announcement) channel, update info about one of the directions
@@ -1861,7 +1860,7 @@ impl<L: Deref> NetworkGraph<L> where L::Target: Logger {
 			($node_id: expr) => {
 				if let BtreeEntry::Occupied(mut entry) = nodes.entry($node_id) {
 					#[cfg(feature = "std")]
-					let time = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time must be > 1970").as_secs();
+					let time = Instant::now();
 					#[cfg(feature = "std")]
 					self.removed_channels.lock().unwrap().insert(short_channel_id, time);
 
@@ -2206,24 +2205,28 @@ mod tests {
 			Err(e) => assert_eq!(e.err, "Already have chain-validated channel")
 		};
 
-		// Insert node 1 into removed_nodes map to track it as removed.
-		gossip_sync.network_graph().removed_nodes.lock().unwrap().insert(NodeId::from_pubkey(
-			&PublicKey::from_secret_key(&secp_ctx, node_1_privkey)), 0);	
+		#[cfg(feature = "std")]
+		{
+			use util::time::tests::MockInstant as Instant;
+			// Insert node 1 into removed_nodes map to track it as removed.
+			gossip_sync.network_graph().removed_nodes.lock().unwrap().insert(NodeId::from_pubkey(
+				&PublicKey::from_secret_key(&secp_ctx, node_1_privkey)), Instant::now());		
 
-		// Return error and ignore valid channel announcement if one of the nodes has been tracked as removed.
-		let valid_announcement = get_signed_channel_announcement(|unsigned_announcement| {
-			unsigned_announcement.short_channel_id += 3;
-		}, node_1_privkey, node_2_privkey, &secp_ctx);
-		match gossip_sync.handle_channel_announcement(&valid_announcement) {
-			Ok(_) => panic!(),
-			Err(e) => assert_eq!(e.err, "Channel or one of its nodes was removed from our network graph recently")
-		}
+			// Return error and ignore valid channel announcement if one of the nodes has been tracked as removed.
+			let valid_announcement = get_signed_channel_announcement(|unsigned_announcement| {
+				unsigned_announcement.short_channel_id += 3;
+			}, node_1_privkey, node_2_privkey, &secp_ctx);
+			match gossip_sync.handle_channel_announcement(&valid_announcement) {
+				Ok(_) => panic!(),
+				Err(e) => assert_eq!(e.err, "Channel or one of its nodes was removed from our network graph recently")
+			}
 
-		// But handle above channel announcement if the node is not being tracked as removed anymore.
-		gossip_sync.network_graph.removed_nodes.lock().unwrap().clear();
-		match gossip_sync.handle_channel_announcement(&valid_announcement) {
-			Ok(res) => assert!(res),
-			_ => panic!()
+			// But handle above channel announcement if the node is not being tracked as removed anymore.
+			gossip_sync.network_graph.removed_nodes.lock().unwrap().clear();
+			match gossip_sync.handle_channel_announcement(&valid_announcement) {
+				Ok(res) => assert!(res),
+				_ => panic!()
+			}
 		}
 
 		// Don't relay valid channels with excess data
@@ -2533,11 +2536,11 @@ mod tests {
 		assert!(gossip_sync.handle_channel_update(&valid_channel_update).is_ok());
 		assert!(network_graph.read_only().channels().get(&short_channel_id).unwrap().one_to_two.is_some());
 
-		network_graph.remove_stale_channels_and_tracking_with_time(100 + STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
+		network_graph.remove_stale_channels_with_time(100 + STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
 		assert_eq!(network_graph.read_only().channels().len(), 1);
 		assert_eq!(network_graph.read_only().nodes().len(), 2);
 
-		network_graph.remove_stale_channels_and_tracking_with_time(101 + STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
+		network_graph.remove_stale_channels_with_time(101 + STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
 		#[cfg(feature = "std")]
 		{
 			// In std mode, a further check is performed before fully removing the channel -
@@ -2551,7 +2554,7 @@ mod tests {
 			use std::time::{SystemTime, UNIX_EPOCH};
 			let announcement_time = SystemTime::now().duration_since(UNIX_EPOCH)
 				.expect("Time must be > 1970").as_secs();
-			network_graph.remove_stale_channels_and_tracking_with_time(
+			network_graph.remove_stale_channels_with_time(
 				announcement_time + 1 + STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
 		}
 
@@ -2560,28 +2563,29 @@ mod tests {
 
 		#[cfg(feature = "std")]
 		{
-			use std::time::{SystemTime, UNIX_EPOCH};
+			use util::time::tests::MockInstant as Instant;
+			use std::time::Duration;
 
 			// Clear tracked nodes and channels for clean slate
 			network_graph.removed_nodes.lock().unwrap().clear();
 			network_graph.removed_channels.lock().unwrap().clear();
 
 			// Track some nodes and channels as removed
-			let started_tracking_time = SystemTime::now().duration_since(UNIX_EPOCH)
-				.expect("Time must be > 1970").as_secs();
+			let started_tracking = Instant::now();
 			network_graph.removed_nodes.lock().unwrap().insert(
 				NodeId::from_pubkey(&PublicKey::from_secret_key(&secp_ctx, node_1_privkey)),
-				started_tracking_time);
-			network_graph.removed_channels.lock().unwrap().insert(1, started_tracking_time);
+				started_tracking);
+			network_graph.removed_channels.lock().unwrap().insert(1, started_tracking);
 
 			// Should not remove from tracking if insufficient time has passed
-			network_graph.remove_stale_channels_and_tracking_with_time(started_tracking_time + 1);
+			Instant::advance(Duration::from_secs(STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS - 1));
+			network_graph.remove_stale_channels_and_tracking();
 			assert_eq!(network_graph.removed_nodes.lock().unwrap().len(), 1);
 			assert_eq!(network_graph.removed_channels.lock().unwrap().len(), 1);
 
-			// Node and channel entry should be removed on next call
-			network_graph.remove_stale_channels_and_tracking_with_time(
-				started_tracking_time + STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS);
+			// Advance mock time so the entries are removed on next call
+			Instant::advance(Duration::from_secs(1)); 
+			network_graph.remove_stale_channels_and_tracking();
 			assert!(network_graph.removed_nodes.lock().unwrap().is_empty());
 			assert!(network_graph.removed_channels.lock().unwrap().is_empty());
 		}