util: add constructors to TokioContext (#3221)

tokio-util/src/context.rs

@@ -12,21 +12,123 @@ use std::{
     pin::Pin,
     task::{Context, Poll},
 };
-use tokio::runtime::Runtime;
+use tokio::runtime::{Handle, Runtime};
 
 pin_project! {
-    /// `TokioContext` allows connecting a custom executor with the tokio runtime.
+    /// `TokioContext` allows running futures that must be inside Tokio's
+    /// context on a non-Tokio runtime.
     ///
-    /// It contains a `Handle` to the runtime. A handle to the runtime can be
-    /// obtain by calling the `Runtime::handle()` method.
-    pub struct TokioContext<'a, F> {
+    /// It contains a [`Handle`] to the runtime. A handle to the runtime can be
+    /// obtain by calling the [`Runtime::handle()`] method.
+    ///
+    /// Note that the `TokioContext` wrapper only works if the `Runtime` it is
+    /// connected to has not yet been destroyed. You must keep the `Runtime`
+    /// alive until the future has finished executing.
+    ///
+    /// **Warning:** If `TokioContext` is used together with a [current thread]
+    /// runtime, that runtime must be inside a call to `block_on` for the
+    /// wrapped future to work. For this reason, it is recommended to use a
+    /// [multi thread] runtime, even if you configure it to only spawn one
+    /// worker thread.
+    ///
+    /// # Examples
+    ///
+    /// This example creates two runtimes, but only [enables time] on one of
+    /// them. It then uses the context of the runtime with the timer enabled to
+    /// execute a [`sleep`] future on the runtime with timing disabled.
+    /// ```
+    /// use tokio::time::{sleep, Duration};
+    /// use tokio_util::context::RuntimeExt;
+    ///
+    /// // This runtime has timers enabled.
+    /// let rt = tokio::runtime::Builder::new_multi_thread()
+    ///     .enable_all()
+    ///     .build()
+    ///     .unwrap();
+    ///
+    /// // This runtime has timers disabled.
+    /// let rt2 = tokio::runtime::Builder::new_multi_thread()
+    ///     .build()
+    ///     .unwrap();
+    ///
+    /// // Wrap the sleep future in the context of rt.
+    /// let fut = rt.wrap(async { sleep(Duration::from_millis(2)).await });
+    ///
+    /// // Execute the future on rt2.
+    /// rt2.block_on(fut);
+    /// ```
+    ///
+    /// [`Handle`]: struct@tokio::runtime::Handle
+    /// [`Runtime::handle()`]: fn@tokio::runtime::Runtime::handle
+    /// [`RuntimeExt`]: trait@crate::context::RuntimeExt
+    /// [`new_static`]: fn@Self::new_static
+    /// [`sleep`]: fn@tokio::time::sleep
+    /// [current thread]: fn@tokio::runtime::Builder::new_current_thread
+    /// [enables time]: fn@tokio::runtime::Builder::enable_time
+    /// [multi thread]: fn@tokio::runtime::Builder::new_multi_thread
+    pub struct TokioContext<F> {
         #[pin]
         inner: F,
-        handle: &'a Runtime,
+        handle: Handle,
     }
 }
 
-impl<F: Future> Future for TokioContext<'_, F> {
+impl<F> TokioContext<F> {
+    /// Associate the provided future with the context of the runtime behind
+    /// the provided `Handle`.
+    ///
+    /// This constructor uses a `'static` lifetime to opt-out of checking that
+    /// the runtime still exists.
+    ///
+    /// # Examples
+    ///
+    /// This is the same as the example above, but uses the `new` constructor
+    /// rather than [`RuntimeExt::wrap`].
+    ///
+    /// [`RuntimeExt::wrap`]: fn@RuntimeExt::wrap
+    ///
+    /// ```
+    /// use tokio::time::{sleep, Duration};
+    /// use tokio_util::context::TokioContext;
+    ///
+    /// // This runtime has timers enabled.
+    /// let rt = tokio::runtime::Builder::new_multi_thread()
+    ///     .enable_all()
+    ///     .build()
+    ///     .unwrap();
+    ///
+    /// // This runtime has timers disabled.
+    /// let rt2 = tokio::runtime::Builder::new_multi_thread()
+    ///     .build()
+    ///     .unwrap();
+    ///
+    /// let fut = TokioContext::new(
+    ///     async { sleep(Duration::from_millis(2)).await },
+    ///     rt.handle().clone(),
+    /// );
+    ///
+    /// // Execute the future on rt2.
+    /// rt2.block_on(fut);
+    /// ```
+    pub fn new(future: F, handle: Handle) -> TokioContext<F> {
+        TokioContext {
+            inner: future,
+            handle,
+        }
+    }
+
+    /// Obtain a reference to the handle inside this `TokioContext`.
+    pub fn handle(&self) -> &Handle {
+        &self.handle
+    }
+
+    /// Remove the association between the Tokio runtime and the wrapped future.
+    pub fn into_inner(self) -> F {
+        self.inner
+    }
+}
+
+impl<F: Future> Future for TokioContext<F> {
     type Output = F::Output;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
@@ -39,40 +141,50 @@ impl<F: Future> Future for TokioContext<'_, F> {
     }
 }
 
