nuclei 0.5.4

//! A shard that stops running: for now (the outermost `block_on` of its
//! thread returned) or for good (dropped). What it watches goes, each to be
//! watched where its user polls next: the sources, each told; the
//! children, each once its registration ended. Its `shard_local!` values go
//! once its operations are settled, each in a catch of its own. A port hands the operations
//! still out to a keeper; so does a ring another thread can reap, as its
//! shard is dropped (`ring::keeper`), and any other ring waits for them a
//! bounded time. A source whose writes left bytes behind that the kernel has
//! not taken goes to a linger that writes them out where it runs, for no
//! task of the handle's may poll that side again (`tla/going/Behind.tla`),
//! and which a runtime that drops waits for (`rt::linger`).

use core::future::poll_fn;
use core::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};

#[cfg(ring)]
use super::Driver;
use super::Shard;
use super::backoff::Backoff;
use crate::io::Source;

impl Shard {
    /// Runs `future` to its end on this shard, a thread's own, entered for
    /// the outermost `block_on` of the thread; then leaves, whether the
    /// call returned or unwound, so what waits on this shard is handed on.
    pub(crate) fn block_on_then_leave<F: core::future::Future>(&self, future: F) -> F::Output {
        let mut future = core::pin::pin!(future);
        crate::rt::context::deeper(self.limits.max_nesting, || {
            let signal = self.signal();
            let output = self.run_then_leave(future.as_mut(), &signal);
            self.keep_signal(signal);
            output
        })
    }

    /// The outermost `block_on` of a thread outside the runtime returned,
    /// and this, its own shard, runs again only at its next one.
    #[inline]
    pub(crate) fn leave(&self) {
        // A shard that never registered a source since it last left has
        // none to let go of.
        if !self.sources.borrow().is_untouched() {
            self.let_go_of_sources();
        }
        if self.nothing_else_to_let_go() {
            return;
        }
        self.let_go_of_the_rest();
    }

    /// Whether a shard that leaves holds nothing more to let go of: on a
    /// ring, no operation out (a watched exit is one) and none posted to
    /// apply; elsewhere no child watched, and no port to hand on.
    #[inline(always)]
    fn nothing_else_to_let_go(&self) -> bool {
        #[cfg(ring)]
        if let Driver::Ring(ring) = &self.driver {
            return ring.in_flight() == 0 && self.shared.io_ops.is_empty();
        }
        !cfg!(port) && self.children.borrow().is_untouched()
    }

    /// Lets go of the sources registered here since the shard last left.
    #[inline(never)]
    fn let_go_of_sources(&self) {
        let gone: Vec<(u64, Arc<Source>)> = self.sources.borrow_mut().drain_tokens().collect();
        evict_all(&gone, self.shared_addr());
        give_up_claims(&gone, &self.shared);
        // On a worker, or with none on the linger thread; a runtime that
        // drops waits for it (`rt::linger`).
        drain_unsent(&gone, |source| crate::rt::linger::hand(drain(source)));
    }

    /// What [`leave`](Self::leave) lets go of past the sources: the
    /// children, and the operations still out.
    #[inline(never)]
    fn let_go_of_the_rest(&self) {
        self.let_go_of_children();
        // Sources that went before (a dropped handle's) may have left a
        // cancel queued, or their operations still out: the settle's first
        // step applies the one, and with none of the other it is over.
        #[cfg(ring)]
        if let Driver::Ring(ring) = &self.driver {
            if self.shared.io_ops.holds_any() {
                self.apply_io_ops();
            }
            if ring.in_flight() != 0 {
                ring.settle(|| self.apply_io_ops());
            }
        }
        // The port's go to its keeper.
        #[cfg(port)]
        if !self.hand_off_port() {
            self.settle_port();
        }
    }
}

/// Drops each shard-local value in a catch of its own: one that panics is
/// logged, and the rest of the values, and of the shard's going, go on.
fn drop_locals(values: Vec<Box<dyn core::any::Any + Send>>) {
    for value in values {
        if let Some(panic) = crate::rt::caught::catch(|| drop(value)) {
            crate::rt::caught::lost(
                format_args!("a shard-local value panicked as its worker stopped"),
                panic,
            );
        }
    }
}

/// Evicts each source of `gone` (registered under its token with the shard
/// at `home`, this one), and waits until each one's registration here has
/// gone. A source a side of which another thread holds is evicted by that
/// holder as it lets go, which comes a driver call later; until then the
/// holder may still hand this shard's driver an operation, so the
/// operations out here are only known once every eviction happened.
fn evict_all(gone: &[(u64, Arc<Source>)], home: usize) {
    for (_, source) in gone {
        source.evict();
    }
    wait_for_evictions(gone, home);
}

