Libp2p stream limits

crates/subspace-networking/src/create.rs

@@ -1,3 +1,6 @@
+#[cfg(test)]
+mod tests;
+
 pub use crate::behavior::custom_record_store::ValueGetter;
 use crate::behavior::custom_record_store::{
     CustomRecordStore, MemoryProviderStorage, NoRecordStorage,
@@ -88,6 +91,7 @@ const SEMAPHORE_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(5);
 
 async fn maintain_semaphore_permits_capacity(
     semaphore: &Semaphore,
+    interval: Duration,
     connected_peers_count_weak: Weak<AtomicUsize>,
     boost_per_peer: usize,
 ) {
@@ -99,13 +103,13 @@ async fn maintain_semaphore_permits_capacity(
     let mut reserved_permits = Vec::new();
     loop {
         let connected_peers_count = match connected_peers_count_weak.upgrade() {
-            Some(connected_peers_count) => connected_peers_count,
+            Some(connected_peers_count) => connected_peers_count.load(Ordering::Relaxed),
             None => {
                 return;
             }
         };
         let expected_total_permits =
-            base_permits + connected_peers_count.load(Ordering::Relaxed) * boost_per_peer;
+            base_permits + connected_peers_count.saturating_sub(1) * boost_per_peer;
 
         // Release reserves to match expected number of permits if necessary
         while total_permits < expected_total_permits && !reserved_permits.is_empty() {
@@ -133,7 +137,7 @@ async fn maintain_semaphore_permits_capacity(
             total_permits = expected_total_permits;
         }
 
-        tokio::time::sleep(SEMAPHORE_MAINTENANCE_INTERVAL).await;
+        tokio::time::sleep(interval).await;
     }
 }
 
@@ -371,6 +375,7 @@ where
             async move {
                 maintain_semaphore_permits_capacity(
                     &kademlia_tasks_semaphore,
+                    SEMAPHORE_MAINTENANCE_INTERVAL,
                     connected_peers_count_weak,
                     KADEMLIA_CONCURRENT_TASKS_BOOST_PER_PEER,
                 )
@@ -384,6 +389,7 @@ where
             async move {
                 maintain_semaphore_permits_capacity(
                     &regular_tasks_semaphore,
+                    SEMAPHORE_MAINTENANCE_INTERVAL,
                     connected_peers_count_weak,
                     REGULAR_CONCURRENT_TASKS_BOOST_PER_PEER,
                 )

crates/subspace-networking/src/create/tests.rs

@@ -0,0 +1,119 @@
+use futures::future::{select, Either};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::Semaphore;
+use tokio::time::sleep;
+
+#[tokio::test]
+async fn maintain_semaphore_permits_capacity() {
+    let base_tasks = 2;
+    let boost_per_peer = 1;
+    let interval = Duration::from_micros(1);
+    let connected_peers_count = Arc::new(AtomicUsize::new(0));
+    let tasks_semaphore = Arc::new(Semaphore::new(base_tasks));
+
+    tokio::spawn({
+        let tasks_semaphore = Arc::clone(&tasks_semaphore);
+        let connected_peers_count_weak = Arc::downgrade(&connected_peers_count);
+
+        async move {
+            super::maintain_semaphore_permits_capacity(
+                &tasks_semaphore,
+                interval,
+                connected_peers_count_weak,
+                boost_per_peer,
+            )
+            .await;
+        }
+    });
+
+    let timeout = Duration::from_millis(100);
+
+    // Let above function time to run at least one loop
+    sleep(timeout).await;
+
+    let permit_1_result = select(
+        Box::pin(tasks_semaphore.acquire()),
+        Box::pin(sleep(timeout)),
+    )
+    .await;
+    if !matches!(permit_1_result, Either::Left(_)) {
+        panic!("Must be able to acquire the permit");
+    }
+
+    let permit_2_result = select(
+        Box::pin(tasks_semaphore.acquire()),
+        Box::pin(sleep(timeout)),
+    )
+    .await;
+    if !matches!(permit_2_result, Either::Left(_)) {
+        panic!("Must be able to acquire the second permit");
+    }
+
+    {
+        let permit_3_result = select(
+            Box::pin(tasks_semaphore.acquire()),
+            Box::pin(sleep(timeout)),
+        )
+        .await;
+        if !matches!(permit_3_result, Either::Right(_)) {
+            panic!("Must not be able to acquire the third permit due to capacity");
+        }
+    }
+
+    // Increase capacity
+    connected_peers_count.fetch_add(1, Ordering::SeqCst);
+
+    {
+        let permit_3_result = select(
+            Box::pin(tasks_semaphore.acquire()),
+            Box::pin(sleep(timeout)),
+        )
+        .await;
+        if !matches!(permit_3_result, Either::Right(_)) {
+            panic!("Must not be able to acquire the third permit due to capacity");
+        }
+    }
+
+    // Increase capacity more
+    connected_peers_count.fetch_add(1, Ordering::SeqCst);
+
+    let permit_3_result = select(
+        Box::pin(tasks_semaphore.acquire()),
+        Box::pin(sleep(timeout)),
+    )
+    .await;
+    if !matches!(permit_3_result, Either::Left(_)) {
+        panic!("Must be able to acquire the third permit due to increased capacity");
+    }
+
+    {
+        let permit_4_result = select(
+            Box::pin(tasks_semaphore.acquire()),
+            Box::pin(sleep(timeout)),
+        )
+        .await;
+        if !matches!(permit_4_result, Either::Right(_)) {
+            panic!("Must not be able to acquire the fourth permit due to capacity");
+        }
+    }
+
+    // Decrease capacity capacity
+    connected_peers_count.fetch_sub(1, Ordering::SeqCst);
+
+    drop(permit_3_result);
+
+    sleep(timeout).await;
+
+    {
+        let permit_3_result = select(
+            Box::pin(tasks_semaphore.acquire()),
+            Box::pin(sleep(timeout)),
+        )
+        .await;
+        if !matches!(permit_3_result, Either::Right(_)) {
+            panic!("Must not be able to acquire the third permit again due to capacity anymore");
+        }
+    }
+}