src/libstd/future.rs

   1 //! Asynchronous values.
   2
   3 use core::cell::Cell;
   4 use core::marker::Unpin;
   5 use core::pin::Pin;
   6 use core::option::Option;
   7 use core::ptr::NonNull;
   8 use core::task::{LocalWaker, Poll};
   9 use core::ops::{Drop, Generator, GeneratorState};
  10
  11 #[doc(inline)]
  12 pub use core::future::*;
  13
  14 /// Wrap a future in a generator.
  15 ///
  16 /// This function returns a `GenFuture` underneath, but hides it in `impl Trait` to give
  17 /// better error messages (`impl Future` rather than `GenFuture<[closure.....]>`).
  18 #[unstable(feature = "gen_future", issue = "50547")]
  19 pub fn from_generator<T: Generator<Yield = ()>>(x: T) -> impl Future<Output = T::Return> {
  20     GenFuture(x)
  21 }
  22
  23 /// A wrapper around generators used to implement `Future` for `async`/`await` code.
  24 #[unstable(feature = "gen_future", issue = "50547")]
  25 #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
  26 struct GenFuture<T: Generator<Yield = ()>>(T);
  27
  28 // We rely on the fact that async/await futures are immovable in order to create
  29 // self-referential borrows in the underlying generator.
  30 impl<T: Generator<Yield = ()>> !Unpin for GenFuture<T> {}
  31
  32 #[unstable(feature = "gen_future", issue = "50547")]
  33 impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
  34     type Output = T::Return;
  35     fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
  36         set_task_waker(lw, || match unsafe { Pin::get_unchecked_mut(self).0.resume() } {
  37             GeneratorState::Yielded(()) => Poll::Pending,
  38             GeneratorState::Complete(x) => Poll::Ready(x),
  39         })
  40     }
  41 }
  42
  43 thread_local! {
  44     static TLS_WAKER: Cell<Option<NonNull<LocalWaker>>> = Cell::new(None);
  45 }
  46
  47 struct SetOnDrop(Option<NonNull<LocalWaker>>);
  48
  49 impl Drop for SetOnDrop {
  50     fn drop(&mut self) {
  51         TLS_WAKER.with(|tls_waker| {
  52             tls_waker.set(self.0.take());
  53         });
  54     }
  55 }
  56
  57 #[unstable(feature = "gen_future", issue = "50547")]
  58 /// Sets the thread-local task context used by async/await futures.
  59 pub fn set_task_waker<F, R>(lw: &LocalWaker, f: F) -> R
  60 where
  61     F: FnOnce() -> R
  62 {
  63     let old_waker = TLS_WAKER.with(|tls_waker| {
  64         tls_waker.replace(Some(NonNull::from(lw)))
  65     });
  66     let _reset_waker = SetOnDrop(old_waker);
  67     f()
  68 }
  69
  70 #[unstable(feature = "gen_future", issue = "50547")]
  71 /// Retrieves the thread-local task waker used by async/await futures.
  72 ///
  73 /// This function acquires exclusive access to the task waker.
  74 ///
  75 /// Panics if no waker has been set or if the waker has already been
  76 /// retrieved by a surrounding call to get_task_waker.
  77 pub fn get_task_waker<F, R>(f: F) -> R
  78 where
  79     F: FnOnce(&LocalWaker) -> R
  80 {
  81     let waker_ptr = TLS_WAKER.with(|tls_waker| {
  82         // Clear the entry so that nested `get_task_waker` calls
  83         // will fail or set their own value.
  84         tls_waker.replace(None)
  85     });
  86     let _reset_waker = SetOnDrop(waker_ptr);
  87
  88     let waker_ptr = waker_ptr.expect(
  89         "TLS LocalWaker not set. This is a rustc bug. \
  90         Please file an issue on https://github.com/rust-lang/rust.");
  91     unsafe { f(waker_ptr.as_ref()) }
  92 }
  93
  94 #[unstable(feature = "gen_future", issue = "50547")]
  95 /// Polls a future in the current thread-local task waker.
  96 pub fn poll_with_tls_waker<F>(f: Pin<&mut F>) -> Poll<F::Output>
  97 where
  98     F: Future
  99 {
 100     get_task_waker(|lw| F::poll(f, lw))
 101 }