From 3fe7cfc32656518d5e6262e580f0a16cd2412dd7 Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Wed, 13 Nov 2019 17:14:50 -0800 Subject: [PATCH] Remove some stack frames from `.async` calls 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: 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: 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. --- src/libstd/future.rs | 45 ++++++++++++-------------------------------- 1 file changed, 12 insertions(+), 33 deletions(-) diff --git a/src/libstd/future.rs b/src/libstd/future.rs index c65f71fb1a4..6de3f1d545b 100644 --- a/src/libstd/future.rs +++ b/src/libstd/future.rs @@ -40,10 +40,11 @@ impl> Future for GenFuture { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { // 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(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: F) -> R +/// Polls a future in the current thread-local task waker. +pub fn poll_with_tls_context(f: Pin<&mut F>) -> Poll 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: 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: Pin<&mut F>) -> Poll -where - F: Future -{ - get_task_context(|cx| F::poll(f, cx)) + unsafe { F::poll(f, cx_ptr.as_mut()) } }