3fe7cfc326
The `Context` argument is currently smuggled through TLS for
async-generated futures. The current infrastructure is closure-based,
and results in an extra 6 stack frames when .awaiting an async-generated
future!
```
12: foo::async_b::{{closure}}
at src/main.rs:10
13: <std::future::GenFuture<T> as core::future::future::Future>::poll::{{closure}}
at /rustc/4560ea788cb760f0a34127156c78e2552949f734/src/libstd/future.rs:43
14: std::future::set_task_context
at /rustc/4560ea788cb760f0a34127156c78e2552949f734/src/libstd/future.rs:79
15: <std::future::GenFuture<T> as core::future::future::Future>::poll
at /rustc/4560ea788cb760f0a34127156c78e2552949f734/src/libstd/future.rs:43
16: std::future::poll_with_tls_context::{{closure}}
at /rustc/4560ea788cb760f0a34127156c78e2552949f734/src/libstd/future.rs:121
17: std::future::get_task_context
at /rustc/4560ea788cb760f0a34127156c78e2552949f734/src/libstd/future.rs:111
18: std::future::poll_with_tls_context
at /rustc/4560ea788cb760f0a34127156c78e2552949f734/src/libstd/future.rs:121
19: foo::async_a::{{closure}}
at src/main.rs:6
```
While the long (medium?) term solution is to remove the use of TLS
entirely, we can improve things a bit in the meantime. In particular,
this commit does 2 things:
1. `get_task_context` has been inlined into `poll_with_tls_context`,
removing 2 frames (16 and 17 above).
2. `set_task_context` now returns a guard type that resets the TLS
rather than taking a closure, removing 2 frames (13 and 14 above).
We can also remove frame 18 by removing `poll_with_tls_context` in favor
of a `get_task_context` function which returns a guard, but that
requires adjusting the code generated for .await, so I've left that off
for now.
101 lines
3.4 KiB
Rust
101 lines
3.4 KiB
Rust
//! Asynchronous values.
|
|
|
|
use core::cell::Cell;
|
|
use core::marker::Unpin;
|
|
use core::pin::Pin;
|
|
use core::option::Option;
|
|
use core::ptr::NonNull;
|
|
use core::task::{Context, Poll};
|
|
use core::ops::{Drop, Generator, GeneratorState};
|
|
|
|
#[doc(inline)]
|
|
#[stable(feature = "futures_api", since = "1.36.0")]
|
|
pub use core::future::*;
|
|
|
|
/// Wrap a generator in a future.
|
|
///
|
|
/// This function returns a `GenFuture` underneath, but hides it in `impl Trait` to give
|
|
/// better error messages (`impl Future` rather than `GenFuture<[closure.....]>`).
|
|
#[doc(hidden)]
|
|
#[unstable(feature = "gen_future", issue = "50547")]
|
|
pub fn from_generator<T: Generator<Yield = ()>>(x: T) -> impl Future<Output = T::Return> {
|
|
GenFuture(x)
|
|
}
|
|
|
|
/// A wrapper around generators used to implement `Future` for `async`/`await` code.
|
|
#[doc(hidden)]
|
|
#[unstable(feature = "gen_future", issue = "50547")]
|
|
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
|
|
#[cfg_attr(not(test), rustc_diagnostic_item = "gen_future")]
|
|
struct GenFuture<T: Generator<Yield = ()>>(T);
|
|
|
|
// We rely on the fact that async/await futures are immovable in order to create
|
|
// self-referential borrows in the underlying generator.
|
|
impl<T: Generator<Yield = ()>> !Unpin for GenFuture<T> {}
|
|
|
|
#[doc(hidden)]
|
|
#[unstable(feature = "gen_future", issue = "50547")]
|
|
impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
|
|
type Output = T::Return;
|
|
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
// Safe because we're !Unpin + !Drop mapping to a ?Unpin value
|
|
let gen = unsafe { Pin::map_unchecked_mut(self, |s| &mut s.0) };
|
|
let _guard = unsafe { set_task_context(cx) };
|
|
match gen.resume() {
|
|
GeneratorState::Yielded(()) => Poll::Pending,
|
|
GeneratorState::Complete(x) => Poll::Ready(x),
|
|
}
|
|
}
|
|
}
|
|
|
|
thread_local! {
|
|
static TLS_CX: Cell<Option<NonNull<Context<'static>>>> = Cell::new(None);
|
|
}
|
|
|
|
struct SetOnDrop(Option<NonNull<Context<'static>>>);
|
|
|
|
impl Drop for SetOnDrop {
|
|
fn drop(&mut self) {
|
|
TLS_CX.with(|tls_cx| {
|
|
tls_cx.set(self.0.take());
|
|
});
|
|
}
|
|
}
|
|
|
|
// Safety: the returned guard must drop before `cx` is dropped and before
|
|
// any previous guard is dropped.
|
|
unsafe fn set_task_context(cx: &mut Context<'_>) -> SetOnDrop {
|
|
// transmute the context's lifetime to 'static so we can store it.
|
|
let cx = core::mem::transmute::<&mut Context<'_>, &mut Context<'static>>(cx);
|
|
let old_cx = TLS_CX.with(|tls_cx| {
|
|
tls_cx.replace(Some(NonNull::from(cx)))
|
|
});
|
|
SetOnDrop(old_cx)
|
|
}
|
|
|
|
#[doc(hidden)]
|
|
#[unstable(feature = "gen_future", issue = "50547")]
|
|
/// Polls a future in the current thread-local task waker.
|
|
pub fn poll_with_tls_context<F>(f: Pin<&mut F>) -> Poll<F::Output>
|
|
where
|
|
F: Future
|
|
{
|
|
let cx_ptr = TLS_CX.with(|tls_cx| {
|
|
// Clear the entry so that nested `get_task_waker` calls
|
|
// will fail or set their own value.
|
|
tls_cx.replace(None)
|
|
});
|
|
let _reset = SetOnDrop(cx_ptr);
|
|
|
|
let mut cx_ptr = cx_ptr.expect(
|
|
"TLS Context not set. This is a rustc bug. \
|
|
Please file an issue on https://github.com/rust-lang/rust.");
|
|
|
|
// Safety: we've ensured exclusive access to the context by
|
|
// removing the pointer from TLS, only to be replaced once
|
|
// we're done with it.
|
|
//
|
|
// The pointer that was inserted came from an `&mut Context<'_>`,
|
|
// so it is safe to treat as mutable.
|
|
unsafe { F::poll(f, cx_ptr.as_mut()) }
|
|
}
|