Reduce the size of the future returned by async get_with and friends

Change the type of the `init` future argument for internal `async fn`s from `Future<...>` to `Pin<&mut Future<...>>`.
moka-rs · Feb 5, 2023 · bc92f4e · bc92f4e
1 parent deec070
commit bc92f4e
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 15 deletions.
diff --git a/src/future/cache.rs b/src/future/cache.rs
@@ -27,13 +27,14 @@ use std::{
     fmt,
     future::Future,
     hash::{BuildHasher, Hash},
+    pin::Pin,
     sync::Arc,
     time::Duration,
 };
 
 //
-// Macros to mitigate the compiler issue that inflates the size of the `init` future
-// in `get_with` and friends methods:
+// Macros to mitigate the compiler issue that inflates the size of future returned
+// from `get_with` and friends methods:
 //
 // - https://github.com/moka-rs/moka/issues/212
 // - https://swatinem.de/blog/future-size/
@@ -43,11 +44,10 @@ use std::{
 // benefits of reducing the future size outweigh this drawback.
 //
 macro_rules! insert_with_hash_and_fun_m(
-    ( $self:ident, $key:expr, $hash:expr, $get:expr, $init:expr, $replace_if:expr, $need_key:expr ) => {
+    ( $self:ident, $key:expr, $hash:expr, $get:expr, $init:ident, $replace_if:expr, $need_key:expr ) => {
         {
             use futures_util::FutureExt;
 
-
             let insert = |v| $self.insert_with_hash($key.clone(), $hash, v).boxed();
 
             let k = if $need_key {
@@ -76,7 +76,7 @@ macro_rules! insert_with_hash_and_fun_m(
 );
 
 macro_rules! get_or_insert_with_hash_and_fun_m(
-    ( $self:ident, $key:expr, $hash:expr, $init:expr, $replace_if:expr, $need_key:expr ) => {
+    ( $self:ident, $key:expr, $hash:expr, $init:ident, $replace_if:expr, $need_key:expr ) => {
         {
             let maybe_entry =
                 $self.base
@@ -95,7 +95,7 @@ macro_rules! get_or_insert_with_hash_and_fun_m(
 );
 
 macro_rules! get_or_insert_with_hash_by_ref_and_fun_m(
-    ( $self:ident, $key:expr, $hash:expr, $init:expr, $replace_if:expr, $need_key:expr ) => {
+    ( $self:ident, $key:expr, $hash:expr, $init:ident, $replace_if:expr, $need_key:expr ) => {
         {
             let maybe_entry =
                 $self.base
@@ -998,6 +998,7 @@ where
     /// `init` futures.
     ///
     pub async fn get_with(&self, key: K, init: impl Future<Output = V>) -> V {
+        futures_util::pin_mut!(init);
         let hash = self.base.hash(&key);
         let key = Arc::new(key);
         let mut replace_if = None as Option<fn(&V) -> bool>;
@@ -1012,6 +1013,7 @@ where
         K: Borrow<Q>,
         Q: ToOwned<Owned = K> + Hash + Eq + ?Sized,
     {
+        futures_util::pin_mut!(init);
         let hash = self.base.hash(key);
         let mut replace_if = None as Option<fn(&V) -> bool>;
         get_or_insert_with_hash_by_ref_and_fun_m!(self, key, hash, init, replace_if, false)
@@ -1027,6 +1029,7 @@ where
         init: impl Future<Output = V>,
         replace_if: impl FnMut(&V) -> bool,
     ) -> V {
+        futures_util::pin_mut!(init);
         let hash = self.base.hash(&key);
         let key = Arc::new(key);
         let mut replace_if = Some(replace_if);
@@ -1129,6 +1132,7 @@ where
     where
         F: Future<Output = Option<V>>,
     {
+        futures_util::pin_mut!(init);
         let hash = self.base.hash(&key);
         let key = Arc::new(key);
         self.get_or_optionally_insert_with_hash_and_fun(key, hash, init, false)
@@ -1146,6 +1150,7 @@ where
         K: Borrow<Q>,
         Q: ToOwned<Owned = K> + Hash + Eq + ?Sized,
     {
+        futures_util::pin_mut!(init);
         let hash = self.base.hash(key);
         self.get_or_optionally_insert_with_hash_by_ref_and_fun(key, hash, init, false)
             .await
@@ -1252,6 +1257,7 @@ where
         F: Future<Output = Result<V, E>>,
         E: Send + Sync + 'static,
     {
+        futures_util::pin_mut!(init);
         let hash = self.base.hash(&key);
         let key = Arc::new(key);
         self.get_or_try_insert_with_hash_and_fun(key, hash, init, false)
@@ -1269,6 +1275,7 @@ where
         K: Borrow<Q>,
         Q: ToOwned<Owned = K> + Hash + Eq + ?Sized,
     {
+        futures_util::pin_mut!(init);
         let hash = self.base.hash(key);
         self.get_or_try_insert_with_hash_by_ref_and_fun(key, hash, init, false)
             .await
@@ -1506,7 +1513,7 @@ where
         &self,
         key: Arc<K>,
         hash: u64,
-        init: impl Future<Output = V>,
+        init: Pin<&mut impl Future<Output = V>>,
         mut replace_if: Option<impl FnMut(&V) -> bool>,
         need_key: bool,
     ) -> Entry<K, V> {
@@ -1517,7 +1524,7 @@ where
         &self,
         key: &Q,
         hash: u64,
-        init: impl Future<Output = V>,
+        init: Pin<&mut impl Future<Output = V>>,
         mut replace_if: Option<impl FnMut(&V) -> bool>,
         need_key: bool,
     ) -> Entry<K, V>
@@ -1571,7 +1578,7 @@ where
         &self,
         key: Arc<K>,
         hash: u64,
-        init: F,
+        init: Pin<&mut F>,
         need_key: bool,
     ) -> Option<Entry<K, V>>
     where
@@ -1590,7 +1597,7 @@ where
         &self,
         key: &Q,
         hash: u64,
-        init: F,
+        init: Pin<&mut F>,
         need_key: bool,
     ) -> Option<Entry<K, V>>
     where
@@ -1612,7 +1619,7 @@ where
         &self,
         key: Arc<K>,
         hash: u64,
-        init: F,
+        init: Pin<&mut F>,
         need_key: bool,
     ) -> Option<Entry<K, V>>
     where
@@ -1654,7 +1661,7 @@ where
         &self,
         key: Arc<K>,
         hash: u64,
-        init: F,
+        init: Pin<&mut F>,
         need_key: bool,
     ) -> Result<Entry<K, V>, Arc<E>>
     where
@@ -1673,7 +1680,7 @@ where
         &self,
         key: &Q,
         hash: u64,
-        init: F,
+        init: Pin<&mut F>,
         need_key: bool,
     ) -> Result<Entry<K, V>, Arc<E>>
     where
@@ -1694,7 +1701,7 @@ where
         &self,
         key: Arc<K>,
         hash: u64,
-        init: F,
+        init: Pin<&mut F>,
         need_key: bool,
     ) -> Result<Entry<K, V>, Arc<E>>
     where

diff --git a/src/future/entry_selector.rs b/src/future/entry_selector.rs
@@ -176,6 +176,7 @@ where
     ///
     /// [get-with-method]: ./struct.Cache.html#method.get_with
     pub async fn or_insert_with(self, init: impl Future<Output = V>) -> Entry<K, V> {
+        futures_util::pin_mut!(init);
         let key = Arc::new(self.owned_key);
         let replace_if = None as Option<fn(&V) -> bool>;
         self.cache
@@ -196,6 +197,7 @@ where
         init: impl Future<Output = V>,
         replace_if: impl FnMut(&V) -> bool,
     ) -> Entry<K, V> {
+        futures_util::pin_mut!(init);
         let key = Arc::new(self.owned_key);
         self.cache
             .get_or_insert_with_hash_and_fun(key, self.hash, init, Some(replace_if), true)
@@ -269,6 +271,7 @@ where
         self,
         init: impl Future<Output = Option<V>>,
     ) -> Option<Entry<K, V>> {
+        futures_util::pin_mut!(init);
         let key = Arc::new(self.owned_key);
         self.cache
             .get_or_optionally_insert_with_hash_and_fun(key, self.hash, init, true)
@@ -344,6 +347,7 @@ where
         F: Future<Output = Result<V, E>>,
         E: Send + Sync + 'static,
     {
+        futures_util::pin_mut!(init);
         let key = Arc::new(self.owned_key);
         self.cache
             .get_or_try_insert_with_hash_and_fun(key, self.hash, init, true)
@@ -521,6 +525,7 @@ where
     ///
     /// [get-with-method]: ./struct.Cache.html#method.get_with
     pub async fn or_insert_with(self, init: impl Future<Output = V>) -> Entry<K, V> {
+        futures_util::pin_mut!(init);
         let replace_if = None as Option<fn(&V) -> bool>;
         self.cache
             .get_or_insert_with_hash_by_ref_and_fun(self.ref_key, self.hash, init, replace_if, true)
@@ -540,6 +545,7 @@ where
         init: impl Future<Output = V>,
         replace_if: impl FnMut(&V) -> bool,
     ) -> Entry<K, V> {
+        futures_util::pin_mut!(init);
         self.cache
             .get_or_insert_with_hash_by_ref_and_fun(
                 self.ref_key,
@@ -617,6 +623,7 @@ where
         self,
         init: impl Future<Output = Option<V>>,
     ) -> Option<Entry<K, V>> {
+        futures_util::pin_mut!(init);
         self.cache
             .get_or_optionally_insert_with_hash_by_ref_and_fun(self.ref_key, self.hash, init, true)
             .await
@@ -692,6 +699,7 @@ where
         F: Future<Output = Result<V, E>>,
         E: Send + Sync + 'static,
     {
+        futures_util::pin_mut!(init);
         self.cache
             .get_or_try_insert_with_hash_by_ref_and_fun(self.ref_key, self.hash, init, true)
             .await

diff --git a/src/future/value_initializer.rs b/src/future/value_initializer.rs
@@ -4,6 +4,7 @@ use std::{
     any::{Any, TypeId},
     future::Future,
     hash::{BuildHasher, Hash},
+    pin::Pin,
     sync::Arc,
 };
 use triomphe::Arc as TrioArc;
@@ -123,7 +124,7 @@ where
         type_id: TypeId,
         // Closure to get an existing value from cache.
         mut get: impl FnMut() -> Option<V>,
-        init: impl Future<Output = O>,
+        init: Pin<&mut impl Future<Output = O>>,
         // Closure to insert a new value into cache.
         mut insert: impl FnMut(V) -> BoxFuture<'a, ()> + Send + 'a,
         // This function will be called after the init future has returned a value of