nuclei 0.5.4

//! Lingers: the work a handle leaves as it drops, which must be done whatever
//! happens to the handle and to the runtime it ran on: the bytes a socket's
//! write-behind holds, a file's writes behind, a gone shard's sources that
//! still hold bytes. Each is bounded by its own deadline (the runtime's
//! [`write_behind_linger`](crate::runtime::Builder::write_behind_linger), see
//! [`time`]) and done by the runtime the thread that hands it is inside (or
//! the process default runtime, if it runs):
//!
//! - while one of its workers takes tasks, as a task there;
//! - with none, on the linger thread, inside the runtime there and counted
//!   by it until it ends: never on the handing thread's own shard, which runs
//!   only at that thread's next `block_on`, if one ever comes
//!   (proofs/tla/going/Unworked.tla);
//! - once the runtime's drop began, in the runtime's list, which the drop
//!   runs on its own thread before it returns.
//!
//! A linger's task dropped before it ended (its shard stopped as the runtime
//! drops) hands what it had left back, so it goes on where it can. A task
//! that waits on a completion as its shard goes (a file's write a ring's
//! keeper or the blocking pool finishes) is dropped only once that
//! completion wakes it, on the completing thread, inside its runtime still
//! (`Sched::drop_future`): the linger's task hands its work back then, and the
//! task that holds a handle hands that handle's linger. So the runtime
//! counts every linger handed to it until it ended, and its drop waits,
//! running its list meanwhile and bounded by the linger's time, until none
//! is left, no file write its workers issued is still to wake its task
//! (`fs::call::WriteCall`), and no task woken on a worker's stopped shard is
//! still to be dropped. One handed where no runtime runs, or after a
//! runtime's drop is over, runs on the linger thread: one for the process,
//! started at the first linger it is handed, with a runtime of its own that
//! runs every linger as a task. No runtime is started behind the user's back
//! (proofs/tla/going/Linger.tla).
//!
//! A runtime that is never dropped (the process default one, or one kept for
//! good) never runs that wait, and the process's exit kills its workers with
//! their lingers. So a file's linger ([`hand_file`]) is counted from its hand
//! to its end, wherever it runs ([`Tracked`]), and the first one hands the
//! process an `atexit` hook ([`wait_at_exit`]) that waits until none is out,
//! for at most the default runtime's linger time, and says what it gave up.

use core::future::Future;
use core::pin::Pin;
use core::sync::atomic::{AtomicBool as ExitFlag, AtomicUsize as ExitCount};
use core::task::{Context, Poll};
use std::sync::{Arc, Condvar, Mutex, Once, PoisonError, Weak};

use kovan_queue::seg_queue::SegQueue;

use super::runtime::Inner;
use crate::sync::atomic_waker::AtomicWaker;
use crate::sync::internal::{AtomicBool, AtomicUsize, Ordering, fence};

/// How long a linger may take: the `write_behind_linger` of the runtime this
/// thread is inside (where the linger runs), or the default.
pub(crate) fn time() -> std::time::Duration {
    super::runtime::io_knobs::current().linger
}

/// A linger's work, boxed so it moves between threads and runtimes.
pub(crate) type Work = Pin<Box<dyn Future<Output = ()> + Send>>;

/// Hands `work` to the runtime this thread is inside, or to the default
/// runtime if it runs; with neither, to a thread of its own.
pub(crate) fn hand<F>(work: F)
where
    F: Future<Output = ()> + Send + 'static,
{
    let work: Work = Box::pin(work);
    match super::runtime::running_inner() {
        Some(inner) => inner.take_linger(work),
        None => alone(work),
    }
}

/// A file's linger: as [`hand`], counted until it ends, which the process's
/// exit waits for.
pub(crate) fn hand_file<F>(work: F)
where
    F: Future<Output = ()> + Send + 'static,
{
    HOOKED.call_once(|| {
        // SAFETY: `wait_at_exit` is a plain function that never unwinds.
        if unsafe { atexit(wait_at_exit) } != 0 {
            log::warn!(
                "the process could not be set to wait at its exit for files dropped with bytes written behind"
            );
        }
    });
    FILES_OUT.fetch_add(1, Ordering::SeqCst);
    hand(Tracked(work));
}

