nuclei 0.5.4

use std::sync::Arc;
use std::task::{Wake, Waker};
use std::time::{Duration, Instant};

use super::{Limits, Shard};
use crate::rt::context;

/// A timer armed in the shard's wheel, done once it fired.
fn sleep_on(shard: &Shard, at: Instant) -> impl core::future::Future<Output = ()> + '_ {
    let mut key = None;
    core::future::poll_fn(move |cx| {
        let mut wheel = shard.wheel().borrow_mut();
        match key {
            None => {
                let armed = wheel.insert(
                    shard.time_of(false, crate::time::clock::nanos(at)),
                    cx.waker().clone(),
                );
                match armed {
                    Ok(k) => key = Some(k),
                    Err(_) => return core::task::Poll::Ready(()),
                }
            }
            Some(k) if wheel.is_fired(k) => {
                drop(wheel.remove(k));
                return core::task::Poll::Ready(());
            }
            Some(k) => drop(wheel.replace_waker(k, cx.waker().clone())),
        }
        let fired = key.is_some_and(|k| wheel.is_fired(k));
        if fired {
            core::task::Poll::Ready(())
        } else {
            core::task::Poll::Pending
        }
    })
}

#[test]
fn a_timer_fires_from_the_run_loop_and_never_early() {
    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let _current = context::enter(&shard);
    for ms in [1u64, 5, 30] {
        let start = Instant::now();
        let at = start + Duration::from_millis(ms);
        shard.block_on(sleep_on(&shard, at));
        assert!(Instant::now() >= at, "a {ms} ms timer fired early");
        assert!(start.elapsed() < Duration::from_secs(5));
    }
}

#[test]
fn the_loop_parks_until_the_next_timer_rather_than_forever() {
    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let _current = context::enter(&shard);
    let start = Instant::now();
    // Nothing but the timer can wake the loop: if it parked without a
    // deadline, this would never return.
    shard.block_on(sleep_on(&shard, start + Duration::from_millis(20)));
    assert!(start.elapsed() >= Duration::from_millis(20));
}

/// A task's timer comes due while the shard is on its way to park: the
/// drive that fires it finds the task it woke, and only looks at the
/// poller instead of waiting there, where nothing would wake the shard.
/// The shard runs on a thread of its own, so a drive that waited fails the
/// test rather than hanging it.
#[test]
fn a_timer_fired_on_the_way_to_park_runs_its_task_rather_than_parking() {
    let (done, finished) = std::sync::mpsc::channel();
    std::thread::spawn(move || {
        let shard = Shard::new(Limits {
            max_nesting: 8,
            buffer_cap: crate::buf::DEFAULT_CAP,
            backend: crate::rt::backend::chosen().unwrap(),
        })
        .unwrap();
        let _current = context::enter(&shard);
        let woke = std::rc::Rc::new(core::cell::Cell::new(false));
        let flag = std::rc::Rc::clone(&woke);
        crate::runtime::spawn_local(async move {
            crate::time::sleep(Duration::from_millis(1)).await;
            flag.set(true);
        })
        .detach();
        // The task arms its timer, which then comes due without firing.
        shard.run_ready();
        let armed_unfired = !woke.get();
        std::thread::sleep(Duration::from_millis(5));
        // Nothing but the timer has work for the shard: a drive that waited
        // on the poller after firing it would never return.
        shard.drive(None).unwrap();
        shard.run_ready();
        let _ = done.send((armed_unfired, woke.get()));
    });
    let (armed_unfired, woke) = finished
        .recv_timeout(Duration::from_secs(10))
        .expect("the drive waited on the poller with the task it woke still to run");
    assert!(armed_unfired, "the timer fired before it was due");
    assert!(woke, "the task the timer woke never ran");
}

