From 966ccd5d5306adf6b6c39721331c2a3c32be6fa8 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Sun, 3 Nov 2019 14:10:14 -0800 Subject: test: unify MockTask and task::spawn (#1728) Delete `MockTask` in favor of `task::spawn`. Both are functionally equivalent. --- tokio-test/Cargo.toml | 1 - tokio-test/src/macros.rs | 55 ++++++------------- tokio-test/src/task.rs | 137 ++++++++++++++++++++++++++--------------------- 3 files changed, 91 insertions(+), 102 deletions(-) (limited to 'tokio-test') diff --git a/tokio-test/Cargo.toml b/tokio-test/Cargo.toml index 95dc5d7f..6d7ebc34 100644 --- a/tokio-test/Cargo.toml +++ b/tokio-test/Cargo.toml @@ -24,7 +24,6 @@ tokio = { version = "=0.2.0-alpha.6", path = "../tokio" } bytes = "0.4" futures-core-preview = "=0.3.0-alpha.19" -pin-convert = "0.1.0" [dev-dependencies] futures-util-preview = "=0.3.0-alpha.19" diff --git a/tokio-test/src/macros.rs b/tokio-test/src/macros.rs index b39b4a02..dbe2280f 100644 --- a/tokio-test/src/macros.rs +++ b/tokio-test/src/macros.rs @@ -13,16 +13,11 @@ /// # Examples /// /// ``` -/// use std::future::Future; -/// use futures_util::{future, pin_mut}; +/// use futures_util::future; /// use tokio_test::{assert_ready, task}; /// -/// task::mock(|cx| { -/// let fut = future::ready(()); -/// -/// pin_mut!(fut); -/// assert_ready!(fut.poll(cx)); -/// }) +/// let mut fut = task::spawn(future::ready(())); +/// assert_ready!(fut.poll()); /// ``` #[macro_export] macro_rules! assert_ready { @@ -57,16 +52,11 @@ macro_rules! assert_ready { /// # Examples /// /// ``` -/// use std::future::Future; -/// use futures_util::{future, pin_mut}; +/// use futures_util::future; /// use tokio_test::{assert_ready_ok, task}; /// -/// task::mock(|cx| { -/// let fut = future::ok::<_, ()>(()); -/// -/// pin_mut!(fut); -/// assert_ready_ok!(fut.poll(cx)); -/// }) +/// let mut fut = task::spawn(future::ok::<_, ()>(())); +/// assert_ready_ok!(fut.poll()); /// ``` #[macro_export] macro_rules! assert_ready_ok { @@ -95,16 +85,11 @@ macro_rules! assert_ready_ok { /// # Examples /// /// ``` -/// use std::future::Future; -/// use futures_util::{future, pin_mut}; +/// use futures_util::future; /// use tokio_test::{assert_ready_err, task}; /// -/// task::mock(|cx| { -/// let fut = future::err::<(), _>(()); -/// -/// pin_mut!(fut); -/// assert_ready_err!(fut.poll(cx)); -/// }) +/// let mut fut = task::spawn(future::err::<(), _>(())); +/// assert_ready_err!(fut.poll()); /// ``` #[macro_export] macro_rules! assert_ready_err { @@ -133,16 +118,11 @@ macro_rules! assert_ready_err { /// # Examples /// /// ``` -/// use std::future::Future; -/// use futures_util::{future, pin_mut}; +/// use futures_util::future; /// use tokio_test::{assert_pending, task}; /// -/// task::mock(|cx| { -/// let fut = future::pending::<()>(); -/// -/// pin_mut!(fut); -/// assert_pending!(fut.poll(cx)); -/// }) +/// let mut fut = task::spawn(future::pending::<()>()); +/// assert_pending!(fut.poll()); /// ``` #[macro_export] macro_rules! assert_pending { @@ -177,16 +157,11 @@ macro_rules! assert_pending { /// # Examples /// /// ``` -/// use std::future::Future; -/// use futures_util::{future, pin_mut}; +/// use futures_util::future; /// use tokio_test::{assert_ready_eq, task}; /// -/// task::mock(|cx| { -/// let fut = future::ready(42); -/// -/// pin_mut!(fut); -/// assert_ready_eq!(fut.poll(cx), 42); -/// }) +/// let mut fut = task::spawn(future::ready(42)); +/// assert_ready_eq!(fut.poll(), 42); /// ``` #[macro_export] macro_rules! assert_ready_eq { diff --git a/tokio-test/src/task.rs b/tokio-test/src/task.rs index 5468749a..c21e31a5 100644 --- a/tokio-test/src/task.rs +++ b/tokio-test/src/task.rs @@ -1,44 +1,19 @@ //! Futures task based helpers -use tokio::executor::enter; - -use pin_convert::AsPinMut; +use futures_core::Stream; use std::future::Future; use std::mem; +use std::ops; use std::pin::Pin; use std::sync::{Arc, Condvar, Mutex}; use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; -/// Run the provided closure in a `MockTask` context. -/// -/// # Examples -/// -/// ``` -/// use std::future::Future; -/// use futures_util::{future, pin_mut}; -/// use tokio_test::task; -/// -/// task::mock(|cx| { -/// let fut = future::ready(()); -/// -/// pin_mut!(fut); -/// assert!(fut.poll(cx).is_ready()); -/// }) -/// ``` -pub fn mock(f: F) -> R -where - F: Fn(&mut Context<'_>) -> R, -{ - let mut task = MockTask::new(); - task.enter(|cx| f(cx)) -} - -/// Mock task -/// -/// A mock task is able to intercept and track wake notifications. -#[derive(Debug, Clone)] -pub struct MockTask { - waker: Arc, +/// TOOD: dox +pub fn spawn(task: T) -> Spawn { + Spawn { + task: MockTask::new(), + future: Box::pin(task), + } } /// Future spawned on a mock task @@ -48,12 +23,12 @@ pub struct Spawn { future: Pin>, } -/// TOOD: dox -pub fn spawn(task: T) -> Spawn { - Spawn { - task: MockTask::new(), - future: Box::pin(task), - } +/// Mock task +/// +/// A mock task is able to intercept and track wake notifications. +#[derive(Debug, Clone)] +struct MockTask { + waker: Arc, } #[derive(Debug)] @@ -66,11 +41,23 @@ const IDLE: usize = 0; const WAKE: usize = 1; const SLEEP: usize = 2; -impl Spawn { - /// Poll a future - pub fn poll(&mut self) -> Poll { - let fut = self.future.as_mut(); - self.task.enter(|cx| fut.poll(cx)) +impl Spawn { + /// Consume `self` returning the inner value + pub fn into_inner(mut self) -> T + where + T: Unpin, + { + drop(self.task); + + // Pin::into_inner is unstable, so we work around it + // + // Safety: `T` is bound by `Unpin`. + unsafe { + let ptr = Pin::get_mut(self.future.as_mut()) as *mut T; + let future = Box::from_raw(ptr); + mem::forget(self.future); + *future + } } /// Returns `true` if the inner future has received a wake notification @@ -85,35 +72,63 @@ impl Spawn { pub fn waker_ref_count(&self) -> usize { self.task.waker_ref_count() } + + /// Enter the task context + pub fn enter(&mut self, f: F) -> R + where + F: FnOnce(&mut Context<'_>, Pin<&mut T>) -> R, + { + let fut = self.future.as_mut(); + self.task.enter(|cx| f(cx, fut)) + } +} + +impl ops::Deref for Spawn { + type Target = T; + + fn deref(&self) -> &T { + &self.future + } +} + +impl ops::DerefMut for Spawn { + fn deref_mut(&mut self) -> &mut T { + &mut self.future + } +} + +impl Spawn { + /// Poll a future + pub fn poll(&mut self) -> Poll { + let fut = self.future.as_mut(); + self.task.enter(|cx| fut.poll(cx)) + } +} + +impl Spawn { + /// Poll a stream + pub fn poll_next(&mut self) -> Poll> { + let stream = self.future.as_mut(); + self.task.enter(|cx| stream.poll_next(cx)) + } } impl MockTask { /// Create a new mock task - pub fn new() -> Self { + fn new() -> Self { MockTask { waker: Arc::new(ThreadWaker::new()), } } - /// Poll a future - pub fn poll(&mut self, mut fut: T) -> Poll - where - T: AsPinMut, - F: Future, - { - self.enter(|cx| fut.as_pin_mut().poll(cx)) - } - /// Run a closure from the context of the task. /// /// Any wake notifications resulting from the execution of the closure are /// tracked. - pub fn enter(&mut self, f: F) -> R + fn enter(&mut self, f: F) -> R where F: FnOnce(&mut Context<'_>) -> R, { - let _enter = enter().unwrap(); - self.waker.clear(); let waker = self.waker(); let mut cx = Context::from_waker(&waker); @@ -123,14 +138,14 @@ impl MockTask { /// Returns `true` if the inner future has received a wake notification /// since the last call to `enter`. - pub fn is_woken(&self) -> bool { + fn is_woken(&self) -> bool { self.waker.is_woken() } /// Returns the number of references to the task waker /// /// The task itself holds a reference. The return value will never be zero. - pub fn waker_ref_count(&self) -> usize { + fn waker_ref_count(&self) -> usize { Arc::strong_count(&self.waker) } @@ -193,7 +208,7 @@ impl ThreadWaker { } } -static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); +static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop_waker); unsafe fn to_raw(waker: Arc) -> RawWaker { RawWaker::new(Arc::into_raw(waker) as *const (), &VTABLE) @@ -225,6 +240,6 @@ unsafe fn wake_by_ref(raw: *const ()) { mem::forget(waker); } -unsafe fn drop(raw: *const ()) { +unsafe fn drop_waker(raw: *const ()) { let _ = from_raw(raw); } -- cgit v1.2.3