sync: make add_permits panic with usize::MAX >> 3 permits (#3188)

tokio-rs · Dec 2, 2020 · a6051a6 · a6051a6
1 parent a8e0f0a
commit a6051a6
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/tokio/src/sync/batch_semaphore.rs b/tokio/src/sync/batch_semaphore.rs
@@ -253,9 +253,9 @@ impl Semaphore {
             }
 
             if rem > 0 && is_empty {
-                let permits = rem << Self::PERMIT_SHIFT;
+                let permits = rem;
                 assert!(
-                    permits < Self::MAX_PERMITS,
+                    permits <= Self::MAX_PERMITS,
                     "cannot add more than MAX_PERMITS permits ({})",
                     Self::MAX_PERMITS
                 );

diff --git a/tokio/tests/sync_semaphore.rs b/tokio/tests/sync_semaphore.rs
@@ -79,3 +79,17 @@ async fn stresstest() {
     let _p5 = sem.try_acquire().unwrap();
     assert!(sem.try_acquire().is_err());
 }
+
+#[test]
+fn add_max_amount_permits() {
+    let s = tokio::sync::Semaphore::new(0);
+    s.add_permits(usize::MAX >> 3);
+    assert_eq!(s.available_permits(), usize::MAX >> 3);
+}
+
+#[test]
+#[should_panic]
+fn add_more_than_max_amount_permits() {
+    let s = tokio::sync::Semaphore::new(1);
+    s.add_permits(usize::MAX >> 3);
+}