/// The sides of each source of `gone` that this shard, `home`, claimed give the claim up, once
/// the shard did all it does with them as their owner (`Source::give_up_claims`).
fn give_up_claims(gone: &[(u64, Arc<Source>)], home: &Arc<crate::rt::shard::Shared>) {
    for (_, source) in gone {
        source.give_up_claims(home);
    }
}

/// Hands each source of `gone`, evicted, whose write side holds bytes the
/// kernel has not taken to `spawn`, for a task that writes them out.
fn drain_unsent(gone: &[(u64, Arc<Source>)], spawn: impl Fn(Arc<Source>)) {
    gone.iter()
        .filter(|(_, source)| source.has_unsent())
        .for_each(|(_, source)| spawn(Arc::clone(source)));
}

/// Writes out what the write side of `source` holds, on the shard it runs on,
/// for at most the linger a dropped handle's bytes get.
async fn drain(source: Arc<Source>) {
    let _ = crate::time::timeout(
        crate::rt::linger::time(),
        poll_fn(|cx| source.poll_drain(cx)),
    )
    .await;
}

/// How long a stopping shard waits for the holders of its sources' sides to let go: far past the
/// driver call after which a holder that lets go has done it.
const EVICTION_WAIT: Duration = Duration::from_secs(1);
/// Waits until no source of `gone` is registered under its token with the
/// shard at `home`, nor has an attach to it still asked (a holder performs
/// it as it lets go, and posts here what it leaves: a ring's install of the
/// file with the lease that keeps the handle's object open, a release of the
/// slot). The ask is looked at first: a ring's attach names its shard and
/// token before it takes the ask (`proofs/tla/going/Pending.tla`). A source
/// evicted from here may register with another shard at once, under a token
/// of the same number: that registration is not waited for, or this shard
/// would wait for as long as it lasts, with the eviction's cancel of the
/// source's operations still in its mailbox.
///
/// The wait is bounded by [`EVICTION_WAIT`], from the first look that has to
/// wait. Past it the shard stops waiting for every source left, says so and
/// counts it: the registration of such a source ends when the holder that
/// did not let go does (`Pending.tla`, `AGiveUp`, `WaitEnds`).
fn wait_for_evictions(gone: &[(u64, Arc<Source>)], home: usize) {
    wait_for_evictions_within(gone, home, EVICTION_WAIT);
}

/// [`wait_for_evictions`] with its bound given.
fn wait_for_evictions_within(gone: &[(u64, Arc<Source>)], home: usize, within: Duration) {
    // Made at the first look that has to wait: a path that finds nothing to wait for reads no clock.
    let mut waited: Option<(Instant, Backoff)> = None;
    for (token, source) in gone {
        loop {
            let asked = source.attach_asked_of(home);
            if !asked && !source.registered_with(home, *token) {
                break;
            }
            #[cfg(test)]
            if asked {
                crate::io::source::ATTACH_WAITS.fetch_add(1, Ordering::SeqCst);
            }
            let (deadline, backoff) =
                waited.get_or_insert_with(|| (Instant::now() + within, Backoff::new()));
            if !backoff.yielding() && Instant::now() >= *deadline {
                crate::rt::counters::eviction_given_up();
                log::warn!(
                    "a worker stopped before another thread let go of one of its sockets or files; its registration ends when that thread does"
                );
                return;
            } else {
                backoff.pause();
            }
        }
    }
}

#[cfg(all(test, unix, not(loom)))]
mod tests {
    use core::future::{Future, poll_fn};
    use core::task::Poll;
    use std::sync::Arc;
    use std::time::Duration;