/// A turn driven by hand runs a bounded number of tasks: a task that yields
/// again and again (waiting for a flag a timer's task sets) would otherwise
/// keep the turn from ever looking at its timers or returning.
#[test]
fn a_turn_runs_a_bounded_number_of_tasks() {
    crate::block_on(async {
        let runs = std::rc::Rc::new(core::cell::Cell::new(0u32));
        let counted = std::rc::Rc::clone(&runs);
        crate::runtime::spawn_local(async move {
            // A cap of its own, so a turn that never stops still ends the test.
            while counted.get() < 100_000 {
                counted.set(counted.get() + 1);
                crate::task::yield_now().await;
            }
        })
        .detach();
        crate::Proactor::get()
            .wait(64, Some(Duration::ZERO))
            .unwrap();
        assert!(
            runs.get() <= 2 * super::DRIVER_INTERVAL,
            "one turn ran the yielding task {} times",
            runs.get()
        );
    });
}

/// A notification of a sleeping shard that fails is tried again, so the
/// shard wakes for its work all the same. Only kqueue's `kevent` and the
/// port's `PostQueuedCompletionStatus` can fail; Linux's eventfd write
/// cannot. Each round's sender waits for the shard to say it parks, then
/// sends with its next notification failing; a round whose send found the
/// shard awake (its notification not made) is run again.
#[cfg(not(any(target_os = "linux", target_os = "android")))]
#[test]
fn a_failed_notification_is_tried_again() {
    let shared = Arc::clone(&context::own_shard().unwrap().shared);
    for _ in 0..50 {
        let (tx, rx) = crate::sync::oneshot::channel::<u32>();
        let sender = {
            let shared = Arc::clone(&shared);
            std::thread::spawn(move || {
                let deadline = Instant::now() + Duration::from_secs(10);
                let mut backoff = Duration::from_micros(50);
                while !shared.is_parked() {
                    assert!(Instant::now() < deadline, "the shard never parked");
                    std::thread::sleep(backoff);
                    backoff = (backoff * 2).min(Duration::from_millis(5));
                }
                // Into its wait on the kernel, past its last look.
                std::thread::sleep(Duration::from_millis(20));
                crate::rt::park::fail_notifies(1);
                tx.send(7).unwrap();
                crate::rt::park::fail_notifies(0) == 0
            })
        };
        let start = Instant::now();
        let got = crate::block_on(crate::time::timeout(Duration::from_secs(10), rx));
        let took = start.elapsed();
        let failed_one = sender.join().unwrap();
        assert_eq!(got.ok().and_then(Result::ok), Some(7));
        if failed_one {
            assert!(
                took < Duration::from_secs(5),
                "the failed notification was lost: the shard slept {took:?}"
            );
            return;
        }
    }
    panic!("no round's notification found the shard asleep");
}

/// The tasks bound to the thread that runs a shard are taken while other
/// threads keep its inbox full: every other interval's turn looks at the
/// home first, where before the home was looked at only once the inbox was
/// empty (proofs/tla/sched/Sched.tla `home_after_empty`).
#[test]
fn a_home_is_served_while_the_inbox_stays_full() {
    use std::sync::atomic::{AtomicBool, Ordering};

    use crate::rt::sched::Sched;

    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let _current = context::enter(&shard);
    let interval = shard.tuning.inbox_interval;
    // More runnables in the inbox than the runs the home may wait.
    let mut tasks = Vec::new();
    for _ in 0..4 * interval {
        // SAFETY: the schedule sends the task back to this shard, whose
        // thread runs it here; the future is 'static.
        let (runnable, task) = unsafe {
            crate::rt::task::spawn_unchecked(
                async {},
                Sched::<false>::counted(Arc::clone(&shard.shared), None, None),
            )
        };
        shard.shared.push_remote(runnable);
        tasks.push(task);
    }
    let ran = Arc::new(AtomicBool::new(false));
    let home = shard.own_home();
    let bound = {
        let ran = Arc::clone(&ran);
        // SAFETY: as above; the task is bound to the shard's own home.
        let (runnable, task) = unsafe {
            crate::rt::task::spawn_unchecked(
                async move { ran.store(true, Ordering::SeqCst) },
                Sched::<true>::counted(Arc::clone(&shard.shared), Some(Arc::clone(&home)), None),
            )
        };
        // SAFETY: this thread is the home's, and the new task is on no list.
        unsafe { home.push_task(runnable.links()) };
        // A wake of the bound task from another thread.
        home.push(runnable);
        task
    };
    let mut runs = 0;
    while !ran.load(Ordering::SeqCst) && runs < 2 * interval {
        let runnable = shard.next().expect("work is queued");
        shard.run(runnable);
        runs += 1;
    }
    assert!(
        ran.load(Ordering::SeqCst),
        "the bound task waited {runs} runs behind a full inbox"
    );
    drop((bound, tasks));
}

