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

Co-authored-by: Eliza Weisman <eliza@buoyant.io>

1 files changed, 11 insertions, 12 deletions

diff --git a/tokio-util/src/context.rs b/tokio-util/src/context.rs
index f6289093..e07538d9 100644
--- a/tokio-util/src/context.rs
+++ b/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,
         }
     }
 }