diff --git a/components/tikv_util/src/yatp_pool/future_pool.rs b/components/tikv_util/src/yatp_pool/future_pool.rs index 2deead30580..c65219ffbb1 100644 --- a/components/tikv_util/src/yatp_pool/future_pool.rs +++ b/components/tikv_util/src/yatp_pool/future_pool.rs @@ -13,6 +13,7 @@ use std::{ use fail::fail_point; use futures::channel::oneshot::{self, Canceled}; +use futures_util::future::FutureExt; use prometheus::{IntCounter, IntGauge}; use tracker::TrackedFuture; use yatp::{queue::Extras, task::future}; @@ -179,11 +180,13 @@ impl PoolInner { metrics_running_task_count.inc(); - let f = async move { - let _ = future.await; + // NB: Prefer FutureExt::map to async block, because an async block + // doubles memory usage. + // See https://github.com/rust-lang/rust/issues/59087 + let f = future.map(move |_| { metrics_handled_task_count.inc(); metrics_running_task_count.dec(); - }; + }); if let Some(extras) = extras { self.pool.spawn(future::TaskCell::new(f, extras)); @@ -208,12 +211,14 @@ impl PoolInner { let (tx, rx) = oneshot::channel(); metrics_running_task_count.inc(); - self.pool.spawn(async move { - let res = future.await; + // NB: Prefer FutureExt::map to async block, because an async block + // doubles memory usage. + // See https://github.com/rust-lang/rust/issues/59087 + self.pool.spawn(future.map(move |res| { metrics_handled_task_count.inc(); metrics_running_task_count.dec(); let _ = tx.send(res); - }); + })); Ok(rx) } } diff --git a/src/read_pool.rs b/src/read_pool.rs index 301ea648274..da251c78aa0 100644 --- a/src/read_pool.rs +++ b/src/read_pool.rs @@ -12,7 +12,10 @@ use std::{ }; use file_system::{set_io_type, IoType}; -use futures::{channel::oneshot, future::TryFutureExt}; +use futures::{ + channel::oneshot, + future::{FutureExt, TryFutureExt}, +}; use kvproto::{errorpb, kvrpcpb::CommandPri}; use online_config::{ConfigChange, ConfigManager, ConfigValue, Result as CfgResult}; use prometheus::{core::Metric, Histogram, IntCounter, IntGauge}; @@ -171,10 +174,9 @@ impl ReadPoolHandle { TaskCell::new( TrackedFuture::new(with_resource_limiter( ControlledFuture::new( - async move { - f.await; + f.map(move |_| { running_tasks.dec(); - }, + }), resource_ctl.clone(), group_name, ), @@ -184,10 +186,9 @@ impl ReadPoolHandle { ) } else { TaskCell::new( - TrackedFuture::new(async move { - f.await; + TrackedFuture::new(f.map(move |_| { running_tasks.dec(); - }), + })), extras, ) }; @@ -211,10 +212,9 @@ impl ReadPoolHandle { { let (tx, rx) = oneshot::channel::(); let res = self.spawn( - async move { - let res = f.await; + f.map(move |res| { let _ = tx.send(res); - }, + }), priority, task_id, metadata, diff --git a/src/storage/txn/sched_pool.rs b/src/storage/txn/sched_pool.rs index 8674a581c72..70d54c24cfa 100644 --- a/src/storage/txn/sched_pool.rs +++ b/src/storage/txn/sched_pool.rs @@ -131,13 +131,7 @@ impl PriorityQueue { extras.set_metadata(metadata.to_vec()); self.worker_pool.spawn_with_extras( with_resource_limiter( - ControlledFuture::new( - async move { - f.await; - }, - self.resource_ctl.clone(), - group_name, - ), + ControlledFuture::new(f, self.resource_ctl.clone(), group_name), resource_limiter, ), extras,