feat: add try_waker and From<&mut Context> for ContextBuilder to allow the extention of contexts by futures

Author: tvallotton

library/alloc/src/task.rs

@@ -165,16 +165,15 @@ fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
 
 /// # Examples
 ///
-/// A
-///
 /// This is a simplified example of a `spawn` and a `block_on` function. The `spawn` function
 /// is used to push new tasks onto the run queue, while the block on function will remove them
 /// and poll them. When a task is woken, it will put itself back on the run queue to be polled by the executor.
 ///
 /// **Note:** A real world example would interlieve poll calls with calls to an io reactor to wait for events instead
-/// of spinning on a loop.
+/// of spinning on a loop. 
 ///
 /// ```rust
+/// #![feature(local_waker)]
 /// use std::task::{LocalWake, ContextBuilder, LocalWaker};
 /// use std::future::Future;
 /// use std::pin::Pin;
@@ -204,9 +203,9 @@ fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
 /// where
 ///     F: Future<Output=()> + 'static + Send + Sync
 /// {
-///     let task = Rc::new(Box::pin(future));
+///     let task = RefCell::new(Box::pin(future));
 ///     RUN_QUEUE.with_borrow_mut(|queue| {
-///         queue.push_back(task)
+///         queue.push_back(Rc::new(Task(task)));
 ///     });
 /// }
 ///
@@ -221,19 +220,22 @@ fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
 ///             return;
 ///         };
 ///         // cast the Rc<Task> into a `LocalWaker`
-///         let waker: LocalWaker = task.into();
+///         let waker: LocalWaker = task.clone().into();
 ///         // Build the context using `ContextBuilder`
-///         let mut cx = ContextBuilder::new()
-///             .local_waker(&waker)
+///         let mut cx = ContextBuilder::from_local_waker(&waker)
 ///             .build();
 ///
 ///         // Poll the task
-///         task.0
+///         let _ = task.0
 ///             .borrow_mut()
 ///             .as_mut()
 ///             .poll(&mut cx);
 ///     }
 /// }
+///
+/// block_on(async {
+///     println!("hello world");
+/// });
 /// ```
 ///
 #[unstable(feature = "local_waker", issue = "none")]

library/core/src/task/wake.rs

@@ -225,21 +225,32 @@ impl<'a> Context<'a> {
     /// # Panics
     /// This function will panic if no `Waker` was set on the context. This happens if
     /// the executor does not support working with thread safe wakers. An alternative
-    /// may be to call [`.local_waker()`](Context::local_waker) instead.
+    /// may be to call [`.local_waker()`](Context::local_waker) instead. For a fallible
+    /// version of this function see [`.try_waker()`](Context::try_waker).
+    #[inline]
+    #[must_use]
     #[stable(feature = "futures_api", since = "1.36.0")]
     #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
-    #[must_use]
-    #[inline]
     pub const fn waker(&self) -> &'a Waker {
         &self
             .waker
             .expect("no waker was set on this context, consider calling `local_waker` instead.")
     }
     /// Returns a reference to the [`LocalWaker`] for the current task.
+    #[inline]
     #[unstable(feature = "local_waker", issue = "none")]
-    pub fn local_waker(&self) -> &'a LocalWaker {
+    #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
+    pub const fn local_waker(&self) -> &'a LocalWaker {
         &self.local_waker
     }
+    /// Returns a `Some(&Waker)` if a waker was defined on the `Context`,
+    /// otherwise it returns `None`.
+    #[inline]
+    #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
+    #[unstable(feature = "local_waker", issue = "none")]
+    pub const fn try_waker(&self) -> Option<&'a Waker> {
+        self.waker
+    }
 }
 
 #[stable(feature = "futures_api", since = "1.36.0")]
@@ -256,18 +267,19 @@ impl fmt::Debug for Context<'_> {
 /// ```
 /// #![feature(local_waker)]
 /// #![feature(noop_waker)]
-/// use std::task::{ContextBuilder, LocalWaker, Waker};
-///
+/// use std::task::{ContextBuilder, LocalWaker, Waker, Poll};
+/// use std::future::Future;
+/// 
 /// let local_waker = LocalWaker::noop();
 /// let waker = Waker::noop();
 ///
-/// let context = ContextBuilder::from_local_waker(&local_waker)
+/// let mut cx = ContextBuilder::from_local_waker(&local_waker)
 ///     .waker(&waker)
 ///     .build();
 ///
-/// let future = pin::pin!(async { 20 });
-/// let poll = future.poll(&mut context);
-/// assert_eq!(poll, task::Poll::Ready(20));
+/// let mut future = std::pin::pin!(async { 20 });
+/// let poll = future.as_mut().poll(&mut cx);
+/// assert_eq!(poll, Poll::Ready(20));
 ///
 /// ```
 #[unstable(feature = "local_waker", issue = "none")]
@@ -323,6 +335,50 @@ impl<'a> ContextBuilder<'a> {
     }
 }
 
+/// Construct a `ContextBuilder`` from a `Context`. This is useful for
+/// overriding values from a context.
+///
+/// # Examples
+/// An example of a future that allows to set a Waker on Context if none was defined.
+/// This can be used to await futures that require a `Waker` even if the runtime does not
+/// support `Waker`.
+/// ```rust
+/// #![feature(noop_waker, local_waker)]
+/// use std::task::{Waker, ContextBuilder};
+/// use std::future::{poll_fn, Future};
+/// use std::pin::pin;
+///
+/// async fn with_waker<F>(f: F, waker: &Waker) -> F::Output
+/// where
+///     F: Future
+/// {
+///     let mut f = pin!(f);
+///     poll_fn(move |cx| {
+///         let has_waker = cx.try_waker().is_some();
+///         if has_waker {
+///             return f.as_mut().poll(cx);
+///         }
+///
+///         let mut cx = ContextBuilder::from(cx)
+///             .waker(&waker)
+///             .build();
+///         f.as_mut().poll(&mut cx)
+///     }).await
+/// }
+/// 
+/// # async fn __() {
+/// with_waker(async { /* ... */ }, &Waker::noop()).await;
+/// # }
+/// ```
+#[unstable(feature = "local_waker", issue = "none")]
+impl<'a> From<&mut Context<'a>> for ContextBuilder<'a> {
+    #[inline]
+    fn from(value: &mut Context<'a>) -> Self {
+        let Context { waker, local_waker, .. } = *value;
+        ContextBuilder { waker, local_waker }
+    }
+}
+
 /// A `Waker` is a handle for waking up a task by notifying its executor that it
 /// is ready to be run.
 ///
@@ -559,12 +615,11 @@ impl fmt::Debug for Waker {
 ///     })
 /// }
 ///
-/// # #[allow(unused_must_use)]
 /// # async fn __() {
 /// yield_now().await;
 /// # }
 /// ```
-/// 
+///
 /// [`Future::poll()`]: core::future::Future::poll
 /// [`Poll::Pending`]: core::task::Poll::Pending
 /// [`local_waker`]: core::task::Context::local_waker