/// A waker that holds a shard's `Shared`, as the one of a worker's loop
/// does.
struct Holds(#[allow(dead_code)] Arc<super::Shared>);

#[expect(
    clippy::manual_noop_waker,
    reason = "the waker holds the Shared, which the no-op waker cannot"
)]
impl Wake for Holds {
    fn wake(self: Arc<Self>) {}
}

/// The waker a worker's loop leaves in its shard's `Shared` (when its look
/// after registering finds it may exit) goes with the shard, so the two do
/// not keep each other.
#[test]
fn a_shard_lets_go_of_the_waker_its_loop_left() {
    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let shared = Arc::downgrade(&shard.shared);
    let waker = Waker::from(Arc::new(Holds(Arc::clone(&shard.shared))));
    shard.shared.stop_waker.register(&waker);
    drop(waker);
    drop(shard);
    assert!(
        shared.upgrade().is_none(),
        "the shard's Shared was kept by its own loop's waker"
    );
}

/// A park publishes the shard's counts for `stats`, even one that does not
/// wait after all (a task queued already); a look between tasks publishes
/// every `PUBLISH_INTERVAL`th time.
#[test]
fn a_park_publishes_and_a_look_between_tasks_every_few() {
    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let _current = context::enter(&shard);
    let published = || shard.shared.stats().polls;
    shard.polls.set(1);
    for look in 1..super::PUBLISH_INTERVAL {
        shard.drive(Some(Duration::ZERO)).unwrap();
        assert_eq!(published(), 0, "look {look}");
    }
    shard.drive(Some(Duration::ZERO)).unwrap();
    assert_eq!(published(), 1);
    shard.polls.set(2);
    let queued = crate::runtime::spawn_local(async {});
    assert!(shard.has_local_work());
    shard.drive(None).unwrap();
    assert_eq!(published(), 2, "a park that found a task queued");
    drop(queued);
}

/// A `block_on` future that keeps waking itself, with no task to run, still
/// has the shard's IO driven: a read it waits for completes.
#[test]
fn a_future_that_keeps_waking_itself_still_has_its_io_driven() {
    use std::io::Write;

    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    let mut client = std::net::TcpStream::connect(listener.local_addr().unwrap()).unwrap();
    let server = crate::Handle::new(listener.accept().unwrap().0).unwrap();
    // The bytes come once the read waits: it needs the driver either way.
    let writer = std::thread::spawn(move || {
        std::thread::sleep(Duration::from_millis(50));
        client.write_all(b"x").unwrap();
        client
    });
    let got = crate::block_on(async {
        let mut buf = [0u8; 1];
        let mut read = core::pin::pin!(server.recv(&mut buf));
        let deadline = Instant::now() + Duration::from_secs(10);
        loop {
            let polled = core::future::poll_fn(|cx| {
                let polled = read.as_mut().poll(cx);
                // Woken again at once: the loop has no reason to wait.
                cx.waker().wake_by_ref();
                core::task::Poll::Ready(polled)
            })
            .await;
            if let core::task::Poll::Ready(n) = polled {
                break n;
            }
            assert!(Instant::now() < deadline, "the read was never driven");
            // Lets the loop come back to this future.
            let mut yielded = false;
            core::future::poll_fn(|cx| {
                if yielded {
                    core::task::Poll::Ready(())
                } else {
                    yielded = true;
                    cx.waker().wake_by_ref();
                    core::task::Poll::Pending
                }
            })
            .await;
        }
    });
    assert_eq!(got.unwrap(), 1);
    drop(writer.join().unwrap());
}

