some_global_executor 0.1.6

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

//! Drain notification and the awaitable drain future.
//!
//! `DrainNotify` tracks how many tasks are in flight and wakes anyone waiting
//! for the count to reach zero. `ExecutorDrain` is the `Future` that an async
//! caller awaits to be notified of that moment. Keeping them together makes the
//! ordering of the register-then-check in `ExecutorDrain::poll` and the
//! release-then-wake in `DrainNotify::task_finished` easy to audit.

use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::AtomicUsize;
use std::sync::{Condvar, Mutex};
use std::task::Waker;

use crate::Executor;

/// A future that completes when all tasks in an executor have finished.
///
/// `ExecutorDrain` is returned by [`Executor::drain_async()`] and implements
/// `Future<Output = ()>`. It polls the executor's internal state to determine
/// when all tasks have completed execution.
///
/// # Examples
///
/// ```
/// # // This example shows async draining but requires an async runtime
/// # // which is not available in doctests, so we use synchronous drain
/// use some_global_executor::Executor;
/// use some_executor::SomeExecutor;
/// use some_executor::task::{Task, Configuration};
/// # if cfg!(target_arch = "wasm32") { return; }
///
/// let mut executor = Executor::new("async-drain".to_string(), 2);
///
/// // Spawn some work
/// let task = Task::without_notifications(
///     "work".to_string(),
///     Configuration::default(),
///     async {
///         // Simulate some work
///         42
///     }
/// );
/// executor.spawn(task);
///
/// // In async context, you would use:
/// // executor.drain_async().await;
/// // Here we use synchronous drain for the example
/// executor.drain();
/// ```
#[derive(Debug)]
pub struct ExecutorDrain {
    pub(crate) executor: Executor,
    pub(crate) waiter_id: Option<usize>,
}

impl Future for ExecutorDrain {
    type Output = ();

    fn poll(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        let this = self.get_mut();
        //need to register prior to check
        let drain_notify = this.executor.imp.drain_notify();
        drain_notify.register(&mut this.waiter_id, cx.waker());
        //Acquire pairs with the Release decrement on task completion, so task side
        //effects are visible once we observe the count at zero
        let running_tasks = drain_notify
            .running_tasks
            .load(std::sync::atomic::Ordering::Acquire);
        #[cfg(feature = "logwise-forensic")]
        logwise::forensic!(
            "some_global_executor.drain.polled",
            running_tasks = support(running_tasks as u64),
        );
        if running_tasks == 0 {
            if let Some(waiter_id) = this.waiter_id.take() {
                drain_notify.unregister(waiter_id);
            }
            std::task::Poll::Ready(())
        } else {
            std::task::Poll::Pending
        }
    }
}

/// Internal notification mechanism for tracking running tasks.
///
/// This structure maintains a count of running tasks and provides
/// a waker mechanism for notifying when all tasks complete.
#[derive(Debug)]
pub(crate) struct DrainNotify {
    pub(crate) running_tasks: AtomicUsize,
    pub(crate) next_waiter_id: AtomicUsize,
    pub(crate) waiters: Mutex<Vec<(usize, Waker)>>,
    pub(crate) completed: Condvar,
}

impl DrainNotify {
    pub(crate) fn new() -> Self {
        Self {
            running_tasks: AtomicUsize::new(0),
            next_waiter_id: AtomicUsize::new(0),
            waiters: Mutex::new(Vec::new()),
            completed: Condvar::new(),
        }
    }

    pub(crate) fn lock_waiters(&self) -> std::sync::MutexGuard<'_, Vec<(usize, Waker)>> {
        self.waiters
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    pub(crate) fn register(&self, waiter_id: &mut Option<usize>, waker: &Waker) {
        let waiter_id = *waiter_id.get_or_insert_with(|| {
            self.next_waiter_id
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
        });
        let mut waiters = self.lock_waiters();
        if let Some((_, registered)) = waiters.iter_mut().find(|(id, _)| *id == waiter_id) {
            if !registered.will_wake(waker) {
                *registered = waker.clone();
            }
        } else {
            waiters.push((waiter_id, waker.clone()));
        }
    }

    pub(crate) fn unregister(&self, waiter_id: usize) {
        self.lock_waiters().retain(|(id, _)| *id != waiter_id);
    }

    pub(crate) fn task_finished(&self) -> usize {
        let old = self
            .running_tasks
            .fetch_sub(1, std::sync::atomic::Ordering::Release);
        debug_assert!(old > 0, "running task counter underflow");
        if old == 1 {
            let waiters = {
                let mut registered = self.lock_waiters();
                self.completed.notify_all();
                std::mem::take(&mut *registered)
            };
            for (_, waker) in waiters {
                if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| waker.wake())).is_err()
                {
                    logwise::log!("Drain waker panicked; ignoring it to keep the executor alive");
                }
            }
        }
        old
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub(crate) fn wait_for_completion(&self) {
        let mut waiters = self.lock_waiters();
        while self
            .running_tasks
            .load(std::sync::atomic::Ordering::Acquire)
            != 0
        {
            waiters = self
                .completed
                .wait(waiters)
                .unwrap_or_else(std::sync::PoisonError::into_inner);
        }
    }
}

impl From<ExecutorDrain> for Executor {
    fn from(mut drain: ExecutorDrain) -> Self {
        if let Some(waiter_id) = drain.waiter_id.take() {
            drain.executor.imp.drain_notify().unregister(waiter_id);
        }
        drain.executor.clone()
    }
}

impl From<Executor> for ExecutorDrain {
    fn from(executor: Executor) -> Self {
        ExecutorDrain {
            executor,
            waiter_id: None,
        }
    }
}

impl Drop for ExecutorDrain {
    fn drop(&mut self) {
        if let Some(waiter_id) = self.waiter_id.take() {
            self.executor.imp.drain_notify().unregister(waiter_id);
        }
    }
}

impl AsRef<Executor> for ExecutorDrain {
    fn as_ref(&self) -> &Executor {
        &self.executor
    }
}

impl AsMut<Executor> for ExecutorDrain {
    fn as_mut(&mut self) -> &mut Executor {
        &mut self.executor
    }
}