Rollup merge of #123203 - jkarneges:context-ext, r=Amanieu

Add `Context::ext`

This change enables `Context` to carry arbitrary extension data via a single `&mut dyn Any` field.

```rust
#![feature(context_ext)]

impl Context {
    fn ext(&mut self) -> &mut dyn Any;
}

impl ContextBuilder {
    fn ext(self, data: &'a mut dyn Any) -> Self;

    fn from(cx: &'a mut Context<'_>) -> Self;
    fn waker(self, waker: &'a Waker) -> Self;
}
```

Basic usage:

```rust
struct MyExtensionData {
    executor_name: String,
}

let mut ext = MyExtensionData {
    executor_name: "foo".to_string(),
};

let mut cx = ContextBuilder::from_waker(&waker).ext(&mut ext).build();

if let Some(ext) = cx.ext().downcast_mut::<MyExtensionData>() {
    println!("{}", ext.executor_name);
}
```

Currently, `Context` only carries a `Waker`, but there is interest in having it carry other kinds of data. Examples include [LocalWaker](https://github.com/rust-lang/rust/issues/118959), [a reactor interface](https://github.com/rust-lang/libs-team/issues/347), and [multiple arbitrary values by type](https://docs.rs/context-rs/latest/context_rs/). There is also a general practice in the ecosystem of sharing data between executors and futures via thread-locals or globals that would arguably be better shared via `Context`, if it were possible.

The `ext` field would provide a low friction (to stabilization) solution to enable experimentation. It would enable experimenting with what kinds of data we want to carry as well as with what data structures we may want to use to carry such data.

Dedicated fields for specific kinds of data could still be added directly on `Context` when we have sufficient experience or understanding about the problem they are solving, such as with `LocalWaker`. The `ext` field would be for data for which we don't have such experience or understanding, and that could be graduated to dedicated fields once proven.

Both the provider and consumer of the extension data must be aware of the concrete type behind the `Any`. This means it is not possible for the field to carry an abstract interface. However, the field can carry a concrete type which in turn carries an interface. There are different ways one can imagine an interface-carrying concrete type to work, hence the benefit of being able to experiment with such data structures.

## Passing interfaces

Interfaces can be placed in a concrete type, such as a struct, and then that type can be casted to `Any`. However, one gotcha is `Any` cannot contain non-static references. This means one cannot simply do:

```rust
struct Extensions<'a> {
    interface1: &'a mut dyn Trait1,
    interface2: &'a mut dyn Trait2,
}

let mut ext = Extensions {
    interface1: &mut impl1,
    interface2: &mut impl2,
};

let ext: &mut dyn Any = &mut ext;
```

To work around this without boxing, unsafe code can be used to create a safe projection using accessors. For example:

```rust
pub struct Extensions {
    interface1: *mut dyn Trait1,
    interface2: *mut dyn Trait2,
}

impl Extensions {
    pub fn new<'a>(
        interface1: &'a mut (dyn Trait1 + 'static),
        interface2: &'a mut (dyn Trait2 + 'static),
        scratch: &'a mut MaybeUninit<Self>,
    ) -> &'a mut Self {
        scratch.write(Self {
            interface1,
            interface2,
        })
    }

    pub fn interface1(&mut self) -> &mut dyn Trait1 {
        unsafe { self.interface1.as_mut().unwrap() }
    }

    pub fn interface2(&mut self) -> &mut dyn Trait2 {
        unsafe { self.interface2.as_mut().unwrap() }
    }
}

let mut scratch = MaybeUninit::uninit();
let ext: &mut Extensions = Extensions::new(&mut impl1, &mut impl2, &mut scratch);

// ext can now be casted to `&mut dyn Any` and back, and used safely
let ext: &mut dyn Any = ext;
```

## Context inheritance

Sometimes when futures poll other futures they want to provide their own `Waker` which requires creating their own `Context`. Unfortunately, polling sub-futures with a fresh `Context` means any properties on the original `Context` won't get propagated along to the sub-futures. To help with this, some additional methods are added to `ContextBuilder`.

Here's how to derive a new `Context` from another, overriding only the `Waker`:

```rust
let mut cx = ContextBuilder::from(parent_cx).waker(&new_waker).build();
```

3 files changed, 102 insertions, 9 deletions

diff --git a/library/core/src/task/wake.rs b/library/core/src/task/wake.rs
index 1b43c46bda5..fe39f6c0b56 100644
--- a/library/core/src/task/wake.rs
+++ b/library/core/src/task/wake.rs
@@ -2,6 +2,7 @@
 
 use crate::mem::transmute;
 
+use crate::any::Any;
 use crate::fmt;
 use crate::marker::PhantomData;
 use crate::ptr;
@@ -220,6 +221,12 @@ impl RawWakerVTable {
     }
 }
 