/// The leaf push from a task's code on a watched shard (`context::push_lifo_marked`): with no
/// claim up it fills the empty LIFO slot and puts the poll's word back, so the poll is outside
/// again; it backs out, the word put back, when the slot holds a task, for another shard, and with
/// a claim up, which it leaves pending for the full bracket to settle.
#[test]
fn a_leaf_push_fills_an_empty_slot_or_backs_out() {
    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let slot = crate::rt::handoff::slot::Slot::new(false);
    let _entered = context::enter_watched(&shard, Some(&slot));
    let addr = shard.shared_addr();
    let info = crate::rt::task::ScheduleInfo::new(false);
    /// A schedule the test never runs: it drops what it is handed.
    struct Ignore;
    // SAFETY: the runnable is not touched after the drop.
    unsafe impl crate::rt::task::Schedule for Ignore {
        fn schedule(&self, runnable: crate::rt::task::Runnable, _: crate::rt::task::ScheduleInfo) {
            drop(runnable);
        }
    }
    let task = || crate::rt::task::spawn(async {}, Ignore).0;
    let (_, polled) = context::poll(&shard.shared, || {
        let word = slot.word();
        assert!(
            context::push_lifo_marked(task(), info, |s| s.shared_addr() == addr).is_ok(),
            "the leaf did not fill the empty slot"
        );
        assert_eq!(slot.word(), word, "the leaf left its mark up");
        assert_eq!(shard.queued(), 1);
        let back = context::push_lifo_marked(task(), info, |s| s.shared_addr() == addr)
            .expect_err("the leaf pushed into a slot that holds a task");
        assert_eq!(slot.word(), word);
        let back = context::push_lifo_marked(back, info, |s| s.shared_addr() == addr ^ 1)
            .expect_err("the leaf pushed for another shard");
        assert_eq!(slot.word(), word);
        drop(shard.lifo.take());
        slot.claim();
        let back = context::push_lifo_marked(back, info, |s| s.shared_addr() == addr)
            .expect_err("the leaf pushed with a claim up");
        assert_eq!(slot.word(), word, "the leaf backed out with its mark up");
        assert_eq!(shard.queued(), 0);
        assert!(slot.take(), "the leaf settled the claim");
        drop(back);
    });
    assert_eq!(polled, context::Polled::Detached);
}

/// The leaf push from a task's code whose shard was taken under a call that has ended since: a
/// nested poll in the call (a nested `block_on`'s) found the take, and the call's end put the
/// slot back as the one the task's code runs outside of. The leaf finds the claim up and backs
/// out, and the full bracket finds no shard.
#[test]
fn a_leaf_push_after_a_take_under_a_call_backs_out_to_no_shard() {
    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let slot = crate::rt::handoff::slot::Slot::new(false);
    let _entered = context::enter_watched(&shard, Some(&slot));
    let addr = shard.shared_addr();
    let info = crate::rt::task::ScheduleInfo::new(false);
    /// A schedule the test never runs: it drops what it is handed.
    struct Ignore;
    // SAFETY: the runnable is not touched after the drop.
    unsafe impl crate::rt::task::Schedule for Ignore {
        fn schedule(&self, runnable: crate::rt::task::Runnable, _: crate::rt::task::ScheduleInfo) {
            drop(runnable);
        }
    }
    let (_, polled) = context::poll(&shard.shared, || {
        let word = slot.word();
        {
            let call = context::bracketed();
            assert!(call.shard().is_some(), "the call found no shard");
            let (_, nested) = context::poll(&shard.shared, || {
                slot.claim();
                assert!(slot.take(), "the nested poll was not taken");
            });
            assert_eq!(nested, context::Polled::Detached);
        }
        assert_eq!(slot.word(), word, "the call's end left another word");
        let task = crate::rt::task::spawn(async {}, Ignore).0;
        let back = context::push_lifo_marked(task, info, |s| s.shared_addr() == addr)
            .expect_err("the leaf pushed on a taken shard");
        assert_eq!(slot.word(), word, "the leaf left its mark up");
        assert_eq!(shard.queued(), 0);
        assert!(
            context::current(|current| current.is_none()),
            "the full bracket found the taken shard"
        );
        drop(back);
    });
    assert_eq!(polled, context::Polled::Detached);
}

