Skip to content

Commit

Permalink
sync::broadcast: create Sender::new
Browse files Browse the repository at this point in the history
  • Loading branch information
marcospb19 committed Jun 26, 2023
1 parent 910a1e2 commit 374bb1b
Showing 1 changed file with 41 additions and 31 deletions.
72 changes: 41 additions & 31 deletions tokio/src/sync/broadcast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -444,42 +444,14 @@ const MAX_RECEIVERS: usize = usize::MAX >> 2;
/// This will panic if `capacity` is equal to `0` or larger
/// than `usize::MAX / 2`.
#[track_caller]
pub fn channel<T: Clone>(mut capacity: usize) -> (Sender<T>, Receiver<T>) {
assert!(capacity > 0, "capacity is empty");
assert!(capacity <= usize::MAX >> 1, "requested capacity too large");

// Round to a power of two
capacity = capacity.next_power_of_two();

let mut buffer = Vec::with_capacity(capacity);

for i in 0..capacity {
buffer.push(RwLock::new(Slot {
rem: AtomicUsize::new(0),
pos: (i as u64).wrapping_sub(capacity as u64),
val: UnsafeCell::new(None),
}));
}

let shared = Arc::new(Shared {
buffer: buffer.into_boxed_slice(),
mask: capacity - 1,
tail: Mutex::new(Tail {
pos: 0,
rx_cnt: 1,
closed: false,
waiters: LinkedList::new(),
}),
num_tx: AtomicUsize::new(1),
});
pub fn channel<T: Clone>(capacity: usize) -> (Sender<T>, Receiver<T>) {
let tx = Sender::new(capacity);

let rx = Receiver {
shared: shared.clone(),
shared: tx.shared.clone(),
next: 0,
};

let tx = Sender { shared };

(tx, rx)
}

Expand All @@ -490,6 +462,44 @@ unsafe impl<T: Send> Send for Receiver<T> {}
unsafe impl<T: Send> Sync for Receiver<T> {}

impl<T> Sender<T> {
/// Creates the sending-half of the [`broadcast`] channel.
///
/// See documentation of [`broadcast::channel`] for errors when calling this function.
///
/// [`broadcast`]: crate::sync::broadcast
/// [`broadcast::channel`]: crate::sync::broadcast
pub fn new(mut capacity: usize) -> Self {
assert!(capacity > 0, "capacity is empty");
assert!(capacity <= usize::MAX >> 1, "requested capacity too large");

// Round to a power of two
capacity = capacity.next_power_of_two();

let mut buffer = Vec::with_capacity(capacity);

for i in 0..capacity {
buffer.push(RwLock::new(Slot {
rem: AtomicUsize::new(0),
pos: (i as u64).wrapping_sub(capacity as u64),
val: UnsafeCell::new(None),
}));
}

let shared = Arc::new(Shared {
buffer: buffer.into_boxed_slice(),
mask: capacity - 1,
tail: Mutex::new(Tail {
pos: 0,
rx_cnt: 1,
closed: false,
waiters: LinkedList::new(),
}),
num_tx: AtomicUsize::new(1),
});

Sender { shared }
}

/// Attempts to send a value to all active [`Receiver`] handles, returning
/// it back if it could not be sent.
///
Expand Down

0 comments on commit 374bb1b

Please sign in to comment.