f simplify disconnection logic a good bit with more utils

Author: TheBlueMatt

lightning/src/ln/channelmanager.rs

@@ -4263,28 +4263,35 @@ where
 		Ok(())
 	}
 
+	fn unfunded_channel_count(
+		peer: &PeerState<<SP::Target as SignerProvider>::Signer>, best_block_height: u32
+	) -> usize {
+		let mut num_unfunded_channels = 0;
+		for (_, chan) in peer.channel_by_id.iter() {
+			if !chan.is_outbound() && chan.minimum_depth().unwrap_or(1) != 0 &&
+				chan.get_funding_tx_confirmations(best_block_height) == 0
+			{
+				num_unfunded_channels += 1;
+			}
+		}
+		num_unfunded_channels
+	}
+
 	/// Gets the number of peers which match the given filter which do not have any funded,
 	/// outbound, or 0-conf channels.
 	///
 	/// The filter is called for each peer and provided with the number of unfunded, inbound, and
 	/// non-0-conf channels we have with the peer.
 	fn peers_without_funded_channels<Filter>(&self, mut filter: Filter) -> usize
-	where Filter: FnMut(&PublicKey, &PeerState<<SP::Target as SignerProvider>::Signer>, usize) -> bool {
+	where Filter: FnMut(&PeerState<<SP::Target as SignerProvider>::Signer>) -> bool {
 		let mut peers_without_funded_channels = 0;
 		let best_block_height = self.best_block.read().unwrap().height();
 		{
 			let peer_state_lock = self.per_peer_state.read().unwrap();
-			for (node_id, peer_mtx) in peer_state_lock.iter() {
+			for (_, peer_mtx) in peer_state_lock.iter() {
 				let peer = peer_mtx.lock().unwrap();
-				let mut num_unfunded_channels = 0;
-				for (_, chan) in peer.channel_by_id.iter() {
-					if !chan.is_outbound() && chan.minimum_depth().unwrap_or(1) != 0 &&
-						chan.get_funding_tx_confirmations(best_block_height) == 0
-					{
-						num_unfunded_channels += 1;
-					}
-				}
-				if !filter(node_id, &*peer, num_unfunded_channels) { continue; }
+				if !filter(&*peer) { continue; }
+				let num_unfunded_channels = Self::unfunded_channel_count(&peer, best_block_height);
 				if num_unfunded_channels == peer.channel_by_id.len() {
 					peers_without_funded_channels += 1;
 				}
@@ -4316,21 +4323,16 @@ where
 		// Get the number of peers with channels, but without funded ones. We don't care too much
 		// about peers that never open a channel, so we filter by peers that have at least one
 		// channel, and then limit the number of those with unfunded channels.
-		let mut this_node_unfunded_channels = 0;
-		let peers_without_funded_channels = self.peers_without_funded_channels(
-			|node_id, node, unfunded_channels| {
-				if node_id == counterparty_node_id {
-					this_node_unfunded_channels = unfunded_channels;
-				}
-				!node.channel_by_id.is_empty()
-			});
-		if this_node_unfunded_channels >= MAX_UNFUNDED_CHANS_PER_PEER {
+		let peers_without_funded_channels = self.peers_without_funded_channels(|node| !node.channel_by_id.is_empty());
+		let mut peer_state_lock = peer_state_mutex_opt.unwrap().lock().unwrap();
+		let peer_state = &mut *peer_state_lock;
+
+		let best_block_height = self.best_block.read().unwrap().height();
+		if Self::unfunded_channel_count(peer_state, best_block_height) >= MAX_UNFUNDED_CHANS_PER_PEER {
 			return Err(MsgHandleErrInternal::send_err_msg_no_close(
 				format!("Refusing more than {} unfunded channels.", MAX_UNFUNDED_CHANS_PER_PEER),
 				msg.temporary_channel_id.clone()));
 		}
-		let mut peer_state_lock = peer_state_mutex_opt.unwrap().lock().unwrap();
-		let peer_state = &mut *peer_state_lock;
 		// If this peer already has some channels, a new channel won't increase our number of peers
 		// with unfunded channels, so as long as we aren't over the maximum number of unfunded
 		// channels per-peer we can accept channels from a peer with existing ones.
@@ -4342,7 +4344,7 @@ where
 
 		let mut channel = match Channel::new_from_req(&self.fee_estimator, &self.entropy_source, &self.signer_provider,
 			counterparty_node_id.clone(), &self.channel_type_features(), &peer_state.latest_features, msg, user_channel_id,
-			&self.default_configuration, self.best_block.read().unwrap().height(), &self.logger, outbound_scid_alias)
+			&self.default_configuration, best_block_height, &self.logger, outbound_scid_alias)
 		{
 			Err(e) => {
 				self.outbound_scid_aliases.lock().unwrap().remove(&outbound_scid_alias);
@@ -6296,25 +6298,19 @@ where
 		let _persistence_guard = PersistenceNotifierGuard::notify_on_drop(&self.total_consistency_lock, &self.persistence_notifier);
 
 		// If we have too many peers connected which don't have funded channels, disconnect the
-		// peer immediately. If we have a bunch of unfunded channels taking up space in memory for
-		// disconnected peers, we still let new peers connect, but we'll reject new channels from
-		// them.
-		let mut this_peer_has_funded_channels = false;
-		let connected_peers_without_funded_channels = self.peers_without_funded_channels(
-			|node_id, node, num_unfunded_channels| {
-				if node_id == counterparty_node_id && num_unfunded_channels != node.channel_by_id.len() {
-					this_peer_has_funded_channels = true;
-				}
-				node.is_connected
-			});
-		if inbound && !this_peer_has_funded_channels && connected_peers_without_funded_channels >= MAX_NO_CHANNEL_PEERS {
-			return Err(());
-		}
+		// peer immediately (as long as it doesn't have funded channels). If we have a bunch of
+		// unfunded channels taking up space in memory for disconnected peers, we still let new
+		// peers connect, but we'll reject new channels from them.
+		let connected_peers_without_funded_channels = self.peers_without_funded_channels(|node| node.is_connected);
+		let inbound_peer_limited = inbound && connected_peers_without_funded_channels >= MAX_NO_CHANNEL_PEERS;
 
 		{
 			let mut peer_state_lock = self.per_peer_state.write().unwrap();
 			match peer_state_lock.entry(counterparty_node_id.clone()) {
 				hash_map::Entry::Vacant(e) => {
+					if inbound_peer_limited {
+						return Err(());
+					}
 					e.insert(Mutex::new(PeerState {
 						channel_by_id: HashMap::new(),
 						latest_features: init_msg.features.clone(),
@@ -6325,6 +6321,14 @@ where
 				hash_map::Entry::Occupied(e) => {
 					let mut peer_state = e.get().lock().unwrap();
 					peer_state.latest_features = init_msg.features.clone();
+
+					let best_block_height = self.best_block.read().unwrap().height();
+					if Self::unfunded_channel_count(&*peer_state, best_block_height) ==
+						peer_state.channel_by_id.len() && inbound_peer_limited
+					{
+						return Err(());
+					}
+
 					debug_assert!(!peer_state.is_connected, "A peer shouldn't be connected twice");
 					peer_state.is_connected = true;
 				},