Add now_or_later future (#672)

* Add now_or_later future

* Update changelog

* Update CHANGELOG.md

CHANGELOG.md

@@ -13,6 +13,9 @@ vNext (Month Day, Year)
 - Improve documentation on collection-aware builders (#664)
 - Add support for Transcribe `StartStreamTranscription` and S3 `SelectObjectContent` operations (#667)
 
+**Internal Changes**
+- Add NowOrLater future to smithy-async (#672)
+
 
 v0.21 (August 19th, 2021)
 -------------------------

rust-runtime/Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 
 members = [
+    "smithy-async",
     "smithy-http",
     "smithy-client",
     "smithy-eventstream",

rust-runtime/smithy-async/Cargo.toml

@@ -14,3 +14,4 @@ tokio = { version = "1.6", features = ["time"], optional = true }
 
 [dev-dependencies]
 tokio = { version = "1.6", features = ["rt", "macros"] }
+futures-util = "0.3.16"

rust-runtime/smithy-async/src/future/mod.rs

@@ -6,4 +6,5 @@
 //! Useful runtime-agnostic future implementations.
 
 pub mod never;
+pub mod now_or_later;
 pub mod timeout;

rust-runtime/smithy-async/src/future/now_or_later.rs

@@ -0,0 +1,166 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+//! Provides the [`NowOrLater`] future with an explicit `Now` variant
+//!
+//! When a future is immediately, ready, this enables avoiding an unnecessary allocation.
+//! This is intended to be used with `Pin<Box<dyn Future>>` or similar as the future variant. For
+//! convenience, [`BoxFuture`] is provided for this use case.
+//!
+//! Typically, this is used when creating a manual async trait. In this case, it's critical that the
+//! lifetime is captured to enable interop with the async-trait macro.
+//!
+//! # Example
+//!
+//! ```rust
+//! mod future {
+//!   use smithy_async::future::now_or_later::{NowOrLater, BoxFuture};
+//!   use std::future::Future;
+//!   pub struct ProvideRegion<'a>(NowOrLater<Option<String>, BoxFuture<'a, Option<String>>>);
+//!   impl<'a> ProvideRegion<'a> {
+//!       pub fn new(f: impl Future<Output = Option<String>> + Send + 'a) -> Self {
+//!           Self(NowOrLater::new(Box::pin(f)))
+//!       }
+//!
+//!       pub fn ready(region: Option<String>) -> Self {
+//!           Self(NowOrLater::ready(region))
+//!       }
+//!   }
+//! }
+//!
+//! pub trait ProvideRegion {
+//!     fn provide_region<'a>(&'a self) -> future::ProvideRegion<'a> where Self: 'a;
+//! }
+//!
+//! struct AsyncRegionProvider;
+//! impl AsyncRegionProvider {
+//!     async fn region(&self) -> Option<String> {
+//!         todo!()
+//!     }
+//! }
+//!
+//! impl ProvideRegion for AsyncRegionProvider {
+//!     fn provide_region<'a>(&'a self) -> future::ProvideRegion<'a> where Self: 'a {
+//!       future::ProvideRegion::new(self.region())
+//!     }
+//! }
+//! ```
+
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use pin_project_lite::pin_project;
+
+pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
+
+/// Zero sized type for using NowOrLater when no future variant exists.
+pub enum OnlyReady {}
+
+pin_project! {
+    /// Future with an explicit `Now` variant
+    ///
+    /// See the [module documentation](crate::future::now_or_later) for more information.
+    pub struct NowOrLater<T, F> {
+        #[pin]
+        inner: Inner<T, F>
+    }
+}
+
+pin_project! {
+    #[project = NowOrLaterProj]
+    enum Inner<T, F> {
+        #[non_exhaustive]
+        Now { value: Option<T> },
+        #[non_exhaustive]
+        Later { #[pin] future: F },
+    }
+}
+
+impl<T, F> NowOrLater<T, F> {
+    pub fn new(future: F) -> Self {
+        Self {
+            inner: Inner::Later { future },
+        }
+    }
+
+    pub fn ready(value: T) -> NowOrLater<T, F> {
+        let value = Some(value);
+        Self {
+            inner: Inner::Now { value },
+        }
+    }
+}
+
+impl<T, F> Future for NowOrLater<T, F>
+where
+    F: Future<Output = T>,
+{
+    type Output = T;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.project().inner.project() {
+            NowOrLaterProj::Now { value } => {
+                Poll::Ready(value.take().expect("cannot be called twice"))
+            }
+            NowOrLaterProj::Later { future } => future.poll(cx),
+        }
+    }
+}
+
+impl<T> Future for NowOrLater<T, OnlyReady> {
+    type Output = T;
+
+    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.project().inner.project() {
+            NowOrLaterProj::Now { value } => {
+                Poll::Ready(value.take().expect("cannot be called twice"))
+            }
+            NowOrLaterProj::Later { .. } => unreachable!(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::future::now_or_later::{NowOrLater, OnlyReady};
+    use futures_util::FutureExt;
+
+    #[test]
+    fn ready_future_immediately_returns() {
+        let a = true;
+        let f = if a {
+            NowOrLater::ready(5)
+        } else {
+            NowOrLater::new(async { 5 })
+        };
+        use futures_util::FutureExt;
+        assert_eq!(f.now_or_never().expect("future was ready"), 5);
+    }
+
+    #[test]
+    fn only_ready_instantiation() {
+        assert_eq!(
+            NowOrLater::<i32, OnlyReady>::ready(5)
+                .now_or_never()
+                .expect("ready"),
+            5
+        );
+    }
+
+    #[tokio::test]
+    async fn box_dyn_future() {
+        let f = async { 5 };
+        let f = Box::pin(f);
+        let wrapped = NowOrLater::new(f);
+        assert_eq!(wrapped.await, 5);
+    }
+
+    #[tokio::test]
+    async fn async_fn_future() {
+        let wrapped = NowOrLater::new(async { 5 });
+        assert_eq!(wrapped.await, 5);
+    }
+}