From 2fcc6c2cb0e38d5e36b40d9a2fdb9871e2825fa7 Mon Sep 17 00:00:00 2001 From: Thomas de Zeeuw Date: Fri, 18 Nov 2022 11:59:05 +0100 Subject: [PATCH] Box Futures in #[tokio::test] This reduces the amount of copies of the Runtime::block_on and related functions the compiler has to generate and LLVM process. We've seen it reduce the compilation time of our tests (some 1900 of them) from 40s down to 12s, with no impact on the runtime of tests. Below is an output of llvm-lines for our tests. Before: Lines Copies Function name ----- ------ ------------- 8954414 156577 (TOTAL) 984626 (11.0%, 11.0%) 9289 (5.9%, 5.9%) std::thread::local::LocalKey::try_with 648093 (7.2%, 18.2%) 1857 (1.2%, 7.1%) tokio::runtime::scheduler::current_thread::CoreGuard::block_on::{{closure}} 557100 (6.2%, 24.5%) 3714 (2.4%, 9.5%) tokio::park::thread::CachedParkThread::block_on 551679 (6.2%, 30.6%) 7430 (4.7%, 14.2%) tokio::coop::with_budget::{{closure}} 514389 (5.7%, 36.4%) 3714 (2.4%, 16.6%) tokio::runtime::scheduler::current_thread::Context::enter 326832 (3.6%, 40.0%) 1857 (1.2%, 17.8%) tokio::runtime::scheduler::current_thread::CurrentThread::block_on 291549 (3.3%, 43.3%) 1857 (1.2%, 19.0%) tokio::runtime::scheduler::current_thread::CoreGuard::enter 261907 (2.9%, 46.2%) 7430 (4.7%, 23.7%) tokio::coop::budget 189468 (2.1%, 48.3%) 7430 (4.7%, 28.5%) tokio::coop::with_budget 137418 (1.5%, 49.8%) 3714 (2.4%, 30.8%) tokio::runtime::enter::Enter::block_on 126276 (1.4%, 51.3%) 1857 (1.2%, 32.0%) tokio::runtime::Runtime::block_on 124419 (1.4%, 52.6%) 1857 (1.2%, 33.2%) tokio::macros::scoped_tls::ScopedKey::set 118897 (1.3%, 54.0%) 3715 (2.4%, 35.6%) core::option::Option::or_else 111420 (1.2%, 55.2%) 1857 (1.2%, 36.8%) tokio::runtime::scheduler::current_thread::CurrentThread::block_on::{{closure}} 109408 (1.2%, 56.4%) 2105 (1.3%, 38.1%) as core::future::future::Future>::poll 105893 (1.2%, 57.6%) 9289 (5.9%, 44.0%) std::thread::local::LocalKey::with 96564 (1.1%, 58.7%) 1857 (1.2%, 45.2%) tokio::runtime::scheduler::current_thread::Context::run_task 90993 (1.0%, 59.7%) 7428 (4.7%, 50.0%) tokio::runtime::scheduler::current_thread::CoreGuard::block_on::{{closure}}::{{closure}} 90515 (1.0%, 60.7%) 2105 (1.3%, 51.3%) core::pin::Pin<&mut T>::map_unchecked_mut 89136 (1.0%, 61.7%) 1857 (1.2%, 52.5%) tokio::runtime::scheduler::multi_thread::MultiThread::block_on After: Lines Copies Function name ----- ------ ------------- 3188618 41634 (TOTAL) 109408 (3.4%, 3.4%) 2105 (5.1%, 5.1%) as core::future::future::Future>::poll 90515 (2.8%, 6.3%) 2105 (5.1%, 10.1%) core::pin::Pin<&mut T>::map_unchecked_mut 56220 (1.8%, 8.0%) 1874 (4.5%, 14.6%) alloc::boxed::Box::pin 48333 (1.5%, 9.5%) 2179 (5.2%, 19.8%) core::ops::function::FnOnce::call_once 28587 (0.9%, 10.4%) 1 (0.0%, 19.8%) XXXXXXXXXXXXXXXXXXX 18730 (0.6%, 11.0%) 1873 (4.5%, 24.3%) alloc::boxed::Box::into_pin 16190 (0.5%, 11.5%) 2 (0.0%, 24.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 15870 (0.5%, 12.0%) 2 (0.0%, 24.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 15250 (0.5%, 12.5%) 1 (0.0%, 24.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXX 12801 (0.4%, 12.9%) 2 (0.0%, 24.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 12801 (0.4%, 13.3%) 2 (0.0%, 24.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 12630 (0.4%, 13.7%) 2105 (5.1%, 29.4%) as core::future::future::Future>::poll::{{closure}} 12613 (0.4%, 14.1%) 2 (0.0%, 29.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 12613 (0.4%, 14.5%) 2 (0.0%, 29.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 12613 (0.4%, 14.9%) 2 (0.0%, 29.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 12613 (0.4%, 15.3%) 2 (0.0%, 29.4%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 11395 (0.4%, 15.7%) 96 (0.2%, 29.7%) alloc::alloc::box_free 11364 (0.4%, 16.0%) 1891 (4.5%, 34.2%) >::into 11238 (0.4%, 16.4%) 1873 (4.5%, 38.7%) alloc::boxed::> for core::pin::Pin>>::from 10735 (0.3%, 16.7%) 2 (0.0%, 38.7%) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX Note that I have replaced our test functions with XXX. As you can clearly see they're not in the top 20 in the before output, while they're in the after oput. Further note that the amount of copies have been reduced from 156577 to 41634. --- tokio-macros/src/entry.rs | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs index 19cb0b1cc..1aa03f34f 100644 --- a/tokio-macros/src/entry.rs +++ b/tokio-macros/src/entry.rs @@ -393,9 +393,37 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To .block_on(body); } }; + + // For test functions wrap the body in a `Box` to reduce the amount + // `Runtime::block_on` (and related functions) copies we generate during + // compilation due to the generic parameter `F` (the future to block on). + // This could have an impact on performance, but because it's only for + // testing and the runtime already boxes the future to run it, it's unlikely + // to be very large. + // + // We don't do this for the main function as it should only be used once so + // there will be no benefit. + let body = if is_test { + let output_type = match &input.sig.output { + // For functions with no return value syn doesn't print anything, + // but that doesn't work as `Output` for our boxed `Future`, so + // default to `()` (the same type as the function output). + syn::ReturnType::Default => quote! { () }, + syn::ReturnType::Type(_, ret_type) => quote! { #ret_type }, + }; + quote! { + let body: ::std::pin::Pin<::std::boxed::Box>> = + ::std::boxed::Box::pin(async #body); + } + } else { + quote! { + let body = async #body; + } + }; + input.block = syn::parse2(quote! { { - let body = async #body; + #body #block_expr } })