Merging upstream version 1.72.1+dfsg1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 232 insertions, 58 deletions

diff --git a/vendor/tokio/src/task/task_local.rs b/vendor/tokio/src/task/task_local.rs
index 6571ffd7b..eeadfbd3e 100644
--- a/vendor/tokio/src/task/task_local.rs
+++ b/vendor/tokio/src/task/task_local.rs
@@ -2,9 +2,10 @@ use pin_project_lite::pin_project;
 use std::cell::RefCell;
 use std::error::Error;
 use std::future::Future;
+use std::marker::PhantomPinned;
 use std::pin::Pin;
 use std::task::{Context, Poll};
-use std::{fmt, thread};
+use std::{fmt, mem, thread};
 
 /// Declares a new task-local key of type [`tokio::task::LocalKey`].
 ///
@@ -47,9 +48,27 @@ macro_rules! task_local {
 }
 
 #[doc(hidden)]
+#[cfg(not(tokio_no_const_thread_local))]
 #[macro_export]
 macro_rules! __task_local_inner {
     ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
+        $(#[$attr])*
+        $vis static $name: $crate::task::LocalKey<$t> = {
+            std::thread_local! {
+                static __KEY: std::cell::RefCell<Option<$t>> = const { std::cell::RefCell::new(None) };
+            }
+
+            $crate::task::LocalKey { inner: __KEY }
+        };
+    };
+}
+
+#[doc(hidden)]
+#[cfg(tokio_no_const_thread_local)]
+#[macro_export]
+macro_rules! __task_local_inner {
+    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
+        $(#[$attr])*
         $vis static $name: $crate::task::LocalKey<$t> = {
             std::thread_local! {
                 static __KEY: std::cell::RefCell<Option<$t>> = std::cell::RefCell::new(None);
@@ -62,7 +81,7 @@ macro_rules! __task_local_inner {
 
 /// A key for task-local data.
 ///
-/// This type is generated by the `task_local!` macro.
+/// This type is generated by the [`task_local!`] macro.
 ///
 /// Unlike [`std::thread::LocalKey`], `tokio::task::LocalKey` will
 /// _not_ lazily initialize the value on first access. Instead, the
@@ -90,7 +109,9 @@ macro_rules! __task_local_inner {
 /// }).await;
 /// # }
 /// ```
+///
 /// [`std::thread::LocalKey`]: struct@std::thread::LocalKey
+/// [`task_local!`]: ../macro.task_local.html
 #[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
 pub struct LocalKey<T: 'static> {
     #[doc(hidden)]
@@ -102,6 +123,11 @@ impl<T: 'static> LocalKey<T> {
     ///
     /// On completion of `scope`, the task-local will be dropped.
     ///
+    /// ### Panics
+    ///
+    /// If you poll the returned future inside a call to [`with`] or
+    /// [`try_with`] on the same `LocalKey`, then the call to `poll` will panic.
+    ///
     /// ### Examples
     ///
     /// ```
@@ -115,21 +141,29 @@ impl<T: 'static> LocalKey<T> {
     /// }).await;
     /// # }
     /// ```
-    pub async fn scope<F>(&'static self, value: T, f: F) -> F::Output
+    ///
+    /// [`with`]: fn@Self::with
+    /// [`try_with`]: fn@Self::try_with
+    pub fn scope<F>(&'static self, value: T, f: F) -> TaskLocalFuture<T, F>
     where
         F: Future,
     {
         TaskLocalFuture {
-            local: &self,
+            local: self,
             slot: Some(value),
-            future: f,
+            future: Some(f),
+            _pinned: PhantomPinned,
         }
-        .await
     }
 
     /// Sets a value `T` as the task-local value for the closure `F`.
     ///
-    /// On completion of `scope`, the task-local will be dropped.
+    /// On completion of `sync_scope`, the task-local will be dropped.
+    ///
+    /// ### Panics
+    ///
+    /// This method panics if called inside a call to [`with`] or [`try_with`]
+    /// on the same `LocalKey`.
     ///
     /// ### Examples
     ///
@@ -144,32 +178,80 @@ impl<T: 'static> LocalKey<T> {
     /// });
     /// # }
     /// ```
+    ///
+    /// [`with`]: fn@Self::with
+    /// [`try_with`]: fn@Self::try_with
+    #[track_caller]
     pub fn sync_scope<F, R>(&'static self, value: T, f: F) -> R
     where
         F: FnOnce() -> R,
     {
-        let mut scope = TaskLocalFuture {
-            local: &self,
-            slot: Some(value),
-            future: (),
-        };
-        Pin::new(&mut scope).with_task(|_| f())
+        let mut value = Some(value);
+        match self.scope_inner(&mut value, f) {
+            Ok(res) => res,
+            Err(err) => err.panic(),
+        }
+    }
+
+    fn scope_inner<F, R>(&'static self, slot: &mut Option<T>, f: F) -> Result<R, ScopeInnerErr>
+    where
+        F: FnOnce() -> R,
+    {
+        struct Guard<'a, T: 'static> {
+            local: &'static LocalKey<T>,
+            slot: &'a mut Option<T>,
+        }
+
+        impl<'a, T: 'static> Drop for Guard<'a, T> {
+            fn drop(&mut self) {
+                // This should not panic.
+                //
+                // We know that the RefCell was not borrowed before the call to
+                // `scope_inner`, so the only way for this to panic is if the
+                // closure has created but not destroyed a RefCell guard.
+                // However, we never give user-code access to the guards, so
+                // there's no way for user-code to forget to destroy a guard.
+                //
+                // The call to `with` also should not panic, since the
+                // thread-local wasn't destroyed when we first called
+                // `scope_inner`, and it shouldn't have gotten destroyed since
+                // then.
+                self.local.inner.with(|inner| {
+                    let mut ref_mut = inner.borrow_mut();
+                    mem::swap(self.slot, &mut *ref_mut);
+                });
+            }
+        }
+
+        self.inner.try_with(|inner| {
+            inner
+                .try_borrow_mut()
+                .map(|mut ref_mut| mem::swap(slot, &mut *ref_mut))
+        })??;
+
+        let guard = Guard { local: self, slot };
+
+        let res = f();
+
+        drop(guard);
+
+        Ok(res)
     }
 
     /// Accesses the current task-local and runs the provided closure.
     ///
     /// # Panics
     ///
-    /// This function will panic if not called within the context
-    /// of a future containing a task-local with the corresponding key.
+    /// This function will panic if the task local doesn't have a value set.
+    #[track_caller]
     pub fn with<F, R>(&'static self, f: F) -> R
     where
         F: FnOnce(&T) -> R,
     {
-        self.try_with(f).expect(
-            "cannot access a Task Local Storage value \
-             without setting it via `LocalKey::set`",
-        )
+        match self.try_with(f) {
+            Ok(res) => res,
+            Err(_) => panic!("cannot access a task-local storage value without setting it first"),
+        }
     }
 
     /// Accesses the current task-local and runs the provided closure.
@@ -181,19 +263,32 @@ impl<T: 'static> LocalKey<T> {
     where
         F: FnOnce(&T) -> R,
     {
-        self.inner.with(|v| {
-            if let Some(val) = v.borrow().as_ref() {
-                Ok(f(val))
-            } else {
-                Err(AccessError { _private: () })
-            }
-        })
+        // If called after the thread-local storing the task-local is destroyed,
+        // then we are outside of a closure where the task-local is set.
+        //
+        // Therefore, it is correct to return an AccessError if `try_with`
+        // returns an error.
+        let try_with_res = self.inner.try_with(|v| {
+            // This call to `borrow` cannot panic because no user-defined code
+            // runs while a `borrow_mut` call is active.
+            v.borrow().as_ref().map(f)
+        });
+
+        match try_with_res {
+            Ok(Some(res)) => Ok(res),
+            Ok(None) | Err(_) => Err(AccessError { _private: () }),
+        }
     }
 }
 
 impl<T: Copy + 'static> LocalKey<T> {
     /// Returns a copy of the task-local value
     /// if the task-local value implements `Copy`.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the task local doesn't have a value set.
+    #[track_caller]
     pub fn get(&'static self) -> T {
         self.with(|v| *v)
     }
@@ -206,55 +301,107 @@ impl<T: 'static> fmt::Debug for LocalKey<T> {
 }
 
 pin_project! {
-    struct TaskLocalFuture<T: StaticLifetime, F> {
+    /// A future that sets a value `T` of a task local for the future `F` during
+    /// its execution.
+    ///
+    /// The value of the task-local must be `'static` and will be dropped on the
+    /// completion of the future.
+    ///
+    /// Created by the function [`LocalKey::scope`](self::LocalKey::scope).
+    ///
+    /// ### Examples
+    ///
+    /// ```
+    /// # async fn dox() {
+    /// tokio::task_local! {
+    ///     static NUMBER: u32;
+    /// }
+    ///
+    /// NUMBER.scope(1, async move {
+    ///     println!("task local value: {}", NUMBER.get());
+    /// }).await;
+    /// # }
+    /// ```
+    pub struct TaskLocalFuture<T, F>
+    where
+        T: 'static,
+    {
         local: &'static LocalKey<T>,
         slot: Option<T>,
         #[pin]
-        future: F,
+        future: Option<F>,
+        #[pin]
+        _pinned: PhantomPinned,
     }
-}
 
-impl<T: 'static, F> TaskLocalFuture<T, F> {
-    fn with_task<F2: FnOnce(Pin<&mut F>) -> R, R>(self: Pin<&mut Self>, f: F2) -> R {
-        struct Guard<'a, T: 'static> {
-            local: &'static LocalKey<T>,
-            slot: &'a mut Option<T>,
-            prev: Option<T>,
-        }
-
-        impl<T> Drop for Guard<'_, T> {
-            fn drop(&mut self) {
-                let value = self.local.inner.with(|c| c.replace(self.prev.take()));
-                *self.slot = value;
+    impl<T: 'static, F> PinnedDrop for TaskLocalFuture<T, F> {
+        fn drop(this: Pin<&mut Self>) {
+            let this = this.project();
+            if mem::needs_drop::<F>() && this.future.is_some() {
+                // Drop the future while the task-local is set, if possible. Otherwise
+                // the future is dropped normally when the `Option<F>` field drops.
+                let mut future = this.future;
+                let _ = this.local.scope_inner(this.slot, || {
+                    future.set(None);
+                });
             }
         }
-
-        let mut project = self.project();
-        let val = project.slot.take();
-
-        let prev = project.local.inner.with(|c| c.replace(val));
-
-        let _guard = Guard {
-            prev,
-            slot: &mut project.slot,
-            local: *project.local,
-        };
-
-        f(project.future)
     }
 }
 
 impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
     type Output = F::Output;
 
+    #[track_caller]
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        self.with_task(|f| f.poll(cx))
+        let this = self.project();
+        let mut future_opt = this.future;
+
+        let res = this
+            .local
+            .scope_inner(this.slot, || match future_opt.as_mut().as_pin_mut() {
+                Some(fut) => {
+                    let res = fut.poll(cx);
+                    if res.is_ready() {
+                        future_opt.set(None);
+                    }
+                    Some(res)
+                }
+                None => None,
+            });
+
+        match res {
+            Ok(Some(res)) => res,
+            Ok(None) => panic!("`TaskLocalFuture` polled after completion"),
+            Err(err) => err.panic(),
+        }
     }
 }
 
-// Required to make `pin_project` happy.
-trait StaticLifetime: 'static {}
-impl<T: 'static> StaticLifetime for T {}
+impl<T: 'static, F> fmt::Debug for TaskLocalFuture<T, F>
+where
+    T: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        /// Format the Option without Some.
+        struct TransparentOption<'a, T> {
+            value: &'a Option<T>,
+        }
+        impl<'a, T: fmt::Debug> fmt::Debug for TransparentOption<'a, T> {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                match self.value.as_ref() {
+                    Some(value) => value.fmt(f),
+                    // Hitting the None branch should not be possible.
+                    None => f.pad("<missing>"),
+                }
+            }
+        }
+
+        f.debug_struct("TaskLocalFuture")
+            .field("value", &TransparentOption { value: &self.slot })
+            .finish()
+    }
+}
 
 /// An error returned by [`LocalKey::try_with`](method@LocalKey::try_with).
 #[derive(Clone, Copy, Eq, PartialEq)]
@@ -275,3 +422,30 @@ impl fmt::Display for AccessError {
 }
 
 impl Error for AccessError {}
+
+enum ScopeInnerErr {
+    BorrowError,
+    AccessError,
+}
+
+impl ScopeInnerErr {
+    #[track_caller]
+    fn panic(&self) -> ! {
+        match self {
+            Self::BorrowError => panic!("cannot enter a task-local scope while the task-local storage is borrowed"),
+            Self::AccessError => panic!("cannot enter a task-local scope during or after destruction of the underlying thread-local"),
+        }
+    }
+}
+
+impl From<std::cell::BorrowMutError> for ScopeInnerErr {
+    fn from(_: std::cell::BorrowMutError) -> Self {
+        Self::BorrowError
+    }
+}
+
+impl From<std::thread::AccessError> for ScopeInnerErr {
+    fn from(_: std::thread::AccessError) -> Self {
+        Self::AccessError
+    }
+}