mikros/rust
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Remove some stack frames from .async calls

Browse Source 
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.
This commit is contained in:
Steven Fackler 2019-11-13 17:14:50 -08:00
parent bc0e288ad0
commit 3fe7cfc326
1 changed files with 12 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
src/libstd/future.rs
View File

@ -40,10 +40,11 @@ impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
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) };
set_task_context(cx, || match gen.resume() {
let _guard = unsafe { set_task_context(cx) };
match gen.resume() {
GeneratorState::Yielded(()) => Poll::Pending,
GeneratorState::Complete(x) => Poll::Ready(x),
})
}
}
}
@ -61,35 +62,23 @@ fn drop(&mut self) {
}
}
#[doc(hidden)]
#[unstable(feature = "gen_future", issue = "50547")]
/// Sets the thread-local task context used by async/await futures.
pub fn set_task_context<F, R>(cx: &mut Context<'_>, f: F) -> R
where
F: FnOnce() -> R
{
// 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 = unsafe {
core::mem::transmute::<&mut Context<'_>, &mut Context<'static>>(cx)
};
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)))
});
let _reset = SetOnDrop(old_cx);
f()
SetOnDrop(old_cx)
}
#[doc(hidden)]
#[unstable(feature = "gen_future", issue = "50547")]
/// Retrieves the thread-local task context used by async/await futures.
///
/// This function acquires exclusive access to the task context.
///
/// Panics if no context has been set or if the context has already been
/// retrieved by a surrounding call to get_task_context.
pub fn get_task_context<F, R>(f: F) -> R
/// 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: FnOnce(&mut Context<'_>) -> R
F: Future
{
let cx_ptr = TLS_CX.with(|tls_cx| {
// Clear the entry so that nested `get_task_waker` calls
@ -108,15 +97,5 @@ pub fn get_task_context<F, R>(f: F) -> R
//
// The pointer that was inserted came from an `&mut Context<'_>`,
// so it is safe to treat as mutable.
unsafe { f(cx_ptr.as_mut()) }
}
#[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
{
get_task_context(|cx| F::poll(f, cx))
unsafe { F::poll(f, cx_ptr.as_mut()) }
}