509 lines
15 KiB
Rust
509 lines
15 KiB
Rust
// Copyright 2024 tison <wander4096@gmail.com>
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
use std::future::Future;
|
|
use std::future::IntoFuture;
|
|
use std::hint::spin_loop;
|
|
use std::mem;
|
|
use std::pin::Pin;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::AtomicU32;
|
|
use std::sync::atomic::AtomicUsize;
|
|
use std::sync::atomic::Ordering;
|
|
use std::task::Context;
|
|
use std::task::Poll;
|
|
use std::task::RawWaker;
|
|
use std::task::RawWakerVTable;
|
|
use std::task::Waker;
|
|
use std::time::Duration;
|
|
use std::time::Instant;
|
|
|
|
use crate::oneshot;
|
|
use crate::oneshot::TryRecvError;
|
|
|
|
struct DropCounterHandle(Arc<AtomicUsize>);
|
|
|
|
impl DropCounterHandle {
|
|
pub fn count(&self) -> usize {
|
|
self.0.load(Ordering::SeqCst)
|
|
}
|
|
}
|
|
|
|
struct DropCounter<T> {
|
|
drop_count: Arc<AtomicUsize>,
|
|
value: Option<T>,
|
|
}
|
|
|
|
impl<T> DropCounter<T> {
|
|
fn new(value: T) -> (Self, DropCounterHandle) {
|
|
let drop_count = Arc::new(AtomicUsize::new(0));
|
|
(
|
|
Self {
|
|
drop_count: drop_count.clone(),
|
|
value: Some(value),
|
|
},
|
|
DropCounterHandle(drop_count),
|
|
)
|
|
}
|
|
|
|
fn value(&self) -> &T {
|
|
self.value.as_ref().unwrap()
|
|
}
|
|
}
|
|
|
|
impl<T> Drop for DropCounter<T> {
|
|
fn drop(&mut self) {
|
|
self.drop_count.fetch_add(1, Ordering::SeqCst);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_before_await() {
|
|
let (sender, receiver) = oneshot::channel();
|
|
assert!(sender.send(19i128).is_ok());
|
|
assert_eq!(receiver.await, Ok(19i128));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn await_with_dropped_sender() {
|
|
let (sender, receiver) = oneshot::channel::<u128>();
|
|
drop(sender);
|
|
receiver.await.unwrap_err();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn await_before_send() {
|
|
let (sender, receiver) = oneshot::channel();
|
|
let (message, counter) = DropCounter::new(79u128);
|
|
let t = tokio::spawn(async move {
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
sender.send(message)
|
|
});
|
|
let returned_message = receiver.await.unwrap();
|
|
assert_eq!(counter.count(), 0);
|
|
assert_eq!(*returned_message.value(), 79u128);
|
|
drop(returned_message);
|
|
assert_eq!(counter.count(), 1);
|
|
t.await.unwrap().unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn await_before_send_then_drop_sender() {
|
|
let (sender, receiver) = oneshot::channel::<u128>();
|
|
let t = tokio::spawn(async {
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
drop(sender);
|
|
});
|
|
assert!(receiver.await.is_err());
|
|
t.await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn poll_receiver_then_drop_it() {
|
|
let (sender, receiver) = oneshot::channel::<()>();
|
|
// This will poll the receiver and then give up after 100 ms.
|
|
tokio::time::timeout(Duration::from_millis(100), receiver)
|
|
.await
|
|
.unwrap_err();
|
|
// Make sure the receiver has been dropped by the runtime.
|
|
assert!(sender.send(()).is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn recv_within_select() {
|
|
let (tx, rx) = oneshot::channel::<&'static str>();
|
|
let mut interval = tokio::time::interval(Duration::from_millis(10));
|
|
|
|
let handle = tokio::spawn(async move {
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
tx.send("shut down").unwrap();
|
|
});
|
|
|
|
let mut recv = rx.into_future();
|
|
loop {
|
|
tokio::select! {
|
|
_ = interval.tick() => println!("another 10ms"),
|
|
msg = &mut recv => {
|
|
println!("Got message: {}", msg.unwrap());
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
handle.await.unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn try_recv_success_then_disconnected() {
|
|
let (tx, rx) = oneshot::channel::<i32>();
|
|
tx.send(10).unwrap();
|
|
|
|
assert_eq!(rx.try_recv(), Ok(10));
|
|
assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
|
|
}
|
|
|
|
#[test]
|
|
fn try_recv_empty_with_live_sender() {
|
|
let (_tx, rx) = oneshot::channel::<()>();
|
|
assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
|
|
}
|
|
|
|
#[test]
|
|
fn try_recv_disconnected_after_drop() {
|
|
let (tx, rx) = oneshot::channel::<()>();
|
|
drop(tx);
|
|
assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
|
|
}
|
|
|
|
#[derive(Default)]
|
|
pub struct WakerHandle {
|
|
clone_count: AtomicU32,
|
|
drop_count: AtomicU32,
|
|
wake_count: AtomicU32,
|
|
}
|
|
|
|
impl WakerHandle {
|
|
pub fn clone_count(&self) -> u32 {
|
|
self.clone_count.load(Ordering::Relaxed)
|
|
}
|
|
|
|
pub fn drop_count(&self) -> u32 {
|
|
self.drop_count.load(Ordering::Relaxed)
|
|
}
|
|
|
|
pub fn wake_count(&self) -> u32 {
|
|
self.wake_count.load(Ordering::Relaxed)
|
|
}
|
|
}
|
|
|
|
fn waker() -> (Waker, Arc<WakerHandle>) {
|
|
let waker_handle = Arc::new(WakerHandle::default());
|
|
let waker_handle_ptr = Arc::into_raw(waker_handle.clone());
|
|
let raw_waker = RawWaker::new(waker_handle_ptr as *const _, waker_vtable());
|
|
(unsafe { Waker::from_raw(raw_waker) }, waker_handle)
|
|
}
|
|
|
|
fn waker_vtable() -> &'static RawWakerVTable {
|
|
&RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw)
|
|
}
|
|
|
|
unsafe fn clone_raw(data: *const ()) -> RawWaker {
|
|
let handle: Arc<WakerHandle> = unsafe { Arc::from_raw(data as *const _) };
|
|
handle.clone_count.fetch_add(1, Ordering::Relaxed);
|
|
mem::forget(handle.clone());
|
|
mem::forget(handle);
|
|
RawWaker::new(data, waker_vtable())
|
|
}
|
|
|
|
unsafe fn wake_raw(data: *const ()) {
|
|
let handle: Arc<WakerHandle> = unsafe { Arc::from_raw(data as *const _) };
|
|
handle.wake_count.fetch_add(1, Ordering::Relaxed);
|
|
handle.drop_count.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
|
|
unsafe fn wake_by_ref_raw(data: *const ()) {
|
|
let handle: Arc<WakerHandle> = unsafe { Arc::from_raw(data as *const _) };
|
|
handle.wake_count.fetch_add(1, Ordering::Relaxed);
|
|
mem::forget(handle)
|
|
}
|
|
|
|
unsafe fn drop_raw(data: *const ()) {
|
|
let handle: Arc<WakerHandle> = unsafe { Arc::from_raw(data as *const _) };
|
|
handle.drop_count.fetch_add(1, Ordering::Relaxed);
|
|
drop(handle)
|
|
}
|
|
|
|
#[test]
|
|
fn poll_then_send() {
|
|
let (sender, receiver) = oneshot::channel::<u128>();
|
|
let mut receiver = receiver.into_future();
|
|
|
|
let (waker, waker_handle) = waker();
|
|
let mut context = Context::from_waker(&waker);
|
|
|
|
assert_eq!(Pin::new(&mut receiver).poll(&mut context), Poll::Pending);
|
|
assert_eq!(waker_handle.clone_count(), 1);
|
|
assert_eq!(waker_handle.drop_count(), 0);
|
|
assert_eq!(waker_handle.wake_count(), 0);
|
|
|
|
sender.send(1234).unwrap();
|
|
assert_eq!(waker_handle.clone_count(), 1);
|
|
assert_eq!(waker_handle.drop_count(), 1);
|
|
assert_eq!(waker_handle.wake_count(), 1);
|
|
|
|
assert_eq!(
|
|
Pin::new(&mut receiver).poll(&mut context),
|
|
Poll::Ready(Ok(1234))
|
|
);
|
|
assert_eq!(waker_handle.clone_count(), 1);
|
|
assert_eq!(waker_handle.drop_count(), 1);
|
|
assert_eq!(waker_handle.wake_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn poll_with_different_wakers() {
|
|
let (sender, receiver) = oneshot::channel::<u128>();
|
|
let mut receiver = receiver.into_future();
|
|
|
|
let (waker1, waker_handle1) = waker();
|
|
let mut context1 = Context::from_waker(&waker1);
|
|
|
|
assert_eq!(Pin::new(&mut receiver).poll(&mut context1), Poll::Pending);
|
|
assert_eq!(waker_handle1.clone_count(), 1);
|
|
assert_eq!(waker_handle1.drop_count(), 0);
|
|
assert_eq!(waker_handle1.wake_count(), 0);
|
|
|
|
let (waker2, waker_handle2) = waker();
|
|
let mut context2 = Context::from_waker(&waker2);
|
|
|
|
assert_eq!(Pin::new(&mut receiver).poll(&mut context2), Poll::Pending);
|
|
assert_eq!(waker_handle1.clone_count(), 1);
|
|
assert_eq!(waker_handle1.drop_count(), 1);
|
|
assert_eq!(waker_handle1.wake_count(), 0);
|
|
|
|
assert_eq!(waker_handle2.clone_count(), 1);
|
|
assert_eq!(waker_handle2.drop_count(), 0);
|
|
assert_eq!(waker_handle2.wake_count(), 0);
|
|
|
|
// Sending should cause the waker from the latest poll to be woken up
|
|
sender.send(1234).unwrap();
|
|
assert_eq!(waker_handle1.clone_count(), 1);
|
|
assert_eq!(waker_handle1.drop_count(), 1);
|
|
assert_eq!(waker_handle1.wake_count(), 0);
|
|
|
|
assert_eq!(waker_handle2.clone_count(), 1);
|
|
assert_eq!(waker_handle2.drop_count(), 1);
|
|
assert_eq!(waker_handle2.wake_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn poll_with_different_wakers_across_threads() {
|
|
let (sender, receiver) = oneshot::channel::<u128>();
|
|
let mut receiver = receiver.into_future();
|
|
|
|
let (waker1, waker_handle1) = waker();
|
|
let mut context1 = Context::from_waker(&waker1);
|
|
|
|
assert_eq!(Pin::new(&mut receiver).poll(&mut context1), Poll::Pending);
|
|
assert_eq!(waker_handle1.clone_count(), 1);
|
|
assert_eq!(waker_handle1.drop_count(), 0);
|
|
assert_eq!(waker_handle1.wake_count(), 0);
|
|
|
|
let receiver_thread = spawn_named("receiver", move || {
|
|
let (waker2, waker_handle2) = waker();
|
|
let mut context2 = Context::from_waker(&waker2);
|
|
|
|
assert_eq!(Pin::new(&mut receiver).poll(&mut context2), Poll::Pending);
|
|
assert_eq!(waker_handle2.clone_count(), 1);
|
|
assert_eq!(waker_handle2.drop_count(), 0);
|
|
assert_eq!(waker_handle2.wake_count(), 0);
|
|
|
|
drop(receiver);
|
|
assert_eq!(waker_handle2.drop_count(), 1);
|
|
});
|
|
|
|
receiver_thread.join().unwrap();
|
|
assert_eq!(waker_handle1.drop_count(), 1);
|
|
assert!(sender.is_closed());
|
|
}
|
|
|
|
#[test]
|
|
fn drop_pending_receiver_closes_channel_and_drops_waker() {
|
|
let (sender, receiver) = oneshot::channel::<u128>();
|
|
let mut receiver = receiver.into_future();
|
|
|
|
let (waker, waker_handle) = waker();
|
|
let mut context = Context::from_waker(&waker);
|
|
|
|
assert_eq!(Pin::new(&mut receiver).poll(&mut context), Poll::Pending);
|
|
assert_eq!(waker_handle.clone_count(), 1);
|
|
assert_eq!(waker_handle.drop_count(), 0);
|
|
assert_eq!(waker_handle.wake_count(), 0);
|
|
|
|
drop(receiver);
|
|
assert_eq!(waker_handle.drop_count(), 1);
|
|
assert_eq!(waker_handle.wake_count(), 0);
|
|
assert!(sender.is_closed());
|
|
|
|
let error = sender.send(1234).unwrap_err();
|
|
assert_eq!(*error.as_inner(), 1234);
|
|
}
|
|
|
|
#[test]
|
|
fn poll_then_drop_receiver_during_send() {
|
|
let (sender, receiver) = oneshot::channel::<u128>();
|
|
let mut receiver = receiver.into_future();
|
|
|
|
let (waker, _waker_handle) = waker();
|
|
let mut context = Context::from_waker(&waker);
|
|
|
|
// Put the channel into the receiving state
|
|
assert_eq!(Pin::new(&mut receiver).poll(&mut context), Poll::Pending);
|
|
|
|
// Spawn a separate thread that sends in parallel
|
|
let t = std::thread::spawn(move || {
|
|
let _ = sender.send(1234);
|
|
});
|
|
|
|
// Drop the receiver.
|
|
drop(receiver);
|
|
|
|
// The send operation should also not have panicked
|
|
t.join().unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn dropping_sender_disconnects_async_receiver() {
|
|
let (sender, receiver) = oneshot::channel::<()>();
|
|
assert!(!sender.is_closed());
|
|
assert!(!receiver.is_closed());
|
|
drop(sender);
|
|
assert!(receiver.is_closed());
|
|
}
|
|
|
|
#[test]
|
|
fn async_receiver_has_message() {
|
|
let (sender, receiver) = oneshot::channel();
|
|
assert!(!receiver.has_message());
|
|
assert!(sender.send(19i128).is_ok());
|
|
assert!(receiver.has_message());
|
|
}
|
|
|
|
#[test]
|
|
fn concurrent_send_and_try_recv_to_completion() {
|
|
let (sender, receiver) = oneshot::channel::<i32>();
|
|
|
|
let receiver_thread = spawn_named("receiver", move || {
|
|
spin_until("message from sender", || match receiver.try_recv() {
|
|
Ok(999) => Some(()),
|
|
Ok(value) => panic!("unexpected value: {value}"),
|
|
Err(TryRecvError::Empty) => None,
|
|
Err(TryRecvError::Disconnected) => panic!("unexpected disconnect"),
|
|
});
|
|
});
|
|
|
|
let sender_thread = spawn_named("sender", move || {
|
|
sender.send(999).unwrap();
|
|
});
|
|
|
|
receiver_thread.join().unwrap();
|
|
sender_thread.join().unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn concurrent_drop_sender_and_try_recv_to_completion() {
|
|
let (sender, receiver) = oneshot::channel::<i32>();
|
|
|
|
let receiver_thread = spawn_named("receiver", move || {
|
|
spin_until("sender disconnect", || match receiver.try_recv() {
|
|
Ok(value) => panic!("unexpected value: {value}"),
|
|
Err(TryRecvError::Empty) => None,
|
|
Err(TryRecvError::Disconnected) => Some(()),
|
|
});
|
|
});
|
|
|
|
let sender_thread = spawn_named("sender", move || {
|
|
drop(sender);
|
|
});
|
|
|
|
receiver_thread.join().unwrap();
|
|
sender_thread.join().unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn concurrent_send_and_poll_to_completion() {
|
|
let (sender, receiver) = oneshot::channel::<i32>();
|
|
|
|
let receiver_thread = spawn_named("receiver", move || {
|
|
let mut receiver = receiver.into_future();
|
|
let (waker, _waker_handle) = waker();
|
|
let mut context = Context::from_waker(&waker);
|
|
|
|
spin_until("poll ready with message", || {
|
|
match Pin::new(&mut receiver).poll(&mut context) {
|
|
Poll::Ready(Ok(999)) => Some(()),
|
|
Poll::Ready(result) => panic!("unexpected result: {result:?}"),
|
|
Poll::Pending => None,
|
|
}
|
|
});
|
|
});
|
|
|
|
let sender_thread = spawn_named("sender", move || {
|
|
sender.send(999).unwrap();
|
|
});
|
|
|
|
receiver_thread.join().unwrap();
|
|
sender_thread.join().unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn concurrent_drop_sender_and_poll_to_completion() {
|
|
let (sender, receiver) = oneshot::channel::<i32>();
|
|
|
|
let receiver_thread = spawn_named("receiver", move || {
|
|
let mut receiver = receiver.into_future();
|
|
let (waker, _waker_handle) = waker();
|
|
let mut context = Context::from_waker(&waker);
|
|
|
|
spin_until("poll ready with disconnect", || {
|
|
match Pin::new(&mut receiver).poll(&mut context) {
|
|
Poll::Ready(Err(oneshot::RecvError::Disconnected)) => Some(()),
|
|
Poll::Ready(result) => panic!("unexpected result: {result:?}"),
|
|
Poll::Pending => None,
|
|
}
|
|
});
|
|
});
|
|
|
|
let sender_thread = spawn_named("sender", move || {
|
|
drop(sender);
|
|
});
|
|
|
|
receiver_thread.join().unwrap();
|
|
sender_thread.join().unwrap();
|
|
}
|
|
|
|
fn spawn_named<F>(name: &str, f: F) -> std::thread::JoinHandle<()>
|
|
where
|
|
F: FnOnce() + Send + 'static,
|
|
{
|
|
std::thread::Builder::new()
|
|
.name(name.to_string())
|
|
.spawn(f)
|
|
.unwrap()
|
|
}
|
|
|
|
fn spin_until<F>(label: &str, mut f: F)
|
|
where
|
|
F: FnMut() -> Option<()>,
|
|
{
|
|
let deadline = Instant::now() + Duration::from_secs(5);
|
|
let mut spins = 0usize;
|
|
|
|
loop {
|
|
if f().is_some() {
|
|
break;
|
|
}
|
|
|
|
assert!(Instant::now() < deadline, "timed out waiting for {label}");
|
|
|
|
if spins % 64 == 0 {
|
|
std::thread::yield_now();
|
|
} else {
|
|
spin_loop();
|
|
}
|
|
|
|
spins += 1;
|
|
}
|
|
}
|