-/// Trait extension that simplifies bundling a `Handle` with a `Future`.
+/// Extension trait that simplifies bundling a `Handle` with a `Future`.
 pub trait RuntimeExt {
-    /// Convenience method that takes a Future and returns a `TokioContext`.
+    /// Create a [`TokioContext`] that wraps the provided future and runs it in
+    /// this runtime's context.
     ///
-    /// # Example: calling Tokio Runtime from a custom ThreadPool
+    /// # Examples
     ///
-    /// ```no_run
-    /// use tokio_util::context::RuntimeExt;
+    /// This example creates two runtimes, but only [enables time] on one of
+    /// them. It then uses the context of the runtime with the timer enabled to
+    /// execute a [`sleep`] future on the runtime with timing disabled.
+    ///
+    /// ```
     /// use tokio::time::{sleep, Duration};
+    /// use tokio_util::context::RuntimeExt;
     ///
+    /// // This runtime has timers enabled.
     /// let rt = tokio::runtime::Builder::new_multi_thread()
     ///     .enable_all()
     ///     .build()
     ///     .unwrap();
     ///
+    /// // This runtime has timers disabled.
     /// let rt2 = tokio::runtime::Builder::new_multi_thread()
     ///     .build()
     ///     .unwrap();
     ///
-    /// let fut = sleep(Duration::from_millis(2));
+    /// // Wrap the sleep future in the context of rt.
+    /// let fut = rt.wrap(async { sleep(Duration::from_millis(2)).await });
     ///
-    /// rt.block_on(
-    ///     rt2
-    ///         .wrap(async { sleep(Duration::from_millis(2)).await }),
-    /// );
-    ///```
-    fn wrap<F: Future>(&self, fut: F) -> TokioContext<'_, F>;
+    /// // Execute the future on rt2.
+    /// rt2.block_on(fut);
+    /// ```
+    ///
+    /// [`TokioContext`]: struct@crate::context::TokioContext
+    /// [`sleep`]: fn@tokio::time::sleep
+    /// [enables time]: fn@tokio::runtime::Builder::enable_time
+    fn wrap<F: Future>(&self, fut: F) -> TokioContext<F>;
 }
 
 impl RuntimeExt for Runtime {
-    fn wrap<F: Future>(&self, fut: F) -> TokioContext<'_, F> {
+    fn wrap<F: Future>(&self, fut: F) -> TokioContext<F> {
         TokioContext {
             inner: fut,
-            handle: self,
+            handle: self.handle().clone(),
         }
     }
 }