some_executor 0.7.2

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Future implementations for task-local storage.

#[cfg(not(target_arch = "wasm32"))]
use std::cell::Cell;
use std::cell::RefCell;
#[cfg(not(target_arch = "wasm32"))]
use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use crate::context::{LocalKey, LocalKeyImmutable};

#[cfg(not(target_arch = "wasm32"))]
thread_local! {
    static ACTIVE_TASK_LOCAL_POLLS: Cell<usize> = const { Cell::new(0) };
    static DEFERRED_UNTIL_TASK_LOCALS_RESTORED: RefCell<VecDeque<Box<dyn FnOnce()>>> =
        RefCell::new(VecDeque::new());
}

#[cfg(not(target_arch = "wasm32"))]
fn enter_task_local_poll() {
    ACTIVE_TASK_LOCAL_POLLS.set(ACTIVE_TASK_LOCAL_POLLS.get() + 1);
}

#[cfg(not(target_arch = "wasm32"))]
fn leave_task_local_poll() {
    let depth = ACTIVE_TASK_LOCAL_POLLS.get();
    debug_assert!(depth > 0, "task-local poll depth underflow");
    ACTIVE_TASK_LOCAL_POLLS.set(depth - 1);
    if depth != 1 {
        return;
    }

    if std::thread::panicking() {
        DEFERRED_UNTIL_TASK_LOCALS_RESTORED.with_borrow_mut(VecDeque::clear);
        return;
    }

    loop {
        let deferred = DEFERRED_UNTIL_TASK_LOCALS_RESTORED.with_borrow_mut(VecDeque::pop_front);
        match deferred {
            Some(deferred) => deferred(),
            None => break,
        }
    }
}

/// Runs work immediately unless a task-local value is currently installed on
/// this thread. In that case, the work is deferred until the outermost scoped
/// poll restores every task-local value.
#[cfg(not(target_arch = "wasm32"))]
pub(crate) fn run_or_defer_until_task_locals_restored(work: impl FnOnce() + 'static) {
    if ACTIVE_TASK_LOCAL_POLLS.get() == 0 {
        work();
    } else {
        DEFERRED_UNTIL_TASK_LOCALS_RESTORED.with_borrow_mut(|deferred| {
            deferred.push_back(Box::new(work));
        });
    }
}

#[cfg(target_arch = "wasm32")]
pub(crate) fn run_or_defer_until_task_locals_restored(work: impl FnOnce() + 'static) {
    // The wasm fallback already schedules work asynchronously, so invoking it
    // here cannot poll the child reentrantly under the parent's task-locals.
    work();
}

/// Installs `slot`'s value into the thread-local `key` for the duration of
/// `future.poll(cx)`, then moves it back into the slot.
///
/// Whatever value was previously installed in the thread-local is saved and
/// restored afterwards, which is what makes nested `scope`s of the same key
/// shadow correctly. Restoration also runs on unwind, so a panicking future
/// cannot leak its value into the thread-local (where unrelated tasks on the
/// same thread would observe it).
fn poll_with_scoped_value<V: 'static, F: Future>(
    key: &'static std::thread::LocalKey<RefCell<Option<V>>>,
    slot: &mut Option<V>,
    future: Pin<&mut F>,
    cx: &mut Context<'_>,
) -> Poll<F::Output> {
    struct Guard<'a, V: 'static> {
        key: &'static std::thread::LocalKey<RefCell<Option<V>>>,
        old_value: Option<Option<V>>,
        slot: &'a mut Option<V>,
    }
    impl<V: 'static> Drop for Guard<'_, V> {
        fn drop(&mut self) {
            let value = self
                .key
                .replace(self.old_value.take().expect("Guard dropped twice"));
            *self.slot = value;
            #[cfg(not(target_arch = "wasm32"))]
            leave_task_local_poll();
        }
    }
    let value = slot.take().expect("No value in slot");
    let old_value = key.replace(Some(value));
    #[cfg(not(target_arch = "wasm32"))]
    enter_task_local_poll();
    let guard = Guard {
        key,
        old_value: Some(old_value),
        slot,
    };
    let r = future.poll(cx);
    drop(guard);
    r
}

