net: Fix inbound connection slot leak

src/net/acceptor.rs

@@ -51,6 +51,26 @@ use crate::{
 /// Atomic pointer to Acceptor
 pub type AcceptorPtr = Arc<Acceptor>;
 
+/// Releases an inbound connection slot when its tracking task exits.
+struct InboundSlotGuard {
+    acceptor: AcceptorPtr,
+    cv: Arc<CondVar>,
+}
+
+impl InboundSlotGuard {
+    fn new(acceptor: AcceptorPtr, cv: Arc<CondVar>) -> Self {
+        Self { acceptor, cv }
+    }
+}
+
+impl Drop for InboundSlotGuard {
+    fn drop(&mut self) {
+        let previous = self.acceptor.conn_count.fetch_sub(1, SeqCst);
+        debug_assert!(previous > 0, "inbound connection counter underflow");
+        self.cv.notify();
+    }
+}
+
 /// Create inbound socket connections
 pub struct Acceptor {
     channel_publisher: PublisherPtr<Result<ChannelPtr>>,
@@ -132,6 +152,11 @@ impl Acceptor {
         self.channel_publisher.clone().subscribe().await
     }
 
+    #[cfg(test)]
+    pub(super) fn connection_count(&self) -> usize {
+        self.conn_count.load(SeqCst)
+    }
+
     /// Run the accept loop in a new thread and error if a connection problem occurs
     fn accept(self: Arc<Self>, listener: Box<dyn PtListener>, ex: ExecutorPtr) {
         let self_ = self.clone();
@@ -191,15 +216,13 @@ impl Acceptor {
                     // This task will subscribe on the new channel and decrement
                     // the connection counter. Along with that, it will notify
                     // the CondVar that might be waiting to allow new connections.
-                    let self_ = self.clone();
                     let channel_ = channel.clone();
-                    let cv_ = cv.clone();
+                    let slot_guard = InboundSlotGuard::new(self.clone(), cv.clone());
                     ex.spawn(async move {
-                        let stop_sub = channel_.subscribe_stop().await?;
-                        stop_sub.receive().await;
-                        self_.conn_count.fetch_sub(1, SeqCst);
-                        cv_.notify();
-                        Ok::<(), crate::Error>(())
+                        if let Ok(stop_sub) = channel_.subscribe_stop().await {
+                            stop_sub.receive().await;
+                        }
+                        drop(slot_guard);
                     })
                     .detach();
 

src/net/channel.rs

@@ -186,12 +186,15 @@ impl Channel {
     pub async fn subscribe_stop(&self) -> Result<Subscription<Error>> {
         debug!(target: "net::channel::subscribe_stop", "START {self:?}");
 
+        let sub = self.stop_publisher.clone().subscribe().await;
+
+        // Subscribe before checking the stopped state so a concurrent stop
+        // cannot happen between the check and subscription registration.
         if self.is_stopped() {
+            sub.unsubscribe().await;
             return Err(Error::ChannelStopped)
         }
 
-        let sub = self.stop_publisher.clone().subscribe().await;
-
         debug!(target: "net::channel::subscribe_stop", "END {self:?}");
 
         Ok(sub)

src/net/session/inbound_session.rs

@@ -130,6 +130,11 @@ impl InboundSession {
         }
     }
 
+    #[cfg(test)]
+    pub(crate) async fn connection_count(&self) -> usize {
+        self.acceptors.lock().await.iter().map(|acceptor| acceptor.connection_count()).sum()
+    }
+
     /// Start accepting connections for inbound session.
     async fn start_accept_session(
         self: Arc<Self>,

src/net/tests.rs

@@ -23,11 +23,12 @@ use std::{
     net::TcpListener,
     panic,
     sync::Arc,
+    time::Duration,
 };
 
 use darkfi_serial::{async_trait, SerialDecodable, SerialEncodable};
 use rand::{prelude::SliceRandom, rngs::ThreadRng, Rng};
-use smol::{channel, future, Executor};
+use smol::{channel, future, net::TcpStream, Executor, Timer};
 use tracing::{error, info, warn};
 use url::Url;
 
@@ -39,7 +40,7 @@ use crate::{
         settings::NetworkProfile,
         P2p, Settings,
     },
-    system::sleep,
+    system::{sleep, timeout::timeout},
     util::logger::{setup_test_logger, Level},
 };
 
@@ -612,3 +613,52 @@ async fn p2p_channel_invalid_message_length_gets_banned_real(ex: Arc<Executor<'s
     node1_p2p.stop().await;
     node2_p2p.stop().await;
 }
+
+#[test]
+fn p2p_inbound_slots_survive_rapid_disconnects() {
+    test_body!(p2p_inbound_slots_survive_rapid_disconnects_real, 2);
+}
+
+async fn p2p_inbound_slots_survive_rapid_disconnects_real(ex: Arc<Executor<'static>>) {
+    const CONNECTION_LIMIT: usize = 8;
+    const CONNECTION_ATTEMPTS: usize = 2048;
+
+    let port = get_random_available_port();
+    let addr = format!("127.0.0.1:{port}");
+    let listen_url = Url::parse(&format!("tcp://{addr}")).unwrap();
+    let settings = Settings {
+        localnet: true,
+        inbound_addrs: vec![listen_url],
+        inbound_connections: CONNECTION_LIMIT,
+        outbound_connections: 0,
+        active_profiles: vec!["tcp".to_string()],
+        ..Default::default()
+    };
+
+    let p2p = P2p::new(settings, ex).await.unwrap();
+    p2p.clone().start().await.unwrap();
+
+    for _ in 0..CONNECTION_ATTEMPTS {
+        let stream = timeout(Duration::from_secs(1), TcpStream::connect(&addr))
+            .await
+            .expect("inbound listener stalled")
+            .expect("failed connecting to inbound listener");
+        drop(stream);
+    }
+
+    timeout(Duration::from_secs(5), async {
+        while p2p.session_inbound().connection_count().await != 0 {
+            Timer::after(Duration::from_millis(10)).await;
+        }
+    })
+    .await
+    .expect("inbound connection slots were not released");
+
+    let stream = timeout(Duration::from_secs(1), TcpStream::connect(&addr))
+        .await
+        .expect("inbound listener did not recover")
+        .expect("failed reconnecting to inbound listener");
+    drop(stream);
+
+    p2p.stop().await;
+}