diff --git a/tokio/src/sync/semaphore.rs b/tokio/src/sync/semaphore.rs index b4c2e3a92..537d0b3c8 100644 --- a/tokio/src/sync/semaphore.rs +++ b/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, - 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 { - 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 { - 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); } } diff --git a/tokio/tests/sync_semaphore.rs b/tokio/tests/sync_semaphore.rs index f204f85c4..113dee064 100644 --- a/tokio/tests/sync_semaphore.rs +++ b/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); +} diff --git a/tokio/tests/sync_semaphore_owned.rs b/tokio/tests/sync_semaphore_owned.rs index f9eeee0cf..1bae1e9e0 100644 --- a/tokio/tests/sync_semaphore_owned.rs +++ b/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); +}