rt: Refactor Runtime::block_on to take &self (#2782)

Co-authored-by: Eliza Weisman <[email protected]>

tokio-util/src/context.rs

@@ -12,21 +12,21 @@ use std::{
     pin::Pin,
     task::{Context, Poll},
 };
-use tokio::runtime::Handle;
+use tokio::runtime::Runtime;
 
 pin_project! {
     /// `TokioContext` allows connecting a custom executor with the 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<F> {
+    pub struct TokioContext<'a, F> {
         #[pin]
         inner: F,
-        handle: Handle,
+        handle: &'a Runtime,
     }
 }
 
-impl<F: Future> Future for TokioContext<F> {
+impl<F: Future> Future for TokioContext<'_, F> {
     type Output = F::Output;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
@@ -39,16 +39,16 @@ impl<F: Future> Future for TokioContext<F> {
 }
 
 /// Trait extension that simplifies bundling a `Handle` with a `Future`.
-pub trait HandleExt {
+pub trait RuntimeExt {
     /// Convenience method that takes a Future and returns a `TokioContext`.
     ///
     /// # Example: calling Tokio Runtime from a custom ThreadPool
     ///
     /// ```no_run
-    /// use tokio_util::context::HandleExt;
+    /// use tokio_util::context::RuntimeExt;
     /// use tokio::time::{delay_for, Duration};
     ///
-    /// let mut rt = tokio::runtime::Builder::new()
+    /// let rt = tokio::runtime::Builder::new()
     ///     .threaded_scheduler()
     ///     .enable_all()
     ///     .build().unwrap();
@@ -61,18 +61,17 @@ pub trait HandleExt {
     ///
     /// rt.block_on(
     ///     rt2
-    ///         .handle()
     ///         .wrap(async { delay_for(Duration::from_millis(2)).await }),
     /// );
     ///```
-    fn wrap<F: Future>(&self, fut: F) -> TokioContext<F>;
+    fn wrap<F: Future>(&self, fut: F) -> TokioContext<'_, F>;
 }
 
-impl HandleExt for Handle {
-    fn wrap<F: Future>(&self, fut: F) -> TokioContext<F> {
+impl RuntimeExt for Runtime {
+    fn wrap<F: Future>(&self, fut: F) -> TokioContext<'_, F> {
         TokioContext {
             inner: fut,
-            handle: self.clone(),
+            handle: self,
         }
     }
 }

tokio-util/tests/context.rs

@@ -2,11 +2,11 @@
 
 use tokio::runtime::Builder;
 use tokio::time::*;
-use tokio_util::context::HandleExt;
+use tokio_util::context::RuntimeExt;
 
 #[test]
 fn tokio_context_with_another_runtime() {
-    let mut rt1 = Builder::new()
+    let rt1 = Builder::new()
         .threaded_scheduler()
         .core_threads(1)
         // no timer!
@@ -21,8 +21,5 @@ fn tokio_context_with_another_runtime() {
 
     // Without the `HandleExt.wrap()` there would be a panic because there is
     // no timer running, since it would be referencing runtime r1.
-    let _ = rt1.block_on(
-        rt2.handle()
-            .wrap(async move { delay_for(Duration::from_millis(2)).await }),
-    );
+    let _ = rt1.block_on(rt2.wrap(async move { delay_for(Duration::from_millis(2)).await }));
 }