/// A future that sets a value `T` of a task local for the future `F` during
/// its execution.
///
/// This future wraps another future and ensures that a task-local value is
/// available during the wrapped future's execution. The task-local value is
/// stored in a slot and swapped into thread-local storage each time the future
/// is polled, then swapped back out when polling completes.
///
/// The value of the task-local must be `'static` and will be dropped on the
/// completion of the future.
///
/// # Implementation Details
///
/// The future maintains the task-local value in a `slot` field. During each
/// poll:
/// 1. The value is moved from the slot into thread-local storage
/// 2. The wrapped future is polled
/// 3. The value is moved back from thread-local storage to the slot
///
/// This ensures the value is available to the future and any code it calls,
/// while keeping it safe from concurrent access.
///
/// Created by the function [`LocalKey::scope`](crate::context::LocalKey::scope).
#[derive(Debug)]
pub(crate) struct TaskLocalFuture<V: 'static, F> {
    pub(crate) slot: Option<V>,
    pub(crate) local_key: &'static LocalKey<V>,
    pub(crate) future: F,
}

/// A future that sets an immutable task-local value for the duration of
/// another future's execution.
///
/// Similar to [`TaskLocalFuture`], but for immutable task-locals that cannot
/// be modified once set within a scope. This provides stronger guarantees about
/// the stability of configuration values during task execution.
///
/// # Implementation Details
///
/// Works identically to `TaskLocalFuture` in terms of the slot/swap mechanism,
/// but the associated `LocalKeyImmutable` type prevents mutation of the value
/// while it's in scope.
///
/// Created by the function [`LocalKeyImmutable::scope`](crate::context::LocalKeyImmutable::scope).
#[derive(Debug)]
pub(crate) struct TaskLocalImmutableFuture<V: 'static, F> {
    pub(crate) slot: Option<V>,
    pub(crate) local_key: &'static LocalKeyImmutable<V>,
    pub(crate) future: F,
}

impl<V, F> TaskLocalFuture<V, F> {
    //     /**
    //     Gets access to the underlying value.
    // */
    //
    //     pub(crate) fn get_val<R>(&self, closure:impl FnOnce(&V) -> R) -> R  {
    //         match self.slot {
    //             Some(ref value) => closure(value),
    //             None => self.local_key.with(|value| {
    //                 closure(value.expect("Value neither in slot nor in thread-local"))
    //             })
    //         }
    //     }
    //     /**
    //     Gets mutable access to the underlying value.
    // */
    //     pub(crate) fn get_val_mut<R>(&mut self, closure:impl FnOnce(&mut V) -> R) -> R  {
    //         match self.slot {
    //             Some(ref mut value) => closure(value),
    //             None => self.local_key.with_mut(|value| {
    //                 closure(value.expect("Value neither in slot nor in thread-local"))
    //             })
    //         }
    //     }

    // pub(crate) fn get_future(&self) -> &F {
    //     &self.future
    // }
    //
    // pub (crate) fn get_future_mut(&mut self) -> &mut F {
    //     &mut self.future
    // }
    //
    // pub(crate) fn into_future(self) -> F {
    //     self.future
    // }
}

impl<V, F> TaskLocalImmutableFuture<V, F> {
    // /**
    // Gets access to the underlying value.
    // */
    //
    // pub(crate) fn get_val<R>(&self, closure:impl FnOnce(&V) -> R) -> R  {
    //     match self.slot {
    //         Some(ref value) => closure(value),
    //         None => self.local_key.with(|value| {
    //             closure(value.expect("Value neither in slot nor in thread-local"))
    //         })
    //     }
    // }

    // pub(crate) fn get_future(&self) -> &F {
    //     &self.future
    // }

    // pub (crate) fn get_future_mut(&mut self) -> &mut F {
    //     &mut self.future
    // }

    // pub(crate) fn into_future(self) -> F {
    //     self.future
    // }
}

/// Implementation of `Future` for `TaskLocalFuture`.
///
/// This implementation handles the mechanics of swapping task-local values
/// in and out of thread-local storage during polling.
impl<V, F> Future for TaskLocalFuture<V, F>
where
    V: Unpin,
    F: Future,
{
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        //destructure
        let (future, slot, local_key) = unsafe {
            let this = self.get_unchecked_mut();
            let future = Pin::new_unchecked(&mut this.future);
            (future, &mut this.slot, this.local_key)
        };

        poll_with_scoped_value(&local_key.0, slot, future, cx)
    }
}