unsafe extern "C" {
    /// The C library's `atexit`: on unix through libc, on Windows the CRT's.
    fn atexit(callback: extern "C" fn()) -> core::ffi::c_int;
}

/// The exit's hook is handed once.
static HOOKED: Once = Once::new();
/// Files' lingers handed and not ended, whichever runtime or thread runs
/// them.
static FILES_OUT: ExitCount = ExitCount::new(0);
/// The exit's wait began: the last linger to end wakes it.
static EXITING: ExitFlag = ExitFlag::new(false);
/// The exit's wait, and its wake; plain std, which needs no thread-local of
/// the exiting thread (glibc drops those before the `atexit` hooks run).
static EXIT_LOCK: Mutex<()> = Mutex::new(());
static EXIT_WAKE: Condvar = Condvar::new();

/// A file's linger, counted until it is dropped, done or not: moved between
/// runtimes and threads it stays counted, and a linger nobody will run ends
/// its count where it is dropped.
struct Tracked<F>(F);

impl<F> Future for Tracked<F>
where
    F: Future<Output = ()>,
{
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        // SAFETY: the future is never moved out of its pin.
        unsafe { self.map_unchecked_mut(|this| &mut this.0) }.poll(cx)
    }
}

impl<F> Drop for Tracked<F> {
    fn drop(&mut self) {
        // SeqCst with the exit's: either the exit sees naught, or this sees
        // the exit waiting, and wakes it under its lock.
        if FILES_OUT.fetch_sub(1, Ordering::SeqCst) == 1 && EXITING.load(Ordering::SeqCst) {
            let _held = EXIT_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
            EXIT_WAKE.notify_all();
        }
    }
}

/// The process's `atexit` hook: waits until no file's linger is out, for at
/// most the default runtime's linger time (the lingers run on their own
/// threads meanwhile: `exit` stops none before its hooks), and says what is
/// still out then (proofs/tla/going/Linger.tla, `M_wait`).
extern "C" fn wait_at_exit() {
    let bound = super::runtime::exit_linger();
    let mut held = EXIT_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
    EXITING.store(true, Ordering::SeqCst);
    let deadline = std::time::Instant::now() + bound;
    loop {
        let out = FILES_OUT.load(Ordering::SeqCst);
        if out == 0 {
            return;
        }
        let now = std::time::Instant::now();
        if now >= deadline {
            crate::rt::counters::write_behind_lost(out as u64);
            log::warn!(
                "{out} dropped files still had bytes written behind when the process exited, after waiting {bound:?}: those bytes were not written"
            );
            return;
        }
        held = EXIT_WAKE
            .wait_timeout(held, deadline - now)
            .unwrap_or_else(PoisonError::into_inner)
            .0;
    }
}

/// A runtime's lingers: the count of those not ended, and the list of those
/// no worker of it takes.
pub(crate) struct Lingering {
    list: SegQueue<Work>,
    /// Lingers handed to the runtime that have not ended: running as tasks,
    /// held by a task that waits to hand them back, or in the list.
    active: AtomicUsize,
    /// The drop's run, woken as a linger comes to the list or ends.
    waker: AtomicWaker,
    /// The runtime's drop began: lingers wait in the list from now on.
    closing: AtomicBool,
    /// The drop is over: one handed later runs alone.
    over: AtomicBool,
}

impl Lingering {
    pub(crate) fn new() -> Self {
        Self {
            list: SegQueue::new(),
            active: AtomicUsize::new(0),
            waker: AtomicWaker::new(),
            closing: AtomicBool::new(false),
            over: AtomicBool::new(false),
        }
    }

    /// Lingers handed to the runtime that have not ended.
    pub(crate) fn active(&self) -> usize {
        self.active.load(Ordering::SeqCst)
    }