    /// A shard that goes waits for its own registrations of its sources
    /// only: a source evicted from it that registered with another shard
    /// since, under a token of the same number (each shard numbers its
    /// own), is not waited for.
    #[test]
    fn a_shard_that_goes_does_not_wait_for_a_registration_elsewhere_under_its_token() {
        let (mine, _peer) = std::os::unix::net::UnixStream::pair().unwrap();
        let handle = crate::io::Handle::new(mine).unwrap();
        let src = Arc::clone(handle.source());
        crate::block_on(async {
            // A receive polled here registers the source with this thread's
            // shard.
            let mut buf = [0u8; 1];
            {
                let mut read = Box::pin(handle.recv(&mut buf));
                let polled = poll_fn(|cx| Poll::Ready(read.as_mut().poll(cx))).await;
                assert!(polled.is_pending());
            }
            let (home, token) = src.registration();
            assert_ne!(home, 0, "the receive did not register the source");
            // Another shard evicted the source under the same token number,
            // and waits for that eviction.
            let (done_tx, done_rx) = std::sync::mpsc::channel();
            let other = {
                let src = Arc::clone(&src);
                std::thread::spawn(move || {
                    crate::block_on(async move {
                        let there =
                            crate::rt::context::current(|s| s.map(super::Shard::shared_addr))
                                .expect("a block_on runs a shard");
                        assert_ne!(there, home);
                        super::wait_for_evictions(&[(token, src)], there);
                    });
                    let _ = done_tx.send(());
                })
            };
            assert!(
                done_rx.recv_timeout(Duration::from_secs(10)).is_ok(),
                "the other shard waited for this shard's registration"
            );
            other.join().unwrap();
            assert_eq!(
                src.registration(),
                (home, token),
                "the registration here moved"
            );
        });
        drop(handle);
    }

    /// A shard that goes stops waiting for a holder that does not let go,
    /// past the bound, and counts it (proofs/tla/going/Pending.tla,
    /// `AGiveUp`).
    #[test]
    fn a_shard_that_goes_gives_up_waiting_for_a_holder_that_does_not_let_go() {
        let (mine, _peer) = std::os::unix::net::UnixStream::pair().unwrap();
        let handle = crate::io::Handle::new(mine).unwrap();
        let src = Arc::clone(handle.source());
        crate::block_on(async {
            let mut buf = [0u8; 1];
            {
                let mut read = Box::pin(handle.recv(&mut buf));
                let polled = poll_fn(|cx| Poll::Ready(read.as_mut().poll(cx))).await;
                assert!(polled.is_pending());
            }
            let (home, token) = src.registration();
            assert_ne!(home, 0, "the receive did not register the source");
            // Nobody evicts the source: the wait ends at its bound, on a thread with a shard of
            // its own to count on.
            let (done_tx, done_rx) = std::sync::mpsc::channel();
            let waiter = {
                let src = Arc::clone(&src);
                std::thread::spawn(move || {
                    crate::block_on(async move {
                        let given_up = || {
                            crate::rt::context::current(|s| {
                                s.expect("a block_on runs a shard")
                                    .shared
                                    .stats()
                                    .errors
                                    .evictions_given_up
                            })
                        };
                        let before = given_up();
                        super::wait_for_evictions_within(
                            &[(token, src)],
                            home,
                            Duration::from_millis(50),
                        );
                        let _ = done_tx.send((before, given_up()));
                    });
                })
            };
            let (before, after) = done_rx
                .recv_timeout(Duration::from_secs(10))
                .expect("the wait did not end at its bound");
            waiter.join().unwrap();
            assert_eq!(after, before + 1, "the give-up was not counted");
        });
        drop(handle);
    }
}

impl Drop for Shard {
    fn drop(&mut self) {
        // This thread runs the shard and no other does: what its sources are
        // asked as it goes is asked as their owner (`context::letting_go`).
        let shared = Arc::clone(&self.shared);
        crate::rt::context::letting_go(&shared, || self.go());
    }
}