/// Implementation of `Future` for `TaskLocalImmutableFuture`.
///
/// This implementation handles the mechanics of swapping immutable task-local
/// values in and out of thread-local storage during polling.
impl<V, F> Future for TaskLocalImmutableFuture<V, F>
where
    V: Unpin,
    F: Future,
{
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        //destructure
        let (future, slot, local_key) = unsafe {
            let this = self.get_unchecked_mut();
            let future = Pin::new_unchecked(&mut this.future);
            (future, &mut this.slot, this.local_key)
        };

        poll_with_scoped_value(&local_key.0, slot, future, cx)
    }
}

#[cfg(test)]
mod tests {
    crate::task_local! {
        static MUTABLE_LEVEL: u32;
        static const IMMUTABLE_LEVEL: u32;
        #[cfg(not(target_arch = "wasm32"))]
        static const REENTRANT_SECRET: u32;
    }

    #[wasm_lite::wasm_lite_test]
    async fn nested_scope_same_key_shadows() {
        MUTABLE_LEVEL
            .scope(1, async {
                assert_eq!(MUTABLE_LEVEL.get(), 1);
                MUTABLE_LEVEL
                    .scope(2, async {
                        assert_eq!(MUTABLE_LEVEL.get(), 2);
                    })
                    .await;
                assert_eq!(MUTABLE_LEVEL.get(), 1);
            })
            .await;
        MUTABLE_LEVEL.with(|v| assert!(v.is_none()));
    }

    #[wasm_lite::wasm_lite_test]
    async fn nested_scope_same_key_shadows_immutable() {
        IMMUTABLE_LEVEL
            .scope(1, async {
                assert_eq!(IMMUTABLE_LEVEL.get(), 1);
                IMMUTABLE_LEVEL
                    .scope(2, async {
                        assert_eq!(IMMUTABLE_LEVEL.get(), 2);
                    })
                    .await;
                assert_eq!(IMMUTABLE_LEVEL.get(), 1);
            })
            .await;
        IMMUTABLE_LEVEL.with(|v| assert!(v.is_none()));
    }

    #[wasm_lite::wasm_lite_test]
    async fn set_outside_scope_then_scope() {
        crate::task_local! {
            static OUTSIDE: u32;
        }
        OUTSIDE.set(7);
        OUTSIDE
            .scope(1, async {
                assert_eq!(OUTSIDE.get(), 1);
            })
            .await;
        // the pre-scope value is restored, not clobbered
        OUTSIDE.with(|v| assert_eq!(v.copied(), Some(7)));
    }

    #[cfg(not(target_arch = "wasm32"))]
    #[test]
    fn panic_in_scope_restores_thread_local() {
        crate::task_local! {
            static PANICKY: u32;
        }
        let result = std::panic::catch_unwind(|| {
            wasm_lite_std::block_on(PANICKY.scope(42, async {
                panic!("boom");
            }));
        });
        assert!(result.is_err());
        // the panicking scope must not leak its value into the thread-local
        PANICKY.with(|v| assert!(v.is_none()));
        // and a subsequent scope on the same thread must still work
        wasm_lite_std::block_on(PANICKY.scope(5, async {
            assert_eq!(PANICKY.get(), 5);
        }));
    }

    #[cfg(not(target_arch = "wasm32"))]
    #[test]
    fn reentrantly_spawned_task_does_not_inherit_task_local() {
        use crate::SomeStaticExecutor;
        use crate::observer::Observer;
        use crate::task::{Configuration, Task};
        use std::sync::Arc;
        use std::sync::atomic::{AtomicBool, Ordering};

        let child_saw_secret = Arc::new(AtomicBool::new(false));
        let child_saw_secret_clone = child_saw_secret.clone();

        wasm_lite_std::block_on(REENTRANT_SECRET.scope(42, async move {
            let child = Task::without_notifications(
                "reentrant child".to_string(),
                Configuration::default(),
                async move {
                    let inherited = REENTRANT_SECRET.with(|secret| secret.is_some());
                    child_saw_secret_clone.store(inherited, Ordering::Relaxed);
                },
            );
            let mut executor = crate::static_last_resort::StaticLastResortExecutor::new();
            executor.spawn_static(child).detach();
        }));

        assert!(
            !child_saw_secret.load(Ordering::Relaxed),
            "a separately spawned task inherited its parent's scoped task-local"
        );
    }
}