/// The time to the precise deadline a shard's poller has its own timer armed
/// for: epoll's timerfd, or the `EVFILT_TIMER` the next wait submits.
#[cfg(any(target_os = "linux", target_os = "android", precise_timer))]
fn precise_armed_in(shard: &Shard) -> Option<Duration> {
    match &shard.driver {
        super::Driver::Poll(poller) => poller.borrow().precise_armed_in(),
        #[cfg(ring)]
        _ => None,
    }
}

// The pollers that arm a timer of their own for a precise deadline: epoll and
// kqueue's. A ring's timeout and a completion port's timer end the wait to
// the microsecond by themselves, so they keep no armed deadline to go stale.
#[cfg(any(target_os = "linux", target_os = "android", precise_timer))]
#[test]
fn an_advance_of_a_running_clock_moves_the_pollers_precise_timer() {
    let rt = crate::runtime::Builder::new()
        .worker_threads(0)
        .build()
        .unwrap();
    let inner = rt.inner();
    // Altered, and running.
    inner.clock.advance(Duration::from_millis(1));
    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::Backend::Poll,
    })
    .unwrap();
    let _runtime = context::set_runtime(inner);
    let _current = context::enter(&shard);
    // On a top-level slot boundary of the microsecond wheel, so its next
    // work is this deadline and not a move down on the way to it.
    let at = (inner.clock.now_nanos() / 1_000 + 6_000_000)
        .next_multiple_of(crate::time::wheel::TOP_SLOT)
        * 1_000;
    let lead = Duration::from_nanos(at - inner.clock.now_nanos());
    let key = shard
        .wheel_of(true)
        .borrow_mut()
        .insert(shard.time_of(true, at), Waker::noop().clone())
        .unwrap();
    assert!(shard.timer_wait(true).is_some());
    let first = precise_armed_in(&shard).expect("the poller armed its timer");
    assert!(first <= lead, "{first:?}");
    let by = lead - Duration::from_secs(5);
    inner.clock.advance(by);
    assert!(shard.timer_wait(true).is_some());
    let second = precise_armed_in(&shard).expect("the poller's timer is armed");
    assert!(second <= Duration::from_secs(5), "{second:?}");
    // Moved by the advance it is `by` nearer than it was, less the time
    // between the two looks, however long that took on a loaded host.
    assert!(
        second + by - Duration::from_secs(1) < first,
        "the poller's timer stayed at the deadline as the system saw it before the advance: {first:?} left before, {second:?} after"
    );
    drop(shard.wheel_of(true).borrow_mut().remove(key));
}

#[test]
fn a_park_that_backs_out_again_and_again_still_looks_at_the_driver() {
    use core::cell::RefCell;
    use core::task::{Poll, Waker};
    use std::rc::Rc;

    let shard = Shard::new(Limits {
        max_nesting: 8,
        buffer_cap: crate::buf::DEFAULT_CAP,
        backend: crate::rt::backend::chosen().unwrap(),
    })
    .unwrap();
    let _current = context::enter(&shard);
    shard.spin.spend();
    let before = shard.looks.get();
    let stored: Rc<RefCell<Option<Waker>>> = Rc::new(RefCell::new(None));
    let hook_stored = Rc::clone(&stored);
    super::run::BEFORE_PARK.with(|hook| {
        *hook.borrow_mut() = Some(Box::new(move || {
            if let Some(waker) = hook_stored.borrow().as_ref() {
                waker.wake_by_ref();
            }
        }));
    });
    let polls = 3 * super::DRIVER_INTERVAL;
    let mut seen = 0u32;
    shard.block_on(core::future::poll_fn(|cx| {
        *stored.borrow_mut() = Some(cx.waker().clone());
        seen += 1;
        if seen >= polls {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }));
    super::run::BEFORE_PARK.with(|hook| *hook.borrow_mut() = None);
    assert!(
        shard.looks.get() - before >= 2,
        "the driver was never looked at while the park kept backing out"
    );
}