+#[derive(Debug)]
+enum ExtData<'a> {
+    Some(&'a mut dyn Any),
+    None(()),
+}
+
 /// The context of an asynchronous task.
 ///
 /// Currently, `Context` only serves to provide access to a [`&Waker`](Waker)
@@ -229,6 +236,7 @@ impl RawWakerVTable {
 pub struct Context<'a> {
     waker: &'a Waker,
     local_waker: &'a LocalWaker,
+    ext: ExtData<'a>,
     // Ensure we future-proof against variance changes by forcing
     // the lifetime to be invariant (argument-position lifetimes
     // are contravariant while return-position lifetimes are
@@ -257,6 +265,7 @@ impl<'a> Context<'a> {
     pub const fn waker(&self) -> &'a Waker {
         &self.waker
     }
+
     /// Returns a reference to the [`LocalWaker`] for the current task.
     #[inline]
     #[unstable(feature = "local_waker", issue = "118959")]
@@ -264,6 +273,17 @@ impl<'a> Context<'a> {
     pub const fn local_waker(&self) -> &'a LocalWaker {
         &self.local_waker
     }
+
+    /// Returns a reference to the extension data for the current task.
+    #[inline]
+    #[unstable(feature = "context_ext", issue = "123392")]
+    #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
+    pub const fn ext(&mut self) -> &mut dyn Any {
+        match &mut self.ext {
+            ExtData::Some(data) => *data,
+            ExtData::None(unit) => unit,
+        }
+    }
 }
 
 #[stable(feature = "futures_api", since = "1.36.0")]
@@ -300,6 +320,7 @@ impl fmt::Debug for Context<'_> {
 pub struct ContextBuilder<'a> {
     waker: &'a Waker,
     local_waker: &'a LocalWaker,
+    ext: ExtData<'a>,
     // Ensure we future-proof against variance changes by forcing
     // the lifetime to be invariant (argument-position lifetimes
     // are contravariant while return-position lifetimes are
@@ -318,7 +339,39 @@ impl<'a> ContextBuilder<'a> {
     pub const fn from_waker(waker: &'a Waker) -> Self {
         // SAFETY: LocalWaker is just Waker without thread safety
         let local_waker = unsafe { transmute(waker) };
-        Self { waker: waker, local_waker, _marker: PhantomData, _marker2: PhantomData }
+        Self {
+            waker: waker,
+            local_waker,
+            ext: ExtData::None(()),
+            _marker: PhantomData,
+            _marker2: PhantomData,
+        }
+    }
+
+    /// Create a ContextBuilder from an existing Context.
+    #[inline]
+    #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
+    #[unstable(feature = "context_ext", issue = "123392")]
+    pub const fn from(cx: &'a mut Context<'_>) -> Self {
+        let ext = match &mut cx.ext {
+            ExtData::Some(ext) => ExtData::Some(*ext),
+            ExtData::None(()) => ExtData::None(()),
+        };
+        Self {
+            waker: cx.waker,
+            local_waker: cx.local_waker,
+            ext,
+            _marker: PhantomData,
+            _marker2: PhantomData,
+        }
+    }
+
+    /// This method is used to set the value for the waker on `Context`.
+    #[inline]
+    #[unstable(feature = "context_ext", issue = "123392")]
+    #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
+    pub const fn waker(self, waker: &'a Waker) -> Self {
+        Self { waker, ..self }
     }
 
     /// This method is used to set the value for the local waker on `Context`.
@@ -329,13 +382,21 @@ impl<'a> ContextBuilder<'a> {
         Self { local_waker, ..self }
     }
 
+    /// This method is used to set the value for the extension data on `Context`.
+    #[inline]
+    #[unstable(feature = "context_ext", issue = "123392")]
+    #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
+    pub const fn ext(self, data: &'a mut dyn Any) -> Self {
+        Self { ext: ExtData::Some(data), ..self }
+    }
+
     /// Builds the `Context`.
     #[inline]
     #[unstable(feature = "local_waker", issue = "118959")]
     #[rustc_const_unstable(feature = "const_waker", issue = "102012")]
     pub const fn build(self) -> Context<'a> {
-        let ContextBuilder { waker, local_waker, _marker, _marker2 } = self;
-        Context { waker, local_waker, _marker, _marker2 }
+        let ContextBuilder { waker, local_waker, ext, _marker, _marker2 } = self;
+        Context { waker, local_waker, ext, _marker, _marker2 }
     }
 }
 
