From d600ab9a8f37e9eff3fa8587069a816b65b6da0b Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Thu, 27 Aug 2020 20:05:48 -0400
Subject: rt: Refactor `Runtime::block_on` to take `&self` (#2782)

Co-authored-by: Eliza Weisman <eliza@buoyant.io>
---
 tokio-util/src/context.rs | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

(limited to 'tokio-util/src/context.rs')

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,
         }
     }
 }
-- 
cgit v1.2.3