diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-17 12:18:32 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-17 12:18:32 +0000 |
commit | 4547b622d8d29df964fa2914213088b148c498fc (patch) | |
tree | 9fc6b25f3c3add6b745be9a2400a6e96140046e9 /library/std/src/sync/mpmc/counter.rs | |
parent | Releasing progress-linux version 1.66.0+dfsg1-1~progress7.99u1. (diff) | |
download | rustc-4547b622d8d29df964fa2914213088b148c498fc.tar.xz rustc-4547b622d8d29df964fa2914213088b148c498fc.zip |
Merging upstream version 1.67.1+dfsg1.
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'library/std/src/sync/mpmc/counter.rs')
-rw-r--r-- | library/std/src/sync/mpmc/counter.rs | 137 |
1 files changed, 137 insertions, 0 deletions
diff --git a/library/std/src/sync/mpmc/counter.rs b/library/std/src/sync/mpmc/counter.rs new file mode 100644 index 000000000..a5a6bdc67 --- /dev/null +++ b/library/std/src/sync/mpmc/counter.rs @@ -0,0 +1,137 @@ +use crate::ops; +use crate::process; +use crate::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +/// Reference counter internals. +struct Counter<C> { + /// The number of senders associated with the channel. + senders: AtomicUsize, + + /// The number of receivers associated with the channel. + receivers: AtomicUsize, + + /// Set to `true` if the last sender or the last receiver reference deallocates the channel. + destroy: AtomicBool, + + /// The internal channel. + chan: C, +} + +/// Wraps a channel into the reference counter. +pub(crate) fn new<C>(chan: C) -> (Sender<C>, Receiver<C>) { + let counter = Box::into_raw(Box::new(Counter { + senders: AtomicUsize::new(1), + receivers: AtomicUsize::new(1), + destroy: AtomicBool::new(false), + chan, + })); + let s = Sender { counter }; + let r = Receiver { counter }; + (s, r) +} + +/// The sending side. +pub(crate) struct Sender<C> { + counter: *mut Counter<C>, +} + +impl<C> Sender<C> { + /// Returns the internal `Counter`. + fn counter(&self) -> &Counter<C> { + unsafe { &*self.counter } + } + + /// Acquires another sender reference. + pub(crate) fn acquire(&self) -> Sender<C> { + let count = self.counter().senders.fetch_add(1, Ordering::Relaxed); + + // Cloning senders and calling `mem::forget` on the clones could potentially overflow the + // counter. It's very difficult to recover sensibly from such degenerate scenarios so we + // just abort when the count becomes very large. + if count > isize::MAX as usize { + process::abort(); + } + + Sender { counter: self.counter } + } + + /// Releases the sender reference. + /// + /// Function `disconnect` will be called if this is the last sender reference. + pub(crate) unsafe fn release<F: FnOnce(&C) -> bool>(&self, disconnect: F) { + if self.counter().senders.fetch_sub(1, Ordering::AcqRel) == 1 { + disconnect(&self.counter().chan); + + if self.counter().destroy.swap(true, Ordering::AcqRel) { + drop(Box::from_raw(self.counter)); + } + } + } +} + +impl<C> ops::Deref for Sender<C> { + type Target = C; + + fn deref(&self) -> &C { + &self.counter().chan + } +} + +impl<C> PartialEq for Sender<C> { + fn eq(&self, other: &Sender<C>) -> bool { + self.counter == other.counter + } +} + +/// The receiving side. +pub(crate) struct Receiver<C> { + counter: *mut Counter<C>, +} + +impl<C> Receiver<C> { + /// Returns the internal `Counter`. + fn counter(&self) -> &Counter<C> { + unsafe { &*self.counter } + } + + /// Acquires another receiver reference. + pub(crate) fn acquire(&self) -> Receiver<C> { + let count = self.counter().receivers.fetch_add(1, Ordering::Relaxed); + + // Cloning receivers and calling `mem::forget` on the clones could potentially overflow the + // counter. It's very difficult to recover sensibly from such degenerate scenarios so we + // just abort when the count becomes very large. + if count > isize::MAX as usize { + process::abort(); + } + + Receiver { counter: self.counter } + } + + /// Releases the receiver reference. + /// + /// Function `disconnect` will be called if this is the last receiver reference. + pub(crate) unsafe fn release<F: FnOnce(&C) -> bool>(&self, disconnect: F) { + if self.counter().receivers.fetch_sub(1, Ordering::AcqRel) == 1 { + disconnect(&self.counter().chan); + + if self.counter().destroy.swap(true, Ordering::AcqRel) { + drop(Box::from_raw(self.counter)); + } + } + } +} + +impl<C> ops::Deref for Receiver<C> { + type Target = C; + + fn deref(&self) -> &C { + &self.counter().chan + } +} + +impl<C> PartialEq for Receiver<C> { + fn eq(&self, other: &Receiver<C>) -> bool { + self.counter == other.counter + } +} |