use std::collections::VecDeque;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Condvar, Mutex, MutexGuard};
use std::task::{Context, Poll, Waker};
struct State<T> {
buffer: VecDeque<T>,
waker: Option<Waker>,
closed: bool,
consumer_gone: bool,
sender_count: usize,
#[cfg(test)]
blocked_producers: usize,
}
struct Shared<T> {
state: Mutex<State<T>>,
capacity_available: Condvar,
capacity: usize,
}
impl<T> Shared<T> {
fn lock_state(&self) -> MutexGuard<'_, State<T>> {
self.state
.lock()
.unwrap_or_else(|_| panic!("BoundedAsyncStream state mutex poisoned"))
}
fn lock_state_for_drop(&self) -> MutexGuard<'_, State<T>> {
self.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[cfg(test)]
fn wait_for_blocked_producers(&self, expected: usize) {
let (state, timeout) = self
.capacity_available
.wait_timeout_while(
self.lock_state(),
std::time::Duration::from_secs(5),
|state| state.blocked_producers < expected,
)
.unwrap_or_else(|_| panic!("BoundedAsyncStream state mutex poisoned"));
assert!(
!timeout.timed_out(),
"timed out waiting for {expected} blocked producer(s)"
);
drop(state);
}
}
pub struct BoundedAsyncStream<T> {
shared: Arc<Shared<T>>,
}
pub struct AsyncStreamSender<T> {
shared: Arc<Shared<T>>,
}
impl<T> Clone for AsyncStreamSender<T> {
fn clone(&self) -> Self {
let mut state = self.shared.lock_state();
state.sender_count = state
.sender_count
.checked_add(1)
.expect("AsyncStreamSender count overflow");
drop(state);
Self {
shared: Arc::clone(&self.shared),
}
}
}
impl<T> fmt::Debug for BoundedAsyncStream<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BoundedAsyncStream")
.field("buffered", &self.buffered_count())
.field("capacity", &self.capacity())
.field("is_closed", &self.is_closed())
.finish_non_exhaustive()
}
}
impl<T> fmt::Debug for AsyncStreamSender<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("AsyncStreamSender").finish_non_exhaustive()
}
}
impl<T> BoundedAsyncStream<T> {
#[must_use]
pub fn new(capacity: usize) -> (Self, AsyncStreamSender<T>) {
assert!(capacity > 0, "BoundedAsyncStream capacity must be > 0");
let shared = Arc::new(Shared {
capacity,
capacity_available: Condvar::new(),
state: Mutex::new(State {
buffer: VecDeque::with_capacity(capacity),
waker: None,
closed: false,
consumer_gone: false,
sender_count: 1,
#[cfg(test)]
blocked_producers: 0,
}),
});
let stream = Self {
shared: Arc::clone(&shared),
};
let sender = AsyncStreamSender { shared };
(stream, sender)
}
#[must_use]
pub const fn next(&self) -> NextItem<'_, T> {
NextItem { stream: self }
}
#[must_use]
pub fn try_next(&self) -> Option<T> {
let item = {
let mut state = self.shared.lock_state();
state.buffer.pop_front()
};
if item.is_some() {
self.shared.capacity_available.notify_one();
}
item
}
#[must_use]
pub fn is_closed(&self) -> bool {
self.shared.lock_state().closed
}
#[must_use]
pub fn buffered_count(&self) -> usize {
self.shared.lock_state().buffer.len()
}
#[must_use]
pub fn capacity(&self) -> usize {
self.shared.capacity
}
pub fn clear_buffer(&self) {
let mut cleared = VecDeque::with_capacity(self.shared.capacity);
{
let mut state = self.shared.lock_state();
std::mem::swap(&mut state.buffer, &mut cleared);
}
self.shared.capacity_available.notify_all();
drop(cleared);
}
fn poll_next_item(&self, cx: &Context<'_>) -> Poll<Option<T>> {
let mut next_waker = Some(cx.waker().clone());
let (result, replaced_waker, made_room) = {
let mut state = self.shared.lock_state();
let outcome = match state.buffer.pop_front() {
Some(item) => (Poll::Ready(Some(item)), None, true),
None if state.closed => (Poll::Ready(None), None, false),
None => {
let replaced_waker = match state.waker.as_ref() {
Some(existing) if existing.will_wake(cx.waker()) => None,
_ => state
.waker
.replace(next_waker.take().expect("next waker present")),
};
(Poll::Pending, replaced_waker, false)
}
};
drop(state);
outcome
};
drop(next_waker);
drop(replaced_waker);
if made_room {
self.shared.capacity_available.notify_one();
}
result
}
}
impl<T> Drop for BoundedAsyncStream<T> {
fn drop(&mut self) {
let stale_waker = {
let mut state = self.shared.lock_state_for_drop();
state.consumer_gone = true;
state.waker.take()
};
self.shared.capacity_available.notify_all();
drop(stale_waker);
}
}
impl<T> AsyncStreamSender<T> {
pub fn push(&self, item: T) {
let (overwritten, waker) = {
let mut state = self.shared.lock_state();
let overwritten = if state.buffer.len() >= self.shared.capacity {
state.buffer.pop_front()
} else {
None
};
state.buffer.push_back(item);
(overwritten, state.waker.take())
};
if let Some(waker) = waker {
waker.wake();
}
drop(overwritten);
}
pub fn push_or_block(&self, item: T) -> Result<(), T> {
let mut state = self.shared.lock_state();
if !state.consumer_gone && state.buffer.len() >= self.shared.capacity {
#[cfg(test)]
{
state.blocked_producers += 1;
self.shared.capacity_available.notify_all();
}
state = self
.shared
.capacity_available
.wait_while(state, |state| {
!state.consumer_gone && state.buffer.len() >= self.shared.capacity
})
.unwrap_or_else(|_| panic!("BoundedAsyncStream state mutex poisoned"));
#[cfg(test)]
{
state.blocked_producers -= 1;
self.shared.capacity_available.notify_all();
}
}
if state.consumer_gone {
drop(state);
return Err(item);
}
state.buffer.push_back(item);
let waker = state.waker.take();
drop(state);
if let Some(waker) = waker {
waker.wake();
}
Ok(())
}
#[must_use]
pub fn buffered_count(&self) -> usize {
self.shared.lock_state().buffer.len()
}
#[must_use]
pub fn is_consumer_gone(&self) -> bool {
self.shared.lock_state().consumer_gone
}
}
impl<T> Drop for AsyncStreamSender<T> {
fn drop(&mut self) {
let waker = {
let mut state = self.shared.lock_state_for_drop();
state.sender_count -= 1;
if state.sender_count == 0 {
state.closed = true;
state.waker.take()
} else {
None
}
};
if let Some(waker) = waker {
waker.wake();
}
}
}
pub struct NextItem<'a, T> {
stream: &'a BoundedAsyncStream<T>,
}
impl<T> fmt::Debug for NextItem<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("NextItem").finish_non_exhaustive()
}
}
impl<T> Future for NextItem<'_, T> {
type Output = Option<T>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.stream.poll_next_item(cx)
}
}
#[cfg(feature = "futures-stream")]
impl<T: 'static> futures_core::Stream for BoundedAsyncStream<T> {
type Item = T;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
self.poll_next_item(cx)
}
}
#[cfg(test)]
mod tests {
#[cfg(feature = "futures-stream")]
use std::future::poll_fn;
use std::future::Future;
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{mpsc, Arc, Barrier, TryLockError, Weak};
use std::task::{Context, Poll, Wake, Waker};
use std::thread::{self, JoinHandle};
use std::time::Duration;
#[cfg(feature = "futures-stream")]
use futures_core::Stream;
use super::{AsyncStreamSender, BoundedAsyncStream, Shared};
const TEST_TIMEOUT: Duration = Duration::from_secs(5);
fn spawn_blocking_push(
sender: AsyncStreamSender<u32>,
item: u32,
) -> (mpsc::Receiver<Result<(), u32>>, JoinHandle<()>) {
let (result_tx, result_rx) = mpsc::channel();
let handle = thread::spawn(move || {
let result = sender.push_or_block(item);
result_tx.send(result).unwrap();
});
(result_rx, handle)
}
#[test]
fn try_next_notifies_blocked_producer() {
let (stream, sender) = BoundedAsyncStream::new(1);
sender.push(1);
let (result_rx, handle) = spawn_blocking_push(sender, 2);
stream.shared.wait_for_blocked_producers(1);
assert_eq!(stream.try_next(), Some(1));
assert_eq!(result_rx.recv_timeout(TEST_TIMEOUT).unwrap(), Ok(()));
handle.join().unwrap();
assert_eq!(stream.try_next(), Some(2));
}
#[test]
fn next_future_notifies_blocked_producer() {
let (stream, sender) = BoundedAsyncStream::new(1);
sender.push(1);
let (result_rx, handle) = spawn_blocking_push(sender, 2);
stream.shared.wait_for_blocked_producers(1);
assert_eq!(pollster::block_on(stream.next()), Some(1));
assert_eq!(result_rx.recv_timeout(TEST_TIMEOUT).unwrap(), Ok(()));
handle.join().unwrap();
assert_eq!(pollster::block_on(stream.next()), Some(2));
}
#[cfg(feature = "futures-stream")]
#[test]
fn stream_poll_notifies_blocked_producer() {
let (mut stream, sender) = BoundedAsyncStream::new(1);
sender.push(1);
let (result_rx, handle) = spawn_blocking_push(sender, 2);
stream.shared.wait_for_blocked_producers(1);
let first = pollster::block_on(poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)));
assert_eq!(first, Some(1));
assert_eq!(result_rx.recv_timeout(TEST_TIMEOUT).unwrap(), Ok(()));
handle.join().unwrap();
let second = pollster::block_on(poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)));
assert_eq!(second, Some(2));
}
#[test]
fn clear_buffer_notifies_blocked_producer() {
let (stream, sender) = BoundedAsyncStream::new(1);
sender.push(1);
let (result_rx, handle) = spawn_blocking_push(sender, 2);
stream.shared.wait_for_blocked_producers(1);
stream.clear_buffer();
assert_eq!(result_rx.recv_timeout(TEST_TIMEOUT).unwrap(), Ok(()));
handle.join().unwrap();
assert_eq!(stream.try_next(), Some(2));
}
#[test]
fn consumer_drop_returns_blocked_item() {
let (stream, sender) = BoundedAsyncStream::new(1);
sender.push(1);
let (result_rx, handle) = spawn_blocking_push(sender, 2);
stream.shared.wait_for_blocked_producers(1);
drop(stream);
assert_eq!(result_rx.recv_timeout(TEST_TIMEOUT).unwrap(), Err(2));
handle.join().unwrap();
}
#[test]
fn concurrent_sender_clone_drops_close_stream() {
let (stream, sender) = BoundedAsyncStream::<u32>::new(1);
let sender_clone = sender.clone();
let barrier = Arc::new(Barrier::new(3));
let first_barrier = Arc::clone(&barrier);
let first = thread::spawn(move || {
first_barrier.wait();
drop(sender);
});
let second_barrier = Arc::clone(&barrier);
let second = thread::spawn(move || {
second_barrier.wait();
drop(sender_clone);
});
barrier.wait();
first.join().unwrap();
second.join().unwrap();
assert!(stream.is_closed());
assert_eq!(pollster::block_on(stream.next()), None);
}
fn shared_state_is_unlocked<T>(shared: &Weak<Shared<T>>) -> bool {
let Some(shared) = shared.upgrade() else {
return true;
};
let unlocked = match shared.state.try_lock() {
Ok(_state) => true,
Err(TryLockError::WouldBlock | TryLockError::Poisoned(_)) => false,
};
unlocked
}
struct ReentrantWaker {
shared: Weak<Shared<u32>>,
wake_was_unlocked: Arc<AtomicBool>,
drop_was_unlocked: Arc<AtomicBool>,
}
impl Wake for ReentrantWaker {
fn wake(self: Arc<Self>) {
self.wake_was_unlocked
.store(shared_state_is_unlocked(&self.shared), Ordering::SeqCst);
}
}
impl Drop for ReentrantWaker {
fn drop(&mut self) {
self.drop_was_unlocked
.store(shared_state_is_unlocked(&self.shared), Ordering::SeqCst);
}
}
#[test]
fn wakes_and_drops_waker_outside_state_mutex() {
let (stream, sender) = BoundedAsyncStream::new(1);
let wake_was_unlocked = Arc::new(AtomicBool::new(false));
let drop_was_unlocked = Arc::new(AtomicBool::new(false));
let probe = Arc::new(ReentrantWaker {
shared: Arc::downgrade(&stream.shared),
wake_was_unlocked: Arc::clone(&wake_was_unlocked),
drop_was_unlocked: Arc::clone(&drop_was_unlocked),
});
let waker = Waker::from(Arc::clone(&probe));
let mut next = stream.next();
{
let mut cx = Context::from_waker(&waker);
assert_eq!(Pin::new(&mut next).poll(&mut cx), Poll::Pending);
}
drop(waker);
drop(probe);
sender.push(1);
assert!(wake_was_unlocked.load(Ordering::SeqCst));
assert!(drop_was_unlocked.load(Ordering::SeqCst));
}
struct ReentrantItem {
shared: Weak<Shared<Self>>,
unlocked_drops: Arc<AtomicUsize>,
locked_drops: Arc<AtomicUsize>,
}
impl Drop for ReentrantItem {
fn drop(&mut self) {
if shared_state_is_unlocked(&self.shared) {
self.unlocked_drops.fetch_add(1, Ordering::SeqCst);
} else {
self.locked_drops.fetch_add(1, Ordering::SeqCst);
}
}
}
#[test]
fn overwritten_and_cleared_items_drop_outside_state_mutex() {
let (stream, sender) = BoundedAsyncStream::new(1);
let unlocked_drops = Arc::new(AtomicUsize::new(0));
let locked_drops = Arc::new(AtomicUsize::new(0));
let make_item = || ReentrantItem {
shared: Arc::downgrade(&stream.shared),
unlocked_drops: Arc::clone(&unlocked_drops),
locked_drops: Arc::clone(&locked_drops),
};
sender.push(make_item());
sender.push(make_item());
assert_eq!(unlocked_drops.load(Ordering::SeqCst), 1);
stream.clear_buffer();
assert_eq!(unlocked_drops.load(Ordering::SeqCst), 2);
assert_eq!(locked_drops.load(Ordering::SeqCst), 0);
}
#[test]
fn poisoned_state_does_not_masquerade_as_close_or_delivery() {
let (stream, sender) = BoundedAsyncStream::new(1);
let shared = Arc::clone(&stream.shared);
assert!(thread::spawn(move || {
let _state = shared.state.lock().unwrap();
panic!("poison stream state");
})
.join()
.is_err());
assert!(catch_unwind(AssertUnwindSafe(|| sender.push(1))).is_err());
assert!(catch_unwind(AssertUnwindSafe(|| stream.try_next())).is_err());
assert!(catch_unwind(AssertUnwindSafe(|| pollster::block_on(stream.next()))).is_err());
}
}