    /// A linger of the runtime ended (or went to a thread of its own).
    fn ended(&self) {
        self.active.fetch_sub(1, Ordering::SeqCst);
        self.waker.wake();
    }

    /// Wakes the drop's run to look again (a worker's shard settled).
    pub(crate) fn nudge(&self) {
        self.waker.wake();
    }

    /// The runtime's drop begins: its workers take no linger from now on.
    pub(crate) fn close(&self) {
        self.closing.store(true, Ordering::SeqCst);
    }

    pub(crate) fn is_closing(&self) -> bool {
        self.closing.load(Ordering::SeqCst)
    }

    /// The drop's run is over (or was given up): lingers go alone.
    pub(crate) fn is_over(&self) -> bool {
        self.over.load(Ordering::SeqCst)
    }

    /// Puts `work` in the list; one that finds the drop's run over takes it
    /// back out and runs it alone, as does the drop's last look.
    pub(crate) fn push(&self, work: Work) {
        self.list.push(work);
        self.waker.wake();
        // Pairs with the fence of `end`: either this sees the run over, or
        // the drop's last look finds the work.
        fence(Ordering::SeqCst);
        if self.over.load(Ordering::SeqCst) {
            self.run_rest_alone();
        }
    }

    /// The drop is over: what is left, and what comes after, runs alone.
    pub(crate) fn end(&self) {
        self.over.store(true, Ordering::SeqCst);
        fence(Ordering::SeqCst);
        self.run_rest_alone();
    }

    fn run_rest_alone(&self) {
        while let Ok(work) = self.list.try_pop() {
            self.ended();
            alone(work);
        }
    }
}

/// The lingers of a runtime's list, run together until the runtime's are
/// all done ([`Inner::lingers_done`]).
pub(crate) struct RunAll<'a> {
    inner: &'a Inner,
    running: Vec<Work>,
}

impl Future for RunAll<'_> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let this = self.get_mut();
        let lingering = this.inner.lingering();
        // Before the looks: a linger handed back or ended after them wakes it.
        lingering.waker.register(cx.waker());
        loop {
            while let Ok(work) = lingering.list.try_pop() {
                this.running.push(work);
            }
            this.running.retain_mut(|work| {
                let ended = work.as_mut().poll(cx).is_ready();
                if ended {
                    lingering.ended();
                }
                !ended
            });
            // A linger that ended may have handed another (a drop in it).
            if lingering.list.is_empty() {
                break;
            }
        }
        if this.inner.lingers_done() {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

impl Drop for RunAll<'_> {
    /// The drop's own bound ran out with lingers still running: each goes on
    /// on a thread of its own, to its own deadline, as the ones the run never
    /// took do (`Lingering::end`).
    fn drop(&mut self) {
        let lingering = self.inner.lingering();
        for work in self.running.drain(..) {
            lingering.ended();
            alone(work);
        }
    }
}

impl Inner {
    /// Whether nothing is left that may hand the runtime a linger, nor a
    /// linger it took not ended: each worker's shard settled (looked at
    /// first: a task its look waits for to drop hands its linger before),
    /// and no linger active.
    pub(crate) fn lingers_done(&self) -> bool {
        self.workers.shards().all(|shard| shard.settled()) && self.lingering().active() == 0
    }

    /// Runs the runtime's list until [`lingers_done`](Self::lingers_done):
    /// a future the drop runs on its own thread.
    pub(crate) fn run_lingers(&self) -> RunAll<'_> {
        RunAll {
            inner: self,
            running: Vec::new(),
        }
    }

    /// Takes a new linger, counted until it ends: as a task of a worker
    /// that takes tasks, else on this thread's shard, until the runtime's
    /// drop began; from then on into the list the drop runs.
    pub(crate) fn take_linger(self: &Arc<Self>, work: Work) {
        self.lingering().active.fetch_add(1, Ordering::SeqCst);
        self.route_linger(work);
    }

    /// Where a linger of this runtime goes now: a worker, or the linger
    /// thread, until the runtime's drop began; the list from then on.
    fn route_linger(self: &Arc<Self>, work: Work) {
        let lingering = self.lingering();
        if !lingering.is_closing() {
            let guarded = Guarded {
                work: Some(work),
                owner: Arc::downgrade(self),
            };
            return super::runtime::spawn_linger(self, guarded);
        }
        lingering.push(work);
    }
}