diff --git a/tests/ui/async-await/async-is-unwindsafe.rs b/tests/ui/async-await/async-is-unwindsafe.rs
index 53009b6e741..d0202f72f00 100644
--- a/tests/ui/async-await/async-is-unwindsafe.rs
+++ b/tests/ui/async-await/async-is-unwindsafe.rs
@@ -11,6 +11,7 @@ fn main() {
 
     is_unwindsafe(async {
         //~^ ERROR the type `&mut Context<'_>` may not be safely transferred across an unwind boundary
+        //~| ERROR the type `&mut (dyn Any + 'static)` may not be safely transferred across an unwind boundary
         use std::ptr::null;
         use std::task::{Context, RawWaker, RawWakerVTable, Waker};
         let waker = unsafe {
diff --git a/tests/ui/async-await/async-is-unwindsafe.stderr b/tests/ui/async-await/async-is-unwindsafe.stderr
index 5d87fc74768..6bb06df9f39 100644
--- a/tests/ui/async-await/async-is-unwindsafe.stderr
+++ b/tests/ui/async-await/async-is-unwindsafe.stderr
@@ -6,19 +6,18 @@ LL |        is_unwindsafe(async {
    |  |_____|
    | ||
 LL | ||
+LL | ||
 LL | ||         use std::ptr::null;
-LL | ||         use std::task::{Context, RawWaker, RawWakerVTable, Waker};
 ...  ||
 LL | ||         drop(cx_ref);
 LL | ||     });
    | ||_____-^ `&mut Context<'_>` may not be safely transferred across an unwind boundary
    |  |_____|
-   |        within this `{async block@$DIR/async-is-unwindsafe.rs:12:19: 29:6}`
+   |        within this `{async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6}`
    |
-   = help: within `{async block@$DIR/async-is-unwindsafe.rs:12:19: 29:6}`, the trait `UnwindSafe` is not implemented for `&mut Context<'_>`, which is required by `{async block@$DIR/async-is-unwindsafe.rs:12:19: 29:6}: UnwindSafe`
-   = note: `UnwindSafe` is implemented for `&Context<'_>`, but not for `&mut Context<'_>`
+   = help: within `{async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6}`, the trait `UnwindSafe` is not implemented for `&mut Context<'_>`, which is required by `{async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6}: UnwindSafe`
 note: future does not implement `UnwindSafe` as this value is used across an await
-  --> $DIR/async-is-unwindsafe.rs:25:18
+  --> $DIR/async-is-unwindsafe.rs:26:18
    |
 LL |         let cx_ref = &mut cx;
    |             ------ has type `&mut Context<'_>` which does not implement `UnwindSafe`
@@ -31,6 +30,38 @@ note: required by a bound in `is_unwindsafe`
 LL | fn is_unwindsafe(_: impl std::panic::UnwindSafe) {}
    |                          ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `is_unwindsafe`
 
-error: aborting due to 1 previous error
+error[E0277]: the type `&mut (dyn Any + 'static)` may not be safely transferred across an unwind boundary
+  --> $DIR/async-is-unwindsafe.rs:12:5
+   |
+LL |        is_unwindsafe(async {
+   |   _____^_____________-
+   |  |_____|
+   | ||
+LL | ||
+LL | ||
+LL | ||         use std::ptr::null;
+...  ||
+LL | ||         drop(cx_ref);
+LL | ||     });
+   | ||_____-^ `&mut (dyn Any + 'static)` may not be safely transferred across an unwind boundary
+   |  |_____|
+   |        within this `{async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6}`
+   |
+   = help: within `{async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6}`, the trait `UnwindSafe` is not implemented for `&mut (dyn Any + 'static)`, which is required by `{async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6}: UnwindSafe`
+note: future does not implement `UnwindSafe` as this value is used across an await
+  --> $DIR/async-is-unwindsafe.rs:26:18
+   |
+LL |         let mut cx = Context::from_waker(&waker);
+   |             ------ has type `Context<'_>` which does not implement `UnwindSafe`
+...
+LL |         async {}.await; // this needs an inner await point
+   |                  ^^^^^ await occurs here, with `mut cx` maybe used later
+note: required by a bound in `is_unwindsafe`
+  --> $DIR/async-is-unwindsafe.rs:3:26
+   |
+LL | fn is_unwindsafe(_: impl std::panic::UnwindSafe) {}
+   |                          ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `is_unwindsafe`
+
+error: aborting due to 2 previous errors
 
 For more information about this error, try `rustc --explain E0277`.