use super::rng::random_below;
use core::cell::Cell;
use std::io;
use std::sync::{Arc, Mutex, OnceLock};
use std::thread::JoinHandle;
use std::time::Duration;
use super::context;
use super::handoff::home::Home;
use super::handoff::life::Life;
use super::handoff::slot::Slot;
use super::handoff::{Gift, Monitor, Owned, Runner, Spare};
use super::runtime::Inner;
use super::shard::{Limits, Shard, Shared, Worked};
use crate::sync::internal::{AtomicPtr, AtomicUsize, Ordering};
thread_local! {
static HOME: Cell<Option<(u64, usize)>> = const { Cell::new(None) };
}
pub(crate) struct Start {
pub(crate) limits: Limits,
pub(crate) tuning: super::tuning::Tuning,
pub(crate) seed: Option<super::rng::RngSeed>,
pub(crate) stack: usize,
pub(crate) name: Arc<dyn Fn() -> String + Send + Sync>,
pub(crate) cores: Vec<usize>,
pub(crate) handoff_after: Option<Duration>,
}
pub(crate) struct Workers {
slots: Box<[OnceLock<Arc<Shared>>]>,
claimed: AtomicUsize,
running: AtomicUsize,
min: usize,
start: Start,
threads: Mutex<Vec<(JoinHandle<()>, Arc<Life>)>>,
exits: kovan_channel::unbounded::Sender<()>,
monitor: Option<Arc<Monitor>>,
monitor_started: OnceLock<()>,
}
struct Run {
place: usize,
shard: Owned,
slot: Option<Arc<Slot>>,
}
impl From<Gift> for Run {
fn from(gift: Gift) -> Self {
Self {
place: gift.place,
shard: gift.shard,
slot: Some(gift.slot),
}
}
}
impl Workers {
pub(crate) fn new(
capacity: usize,
min: usize,
start: Start,
exits: kovan_channel::unbounded::Sender<()>,
) -> Self {
let monitor = start
.handoff_after
.filter(|_| capacity > 0 && super::handoff::SUPPORTED)
.map(|after| {
Arc::new(Monitor::new(
after,
capacity,
super::runtime::allowed_cores(),
))
});
Self {
slots: (0..capacity).map(|_| OnceLock::new()).collect(),
claimed: AtomicUsize::new(0),
running: AtomicUsize::new(0),
min,
start,
threads: Mutex::new(Vec::new()),
exits,
monitor,
monitor_started: OnceLock::new(),
}
}
pub(crate) fn running(&self) -> usize {
self.running.load(Ordering::Relaxed)
}
pub(crate) fn shards(&self) -> impl Iterator<Item = &Arc<Shared>> {
let claimed = self.claimed.load(Ordering::Acquire).min(self.slots.len());
self.slots[..claimed].iter().filter_map(OnceLock::get)
}
pub(crate) fn place(&self, runtime: u64, margin: usize) -> Option<&Arc<Shared>> {
let n = self.claimed.load(Ordering::Acquire).min(self.slots.len());
let take = |i: usize| self.slots[i].get().filter(|s| s.takes_tasks());
match n {
0 => return None,
1 => return take(0),
_ => {}
}
let load = |s: &Shared| s.load.load(Ordering::Relaxed);
let stay_or_move = |home: usize| {
let mine = take(home)?;
let other = (home + 1 + random_below(n - 1)) % n;
Some(match take(other) {
Some(o) if load(mine) > load(o).saturating_add(margin) => (other, o),
_ => (home, mine),
})
};
let local = context::current(|current| current.and_then(|s| s.worker_index(runtime)));
if let Some((_, shard)) = local.filter(|&i| i < n).and_then(stay_or_move) {
return Some(shard);
}
let home = HOME
.try_with(Cell::get)
.ok()
.flatten()
.filter(|&(owner, i)| owner == runtime && i < n)
.and_then(|(_, i)| stay_or_move(i));
let (index, shard) = match home {
Some(chosen) => chosen,
None => {
let a = random_below(n);
let b = (a + 1 + random_below(n - 1)) % n;
match (take(a), take(b)) {
(Some(x), Some(y)) if load(y) < load(x) => (b, y),
(Some(x), _) => (a, x),
(None, Some(y)) => (b, y),
(None, None) => (0..n).find_map(|i| take(i).map(|s| (i, s)))?,
}
}
};
let _ = HOME.try_with(|h| h.set(Some((runtime, index))));
Some(shard)
}
pub(crate) fn start(inner: &Arc<Inner>, count: usize) -> io::Result<usize> {
let mut started = 0;
while started < count && Self::start_one(inner)? {
started += 1;
}
Ok(started)
}
fn start_one(inner: &Arc<Inner>) -> io::Result<bool> {
let workers = &inner.workers;
let index = workers.claimed.fetch_add(1, Ordering::AcqRel);
if index >= workers.slots.len() {
workers.claimed.fetch_sub(1, Ordering::AcqRel);
return Ok(false);
}
let mut limits = workers.start.limits;
if workers.monitor.is_some() {
limits.backend = limits.backend.handed_over();
}
#[expect(
clippy::arc_with_non_send_sync,
reason = "the count is what is shared: the shard itself moves between threads only at a hand-over (`rt::handoff`)"
)]
let shard = Arc::new(Shard::worker(
limits,
workers.start.tuning.clone(),
inner.id(),
index,
workers.monitor.clone(),
)?);
let shared = Arc::clone(&shard.shared);
let _ = workers.slots[index].set(Arc::clone(&shared));
workers.running.fetch_add(1, Ordering::AcqRel);
let life = Arc::new(Life::new());
let slot = workers.monitor.as_ref().map(|monitor| {
let slot = monitor.slot();
monitor.started(
index,
&shard,
Runner {
slot: Arc::clone(&slot),
life: Arc::clone(&life),
},
);
slot
});
let run = Run {
place: index,
shard: Owned(shard),
slot,
};
match Self::spawn_thread(inner, run, life) {
Ok(()) => {
Self::start_monitor(inner)?;
Ok(true)
}
Err((err, _)) => {
shared.stop();
workers.running.fetch_sub(1, Ordering::AcqRel);
Err(err)
}
}
}
fn start_monitor(inner: &Arc<Inner>) -> io::Result<()> {
let workers = &inner.workers;
let Some(monitor) = &workers.monitor else {
return Ok(());
};
if workers.monitor_started.set(()).is_err() {
return Ok(());
}
let (monitor, runtime) = (Arc::clone(monitor), Arc::downgrade(inner));
let exit = Exit(workers.exits.clone());
let thread = std::thread::Builder::new()
.name("nuclei-monitor".to_string())
.spawn(move || {
let _exit = exit;
super::handoff::run(&monitor, &runtime);
})?;
lock(&workers.threads).push((thread, Arc::new(Life::new())));
Ok(())
}
pub(crate) fn spawn_runner(
inner: &Arc<Inner>,
gift: Gift,
life: Arc<Life>,
) -> Result<(), Gift> {
Self::spawn_thread(inner, Run::from(gift), life).map_err(|(_, run)| Gift {
place: run.place,
shard: run.shard,
slot: run.slot.unwrap_or_else(|| inner.workers.monitor_slot()),
})
}
fn monitor_slot(&self) -> Arc<Slot> {
self.monitor
.as_ref()
.map_or_else(|| Arc::new(Slot::new(true)), |monitor| monitor.slot())
}
fn spawn_thread(inner: &Arc<Inner>, run: Run, life: Arc<Life>) -> Result<(), (io::Error, Run)> {
let workers = &inner.workers;
let handed = Arc::new(Handed(AtomicPtr::new(Box::into_raw(Box::new(run)))));
let (taken, thread_inner, thread_life) =
(Arc::clone(&handed), Arc::clone(inner), Arc::clone(&life));
let spawned = std::thread::Builder::new()
.name((workers.start.name)())
.stack_size(workers.start.stack)
.spawn(move || {
if let Some(run) = taken.take() {
runner(&thread_inner, run, &thread_life);
}
});
match spawned {
Ok(thread) => {
lock(&workers.threads).push((thread, life));
Ok(())
}
Err(err) => match handed.take() {
Some(run) => Err((err, run)),
None => Ok(()),
},
}
}
fn claim_stop(&self) -> bool {
let mut running = self.running.load(Ordering::Acquire);
while running > self.min {
match self.running.compare_exchange_weak(
running,
running - 1,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return true,
Err(now) => running = now,
}
}
false
}
pub(crate) fn stop_one(&self) -> bool {
if !self.claim_stop() {
return false;
}
match self.shards().filter(|s| s.takes_tasks()).last() {
Some(shard) => {
shard.drain();
true
}
None => {
self.running.fetch_add(1, Ordering::AcqRel);
false
}
}
}
pub(crate) fn stop(&self, shard: &Shared) -> bool {
if !shard.takes_tasks() || !self.claim_stop() {
return false;
}
shard.drain();
true
}
pub(crate) fn stop_all(&self) {
if let Some(monitor) = &self.monitor {
monitor.close();
}
self.shards().for_each(|s| s.stop());
}
pub(crate) fn take_threads(&self) -> Vec<JoinHandle<()>> {
core::mem::take(&mut *lock(&self.threads))
.into_iter()
.filter(|(_, life)| !life.is_detached())
.map(|(thread, _)| thread)
.collect()
}
}
struct Exit(kovan_channel::unbounded::Sender<()>);
impl Drop for Exit {
fn drop(&mut self) {
self.0.send(());
}
}
struct Handed(AtomicPtr<Run>);
impl Handed {
fn take(&self) -> Option<Run> {
let run = self.0.swap(core::ptr::null_mut(), Ordering::AcqRel);
(!run.is_null()).then(|| *unsafe { Box::from_raw(run) })
}
}
impl Drop for Handed {
fn drop(&mut self) {
drop(self.take());
}
}
fn runner(inner: &Arc<Inner>, first: Run, life: &Arc<Life>) {
let workers = &inner.workers;
let _exit = Exit(workers.exits.clone());
let _stopping = workers.start.tuning.hooks.started();
life.here();
let _runtime = context::set_runtime(inner);
let spare = workers
.monitor
.as_ref()
.map(|_| Spare::here(Arc::clone(life)));
let mut next = Some(first);
while let Some(run) = next.take() {
if run_shard(inner, run) == Worked::Detached {
life.back();
}
if let (Some(monitor), Some(spare)) = (&workers.monitor, &spare) {
next = monitor.wait_for_shard(spare).map(Run::from);
}
}
}
fn run_shard(inner: &Arc<Inner>, run: Run) -> Worked {
let workers = &inner.workers;
let Run {
place,
shard: Owned(shard),
slot,
} = run;
let cores = &workers.start.cores;
if !cores.is_empty() {
super::runtime::pin_to(cores[place % cores.len()]);
}
let _seeded = workers
.start
.seed
.map(|seed| super::rng::Reseeded::new(seed, place as u64 + 1));
let home = Arc::new(Home::new(Arc::clone(&shard.shared)));
context::set_home(Arc::clone(&home));
let worked = {
let _current = context::enter_watched(&shard, slot.as_deref());
shard.run_worker()
};
if worked == Worked::Detached
&& let Some(monitor) = &workers.monitor
{
finish_home(monitor, &home);
}
drop(shard);
worked
}
#[expect(
clippy::panic,
reason = "a thread whose shard cannot start cannot run its bound tasks, which would wait forever"
)]
fn finish_home(monitor: &Monitor, home: &Home) {
if home.is_idle() {
return;
}
let own = match context::own_shard() {
Ok(own) => own,
Err(err) => panic!("this thread's shard could not start: {err}"),
};
home.target().move_to(&own.shared);
let _current = context::enter(&own);
#[cfg(port)]
own.take_port_back();
own.block_on_then_leave(monitor.home_done(home));
}
fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[cfg(all(test, not(loom)))]
mod tests {
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use super::Workers;
use crate::runtime::{Builder, Runtime};
fn runtime(workers: usize, max: usize) -> Runtime {
Builder::new()
.worker_threads(workers)
.max_worker_threads(max)
.pin_workers(false)
.build()
.unwrap()
}
fn eventually(what: &str, mut done: impl FnMut() -> bool) {
let deadline = Instant::now() + Duration::from_secs(10);
while !done() {
assert!(Instant::now() < deadline, "{what} never happened");
std::thread::yield_now();
}
}
fn threads_used(rt: &Runtime, tasks: usize) -> HashSet<std::thread::ThreadId> {
let ids: Vec<_> = (0..tasks)
.map(|_| {
rt.spawn(async {
crate::time::sleep(Duration::from_millis(1)).await;
std::thread::current().id()
})
})
.collect();
rt.block_on(async {
let mut set = HashSet::new();
for id in ids {
set.insert(id.await);
}
set
})
}
#[test]
fn a_runtime_watches_its_workers_only_with_the_handoff_on() {
let off = Builder::new()
.worker_threads(1)
.pin_workers(false)
.handoff_after(None)
.build()
.unwrap();
assert!(off.inner().workers.monitor.is_none());
let on = runtime(1, 1);
assert!(on.inner().workers.monitor.is_some(), "on by default");
assert!(on.inner().workers.monitor_started.get().is_some());
}
#[test]
fn more_workers_start_up_to_the_maximum_and_take_tasks() {
let rt = runtime(1, 3);
assert_eq!(threads_used(&rt, 32).len(), 1);
assert_eq!(
Workers::start(rt.inner(), 5).unwrap(),
2,
"two fit under the maximum"
);
assert_eq!(rt.inner().workers.running(), 3);
assert_eq!(Workers::start(rt.inner(), 1).unwrap(), 0, "none is left");
assert!(threads_used(&rt, 300).len() > 1, "new workers take tasks");
}
#[test]
fn a_stop_leaves_the_minimum_running() {
let rt = runtime(1, 3);
Workers::start(rt.inner(), 2).unwrap();
assert!(rt.inner().workers.stop_one());
assert!(rt.inner().workers.stop_one());
assert!(!rt.inner().workers.stop_one(), "the minimum stays");
assert_eq!(rt.inner().workers.running(), 1);
assert_eq!(threads_used(&rt, 64).len(), 1);
}
#[test]
fn a_stopped_worker_finishes_its_task_then_exits() {
let rt = runtime(1, 2);
Workers::start(rt.inner(), 1).unwrap();
let workers = &rt.inner().workers;
let second = workers.shards().nth(1).map(Arc::clone).unwrap();
let release = Arc::new(AtomicBool::new(false));
let held = (0..1000)
.find_map(|_| {
let (r, s) = (Arc::clone(&release), Arc::clone(&second));
let (tell, landed) = std::sync::mpsc::channel();
let probe = async move {
let here = crate::rt::context::current(|c| c.is_some_and(|shard| shard.is(&s)));
let _ = tell.send(here);
while here && !r.load(Ordering::SeqCst) {
crate::time::sleep(Duration::from_millis(1)).await;
}
here
};
let task =
std::thread::scope(|scope| scope.spawn(|| rt.spawn(probe)).join().unwrap());
let here = landed.recv_timeout(Duration::from_secs(10)).unwrap();
here.then_some(task)
})
.expect("no task landed on the second worker");
assert!(workers.stop(&second));
assert!(!second.takes_tasks());
assert!(!second.is_stopped());
release.store(true, Ordering::SeqCst);
assert!(rt.block_on(held));
eventually("the stopped worker's exit", || second.is_stopped());
}
#[test]
fn a_source_on_a_worker_that_exits_registers_again_where_it_is_polled() {
use std::io::Write;
let rt = runtime(2, 2);
#[cfg(unix)]
let (mut peer, mine) = std::os::unix::net::UnixStream::pair().unwrap();
#[cfg(windows)]
let (mut peer, mine) = {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let peer = std::net::TcpStream::connect(listener.local_addr().unwrap()).unwrap();
(peer, listener.accept().unwrap().0)
};
let handle = Arc::new(crate::io::Handle::new(mine).unwrap());
let h = Arc::clone(&handle);
let first = rt.spawn(async move {
let mut buf = [0u8; 4];
let n = h.recv(&mut buf).await.unwrap();
(buf, n, std::thread::current().id())
});
peer.write_all(b"one!").unwrap();
let (_, n, _) = rt.block_on(first);
assert_eq!(n, 4);
for shard in rt.inner().workers.shards() {
shard.stop();
}
eventually("the workers' exit", || {
rt.inner().workers.shards().all(|s| s.is_stopped())
});
let h = Arc::clone(&handle);
let writer = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(10));
peer.write_all(b"two!").unwrap();
peer
});
let mut buf = [0u8; 4];
let n = crate::runtime::block_on(async {
crate::time::timeout(Duration::from_secs(10), h.recv(&mut buf)).await
})
.unwrap_or_else(|_| {
let source = handle.source();
panic!(
"the receive never completed: read side {}; write side {}",
source.read_report(),
source.write_report()
)
})
.unwrap();
assert_eq!(&buf[..n], b"two!");
drop(writer.join().unwrap());
}
}
#[cfg(all(test, not(loom)))]
mod placement_tests {
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use crate::runtime::{Builder, Runtime};
fn runtime(workers: usize) -> Runtime {
Builder::new()
.worker_threads(workers)
.pin_workers(false)
.build()
.unwrap()
}
fn spread(rt: &Runtime, count: usize, release: &Arc<AtomicBool>) -> Vec<usize> {
let (tell, landed) = std::sync::mpsc::channel();
let tasks: Vec<_> = (0..count)
.map(|_| {
let (tell, release) = (tell.clone(), Arc::clone(release));
rt.spawn(async move {
let _ = tell.send(std::thread::current().id());
while !release.load(Ordering::SeqCst) {
crate::time::sleep(Duration::from_millis(1)).await;
}
})
})
.collect();
let mut per_thread: HashMap<_, usize> = HashMap::new();
for _ in 0..count {
let id = landed.recv_timeout(Duration::from_secs(10)).unwrap();
*per_thread.entry(id).or_default() += 1;
}
tasks.into_iter().for_each(|t| t.detach());
let mut counts: Vec<_> = per_thread.into_values().collect();
counts.sort_unstable();
counts
}
#[test]
fn spawns_from_outside_the_runtime_stay_together() {
let rt = runtime(4);
let release = Arc::new(AtomicBool::new(true));
assert_eq!(spread(&rt, 20, &release), [20]);
}
#[test]
fn spawns_from_outside_spread_once_a_shard_is_a_margin_ahead() {
let rt = runtime(4);
let release = Arc::new(AtomicBool::new(false));
let counts = spread(&rt, 400, &release);
release.store(true, Ordering::SeqCst);
assert_eq!(counts.len(), 4, "every worker took tasks: {counts:?}");
let (min, max) = (counts[0], counts[counts.len() - 1]);
assert!(
max - min <= 3 * 32,
"the load stays within the margin: {counts:?}"
);
}
}