From 4ef45640752a7918349de6f81014280bae567dbe Mon Sep 17 00:00:00 2001
From: Alisa Sireneva <me@purplesyringa.moe>
Date: Sat, 30 Nov 2024 08:44:12 +0300
Subject: [PATCH] Reduce memory footprint

By creating the future manually instead of relying on `async { .. }`, we
workaround rustc's inefficient future layouting. On
[a simple benchmark](https://github.com/hez2010/async-runtimes-benchmarks-2024)
spawning 1M of tasks, this reduces memory use from about 512 bytes per
future to about 340 bytes per future.

More context: https://github.com/hez2010/async-runtimes-benchmarks-2024/pull/1
---
 Cargo.toml |  1 +
 src/lib.rs | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 4025c43..2f97cc1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,6 +23,7 @@ async-task = "4.4.0"
 concurrent-queue = "2.5.0"
 fastrand = "2.0.0"
 futures-lite = { version = "2.0.0", default-features = false }
+pin-project-lite = "0.2"
 slab = "0.4.4"
 
 [target.'cfg(target_family = "wasm")'.dependencies]
diff --git a/src/lib.rs b/src/lib.rs
index 710a650..baeda1d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -42,14 +42,16 @@
 use std::fmt;
 use std::marker::PhantomData;
 use std::panic::{RefUnwindSafe, UnwindSafe};
+use std::pin::Pin;
 use std::rc::Rc;
 use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
 use std::sync::{Arc, Mutex, MutexGuard, RwLock, TryLockError};
-use std::task::{Poll, Waker};
+use std::task::{Context, Poll, Waker};
 
 use async_task::{Builder, Runnable};
 use concurrent_queue::ConcurrentQueue;
 use futures_lite::{future, prelude::*};
+use pin_project_lite::pin_project;
 use slab::Slab;
 
 #[cfg(feature = "static")]
@@ -245,10 +247,7 @@ impl<'a> Executor<'a> {
         let entry = active.vacant_entry();
         let index = entry.key();
         let state = self.state_as_arc();
-        let future = async move {
-            let _guard = CallOnDrop(move || drop(state.active().try_remove(index)));
-            future.await
-        };
+        let future = AsyncCallOnDrop::new(future, move || drop(state.active().try_remove(index)));
 
         // Create the task and register it in the set of active tasks.
         //
@@ -1155,6 +1154,32 @@ impl<F: FnMut()> Drop for CallOnDrop<F> {
     }
 }
 
+pin_project! {
+    /// A wrapper around a future, running a closure when dropped.
+    struct AsyncCallOnDrop<Fut, Cleanup: FnMut()> {
+        #[pin]
+        future: Fut,
+        cleanup: CallOnDrop<Cleanup>,
+    }
+}
+
+impl<Fut, Cleanup: FnMut()> AsyncCallOnDrop<Fut, Cleanup> {
+    fn new(future: Fut, cleanup: Cleanup) -> Self {
+        Self {
+            future,
+            cleanup: CallOnDrop(cleanup),
+        }
+    }
+}
+
+impl<Fut: Future, Cleanup: FnMut()> Future for AsyncCallOnDrop<Fut, Cleanup> {
+    type Output = Fut::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.project().future.poll(cx)
+    }
+}
+
 fn _ensure_send_and_sync() {
     use futures_lite::future::pending;