impl Shard {
    /// [`Drop`]'s work: the shard goes for good.
    fn go(&mut self) {
        #[cfg(ring)]
        let started = std::time::Instant::now();
        self.shared.counters.went();
        // Every task still queued here is dropped on this thread, the one its
        // future may be dropped on.
        self.shared.draining.store(true, Ordering::SeqCst);
        self.shared.stopped.store(true, Ordering::SeqCst);
        // The waker a worker's loop leaves when its look after registering
        // finds it may exit: it holds the shard's `Shared`, which holds it.
        drop(self.shared.stop_waker.take());
        // The sources registered here lose their driver: each is told, and
        // registers where it is polled next.
        let gone: Vec<(u64, Arc<Source>)> = self.sources.get_mut().drain_tokens().collect();
        evict_all(&gone, self.shared_addr());
        // A worker that takes tasks drains them, or with none the linger
        // thread; once the runtime's drop began, its drop does (`rt::linger`).
        drain_unsent(&gone, |source| crate::rt::linger::hand(drain(source)));
        self.let_go_of_children();
        // The ring's operations still out are cancelled, and their completions waited for or
        // left to a keeper, so no memory they point at goes while the kernel may write it.
        #[cfg(ring)]
        if let Driver::Ring(ring) = &self.driver {
            if ring.reapable() {
                // Cancels a holder of a source on another thread left here; the keeper applies
                // what comes later.
                self.apply_io_ops();
                ring.cancel_all();
            } else {
                ring.settle(|| self.apply_io_ops());
                self.pool
                    .unreserve(ring.shut_down_from(started, || self.apply_io_ops()));
            }
        }
        // The port's operations to a keeper, unless one has them already.
        #[cfg(port)]
        if !self.port_kept() && !self.hand_off_port() {
            self.settle_port();
        }
        // The shard-local values go once the operations are settled, so their
        // drops, which are user code, run while nothing the kernel may still
        // write into hangs on the rest of this going; and before the tasks, so a
        // task one wakes is dropped below. A handle's sources were let go of as
        // any other's (`tla/going/Locals.tla`).
        drop_locals(self.locals.take());
        // No timer fires here again: the arm requests never applied orphan
        // their timers now, and those in the wheel as the loop below lets go
        // of them (`sleep::Armed`), so each sleep arms again where it is
        // polled, or its task, woken here, goes too.
        self.shared.drain_timer_ops();
        // What goes with a task may wake another of the shard's (a timer's
        // waker, a channel its future held): that one comes back through the
        // inbox, and is dropped here too, until none is left.
        loop {
            let lifo = self.lifo.take();
            let queue = core::mem::take(&mut *self.queue.borrow_mut());
            let timers = self.wheel.borrow_mut().take_wakers();
            let precise = self.precise.borrow_mut().take_wakers();
            let fired = core::mem::take(&mut *self.fired.borrow_mut());
            let mut quiet = lifo.is_none()
                && queue.is_empty()
                && timers.is_empty()
                && precise.is_empty()
                && fired.is_empty();
            drop((lifo, queue, timers, precise, fired));
            while let Some(runnable) = self.pop_inbox() {
                drop(runnable);
                quiet = false;
            }
            // SAFETY: this thread is the shard's last holder, the owner of its list, and of its
            // home's when the shard is a thread's own.
            quiet &= unsafe { self.shared.tasks.cancel_all() } == 0;
            if let Some(home) = self.home.get() {
                // SAFETY: as above.
                quiet &= unsafe { home.cancel_listed() } == 0;
            }
            if quiet {
                break;
            }
        }
        // Bound tasks of this thread whose wake is still on its way are waited for, to a bound.
        if let Some(home) = self.home.get()
            // SAFETY: this thread is the home's.
            && !unsafe { home.is_idle() }
        {
            // SAFETY: as above.
            unsafe {
                home.cancel_all(
                    Some(self),
                    Instant::now() + crate::rt::workers::TEARDOWN_WAIT,
                );
            }
        }
        // Every task is off the list before the inbox's role goes: `release_inbox` lets a pusher
        // drain and drop on its own thread, and the runnables still on their way then find their
        // tasks on no list. One still in the inbox is dropped by the drain `release_inbox` does,
        // on this thread, its task already off the list.
        // SAFETY: as above.
        unsafe { self.shared.tasks.forget_all() };
        // The inbox's taker's role goes, looking at the inbox once more: a
        // pusher that saw the flag while the role was held left its
        // runnable to this drain.
        self.shared.release_inbox();
        // Done with the sources as their owner (the dropped tasks' closes included): what it
        // claimed of them goes.
        give_up_claims(&gone, &self.shared);
        // The mailbox closes, unless a keeper applies it from now on (and
        // closes it once it is done).
        #[cfg(ring)]
        let kept = self.keep_ring();
        #[cfg(not(ring))]
        let kept = false;
        if !kept {
            self.shared.close_io();
        }
    }
}

impl Shard {
    /// The shard is dropped: a ring another thread can reap goes to a keeper with what is still
    /// out once the ready completions are applied. Where no keeper can be had, the shard waits
    /// for them a bounded time, and what is still out keeps its memory for good.
    /// Whether the ring went to a keeper, which applies the shard's mailbox
    /// from then on.
    #[cfg(ring)]
    fn keep_ring(&mut self) -> bool {
        // A ring no keeper can reap was shut down as the shard went.
        let Driver::Ring(ring) = &self.driver else {
            return false;
        };
        if !ring.reapable() {
            return false;
        }
        // A look that fails leaves the ready completions to the keeper, whose wait meets the
        // same failure and keeps what is out for good.
        let looked = ring.wait(None, Some(std::time::Duration::ZERO)).is_ok();
        if looked && ring.in_flight() == 0 {
            return false;
        }
        let Driver::Ring(ring) = core::mem::replace(&mut self.driver, Driver::Gone) else {
            return false;
        };
        super::ring::keeper::hand_off(ring, Arc::clone(&self.shared));
        true
    }
}