/// A linger run as a task: dropped before it ended, it hands what it left
/// back to its runtime (or to a thread of its own, the runtime gone).
pub(crate) struct Guarded {
    work: Option<Work>,
    owner: Weak<Inner>,
}

impl Future for Guarded {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let this = self.get_mut();
        let Some(work) = this.work.as_mut() else {
            return Poll::Ready(());
        };
        let poll = work.as_mut().poll(cx);
        if poll.is_ready() {
            this.work = None;
            if let Some(inner) = this.owner.upgrade() {
                inner.lingering().ended();
            }
        }
        poll
    }
}

impl Drop for Guarded {
    fn drop(&mut self) {
        let Some(work) = self.work.take() else {
            return;
        };
        // Never on this thread's shard (it may be the one going): a worker
        // that takes tasks, the linger thread, or the list.
        match self.owner.upgrade() {
            Some(inner) => inner.route_linger(work),
            None => alone(work),
        }
    }
}

/// Runs `linger`, one no worker of its runtime takes, on the linger thread:
/// polled inside its runtime there (its clock, its blocking pool for a file's
/// calls), and counted by it until it ends, so a drop of that runtime waits
/// for it (proofs/tla/going/Unworked.tla, `DropWaits`).
pub(crate) fn on_thread(linger: Guarded) {
    alone(Box::pin(OnThread(linger)));
}

/// A linger of a runtime, run on the linger thread with that runtime set.
struct OnThread(Guarded);

impl Future for OnThread {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let guarded = &mut self.get_mut().0;
        let owner = guarded.owner.upgrade();
        let _runtime = owner.as_ref().map(super::context::set_runtime);
        Pin::new(guarded).poll(cx)
    }
}

/// Runs `work` on the linger thread; where that thread cannot start, on
/// this one, as a last resort said each time.
pub(crate) fn alone(work: Work) {
    let unsent = match linger_thread() {
        Some(thread) => thread.send(work).err().map(|unsent| unsent.0),
        None => Some(work),
    };
    if let Some(work) = unsent {
        run_with_own_runtime(work);
    }
}

/// The linger thread's inbox, the thread started now if it is not running:
/// one thread for the process, whose runtime (no worker) runs every work it
/// is sent as a task of its shard, inside a `block_on` that never returns;
/// `None` where no thread can start.
fn linger_thread() -> Option<&'static crate::sync::mpsc::UnboundedSender<Work>> {
    static THREAD: std::sync::OnceLock<crate::sync::mpsc::UnboundedSender<Work>> =
        std::sync::OnceLock::new();
    if let Some(thread) = THREAD.get() {
        return Some(thread);
    }
    let (send, mut works) = crate::sync::mpsc::unbounded_channel::<Work>();
    let started = std::thread::Builder::new()
        .name("nuclei-linger".into())
        .spawn(move || {
            run_with_own_runtime(Box::pin(async move {
                while let Some(work) = works.recv().await {
                    crate::spawn(work).detach();
                }
            }));
        });
    match started {
        // Two hand-offs that started a thread at once: the one set wins, and
        // the other's ends as its inbox's sender goes, here.
        Ok(_) => {
            let _ = THREAD.set(send);
            THREAD.get()
        }
        Err(err) => {
            log::warn!(
                "no thread could start to deliver the bytes a dropped handle left ({err}): they are delivered on the thread that dropped it"
            );
            None
        }
    }
}

fn run_with_own_runtime(work: Work) {
    match crate::Runtime::builder().worker_threads(0).build() {
        Ok(rt) => rt.block_on(work),
        Err(err) => log::error!(
            "the bytes a dropped handle left were not delivered: no runtime could start for them ({err})"
        ),
    }
}