[vfs] Reduce the size of spawned futures

The size of the wrapping future created in ExecutionScope::spawn was
twice the size of the task passed into it. This was because the wrapping
future was capturing the task in an upvar and then moving it to a local
when polling it requiring space for 2 copies of the task.
rust-lang/rust#108906

Writing our own future prevents duplicating the space requirements. The
task spawned in VolumesDirectory::add_directory_entry is now 15312
bytes, down from 30640 bytes.

Bug: b/393365596
Change-Id: Ib90a3ad99d42bd937effcac6bfa39d780ef188f4
Reviewed-on: https://fuchsia-review.googlesource.com/c/fuchsia/+/1198686
Commit-Queue: Chris Drouillard <[email protected]>
Reviewed-by: Chris Suter <[email protected]>

Author: Chris Drouillard

src/storage/lib/vfs/rust/src/execution_scope.rs

@@ -21,6 +21,7 @@ use crate::token_registry::TokenRegistry;
 use fuchsia_async::{JoinHandle, Scope, Task};
 use futures::task::{self, Poll};
 use futures::Future;
+use pin_project::pin_project;
 use std::future::{pending, poll_fn};
 use std::pin::pin;
 use std::sync::{Arc, Mutex, OnceLock};
@@ -102,26 +103,7 @@ impl ExecutionScope {
     /// This way `ExecutionScope` can actually also implement [`futures::task::Spawn`] - it just was
     /// not necessary for now.
     pub fn spawn(&self, task: impl Future<Output = ()> + Send + 'static) -> JoinHandle<()> {
-        let executor = self.executor.clone();
-        self.executor.scope().spawn(async move {
-            let mut task = std::pin::pin!(task);
-            poll_fn(|cx| {
-                let shutdown_state = executor.inner.lock().unwrap().shutdown_state;
-                match task.as_mut().poll(cx) {
-                    Poll::Ready(()) => Poll::Ready(()),
-                    Poll::Pending => match shutdown_state {
-                        ShutdownState::Active => Poll::Pending,
-                        ShutdownState::Shutdown
-                            if executor.inner.lock().unwrap().active_count > 0 =>
-                        {
-                            Poll::Pending
-                        }
-                        _ => Poll::Ready(()),
-                    },
-                }
-            })
-            .await;
-        })
+        self.executor.scope().spawn(FutureWithShutdown { executor: self.executor.clone(), task })
     }
 
     pub fn token_registry(&self) -> &TokenRegistry {
@@ -305,6 +287,33 @@ pub async fn yield_to_executor() {
     .await;
 }
 
+/// A future that wraps another future and watches for the shutdown signal.
+#[pin_project]
+struct FutureWithShutdown<Task: Future<Output = ()> + Send + 'static> {
+    executor: Arc<Executor>,
+    #[pin]
+    task: Task,
+}
+
+impl<Task: Future<Output = ()> + Send + 'static> Future for FutureWithShutdown<Task> {
+    type Output = ();
+
+    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let shutdown_state = this.executor.inner.lock().unwrap().shutdown_state;
+        match this.task.poll(cx) {
+            Poll::Ready(()) => Poll::Ready(()),
+            Poll::Pending => match shutdown_state {
+                ShutdownState::Active => Poll::Pending,
+                ShutdownState::Shutdown if this.executor.inner.lock().unwrap().active_count > 0 => {
+                    Poll::Pending
+                }
+                _ => Poll::Ready(()),
+            },
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::{yield_to_executor, ExecutionScope};