sync: use usize for permit count in SemaphorePermit (#8405)

tokio/src/sync/semaphore.rs

@@ -441,7 +441,7 @@ pub struct Semaphore {
 #[derive(Debug)]
 pub struct SemaphorePermit<'a> {
     sem: &'a Semaphore,
-    permits: u32,
+    permits: usize,
 }
 
 /// An owned permit from the semaphore.
@@ -454,7 +454,7 @@ pub struct SemaphorePermit<'a> {
 #[derive(Debug)]
 pub struct OwnedSemaphorePermit {
     sem: Arc<Semaphore>,
-    permits: u32,
+    permits: usize,
 }
 
 #[test]
@@ -674,7 +674,7 @@ impl Semaphore {
 
         Ok(SemaphorePermit {
             sem: self,
-            permits: n,
+            permits: n as usize,
         })
     }
 
@@ -745,7 +745,7 @@ impl Semaphore {
         match self.ll_sem.try_acquire(n as usize) {
             Ok(()) => Ok(SemaphorePermit {
                 sem: self,
-                permits: n,
+                permits: n as usize,
             }),
             Err(e) => Err(e),
         }
@@ -960,7 +960,7 @@ impl Semaphore {
         inner.await?;
         Ok(OwnedSemaphorePermit {
             sem: self,
-            permits: n,
+            permits: n as usize,
         })
     }
 
@@ -1042,7 +1042,7 @@ impl Semaphore {
         match self.ll_sem.try_acquire(n as usize) {
             Ok(()) => Ok(OwnedSemaphorePermit {
                 sem: self,
-                permits: n,
+                permits: n as usize,
             }),
             Err(e) => Err(e),
         }
@@ -1232,7 +1232,10 @@ impl<'a> SemaphorePermit<'a> {
             std::ptr::eq(self.sem, other.sem),
             "merging permits from different semaphore instances"
         );
-        self.permits += other.permits;
+        self.permits = self
+            .permits
+            .checked_add(other.permits)
+            .expect("number of permits overflowed");
         other.permits = 0;
     }
 
@@ -1255,8 +1258,6 @@ impl<'a> SemaphorePermit<'a> {
     /// assert_eq!(p2.num_permits(), 1);
     /// ```
     pub fn split(&mut self, n: usize) -> Option<Self> {
-        let n = u32::try_from(n).ok()?;
-
         if n > self.permits {
             return None;
         }
@@ -1271,7 +1272,7 @@ impl<'a> SemaphorePermit<'a> {
 
     /// Returns the number of permits held by `self`.
     pub fn num_permits(&self) -> usize {
-        self.permits as usize
+        self.permits
     }
 }
 
@@ -1339,7 +1340,10 @@ impl OwnedSemaphorePermit {
             Arc::ptr_eq(&self.sem, &other.sem),
             "merging permits from different semaphore instances"
         );
-        self.permits += other.permits;
+        self.permits = self
+            .permits
+            .checked_add(other.permits)
+            .expect("number of permits overflowed");
         other.permits = 0;
     }
 
@@ -1366,8 +1370,6 @@ impl OwnedSemaphorePermit {
     /// assert_eq!(p2.num_permits(), 1);
     /// ```
     pub fn split(&mut self, n: usize) -> Option<Self> {
-        let n = u32::try_from(n).ok()?;
-
         if n > self.permits {
             return None;
         }
@@ -1387,18 +1389,18 @@ impl OwnedSemaphorePermit {
 
     /// Returns the number of permits held by `self`.
     pub fn num_permits(&self) -> usize {
-        self.permits as usize
+        self.permits
     }
 }
 
 impl Drop for SemaphorePermit<'_> {
     fn drop(&mut self) {
-        self.sem.add_permits(self.permits as usize);
+        self.sem.add_permits(self.permits);
     }
 }
 
 impl Drop for OwnedSemaphorePermit {
     fn drop(&mut self) {
-        self.sem.add_permits(self.permits as usize);
+        self.sem.add_permits(self.permits);
     }
 }

tokio/tests/sync_semaphore.rs

@@ -326,3 +326,21 @@ async fn blocking_acquire_many_owned_in_async_context() {
     // Calling a blocking method from an async context must panic.
     let _permit = sem.blocking_acquire_many_owned(1);
 }
+
+#[test]
+#[cfg(target_pointer_width = "64")]
+fn merge_many_permits() {
+    let sem = Arc::new(Semaphore::new(6_000_000_000));
+    let mut a = sem.try_acquire_many(3_000_000_000).unwrap();
+    let b = sem.try_acquire_many(3_000_000_000).unwrap();
+    assert_eq!(a.num_permits(), 3_000_000_000);
+    assert_eq!(b.num_permits(), 3_000_000_000);
+    assert_eq!(sem.available_permits(), 0);
+
+    a.merge(b);
+    assert_eq!(a.num_permits(), 6_000_000_000);
+    assert_eq!(sem.available_permits(), 0);
+
+    drop(a);
+    assert_eq!(sem.available_permits(), 6_000_000_000);
+}

tokio/tests/sync_semaphore_owned.rs

@@ -162,3 +162,21 @@ async fn stress_test() {
     let _p5 = sem.clone().try_acquire_owned().unwrap();
     assert!(sem.try_acquire_owned().is_err());
 }
+
+#[test]
+#[cfg(target_pointer_width = "64")]
+fn merge_many_permits_owned() {
+    let sem = Arc::new(Semaphore::new(6_000_000_000));
+    let mut a = sem.clone().try_acquire_many_owned(3_000_000_000).unwrap();
+    let b = sem.clone().try_acquire_many_owned(3_000_000_000).unwrap();
+    assert_eq!(a.num_permits(), 3_000_000_000);
+    assert_eq!(b.num_permits(), 3_000_000_000);
+    assert_eq!(sem.available_permits(), 0);
+
+    a.merge(b);
+    assert_eq!(a.num_permits(), 6_000_000_000);
+    assert_eq!(sem.available_permits(), 0);
+
+    drop(a);
+    assert_eq!(sem.available_permits(), 6_000_000_000);
+}