use core::cell::Cell;
use core::ptr::{self, NonNull};
use core::sync::atomic::{AtomicPtr, AtomicU32, AtomicU64, AtomicUsize, Ordering};
use crate::{Futex, WaitGroup};
use bun_core::Output;
#[inline]
fn stats_enabled() -> bool {
static CELL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
*CELL.get_or_init(|| bun_core::getenv_z(bun_core::zstr!("BUN_THREADPOOL_STATS")).is_some())
}
#[derive(Default)]
struct PoolStats {
idle_ns: AtomicU64,
busy_ns: AtomicU64,
tasks: AtomicU64,
sleeps: AtomicU64,
last_dump_ns: AtomicU64,
}
#[repr(transparent)]
#[derive(Copy, Clone, PartialEq, Eq)]
struct Sync(u32);
#[repr(u8)]
#[derive(Copy, Clone, PartialEq, Eq)]
enum SyncState {
Pending = 0,
Signaled = 1,
Waking = 2,
Shutdown = 3,
}
impl Sync {
const IDLE_MASK: u32 = 0x3FFF;
const SPAWNED_SHIFT: u32 = 14;
const SPAWNED_MASK: u32 = 0x3FFF << Self::SPAWNED_SHIFT;
const NOTIFIED_BIT: u32 = 1 << 29;
const STATE_SHIFT: u32 = 30;
const STATE_MASK: u32 = 0b11 << Self::STATE_SHIFT;
const fn zero() -> Self {
Sync(0)
}
#[inline]
fn idle(self) -> u16 {
(self.0 & Self::IDLE_MASK) as u16
}
#[inline]
fn set_idle(&mut self, v: u16) {
self.0 = (self.0 & !Self::IDLE_MASK) | (v as u32 & Self::IDLE_MASK);
}
#[inline]
fn spawned(self) -> u16 {
((self.0 & Self::SPAWNED_MASK) >> Self::SPAWNED_SHIFT) as u16
}
#[inline]
fn set_spawned(&mut self, v: u16) {
self.0 = (self.0 & !Self::SPAWNED_MASK) | ((v as u32 & 0x3FFF) << Self::SPAWNED_SHIFT);
}
#[inline]
fn notified(self) -> bool {
self.0 & Self::NOTIFIED_BIT != 0
}
#[inline]
fn set_notified(&mut self, v: bool) {
if v {
self.0 |= Self::NOTIFIED_BIT;
} else {
self.0 &= !Self::NOTIFIED_BIT;
}
}
#[inline]
fn state(self) -> SyncState {
match (self.0 >> Self::STATE_SHIFT) & 0b11 {
0 => SyncState::Pending,
1 => SyncState::Signaled,
2 => SyncState::Waking,
_ => SyncState::Shutdown,
}
}
#[inline]
fn set_state(&mut self, s: SyncState) {
self.0 = (self.0 & !Self::STATE_MASK) | ((s as u32) << Self::STATE_SHIFT);
}
}
#[repr(transparent)]
struct AtomicSync(AtomicU32);
impl AtomicSync {
const fn new(v: Sync) -> Self {
AtomicSync(AtomicU32::new(v.0))
}
#[inline]
fn load(&self, order: Ordering) -> Sync {
Sync(self.0.load(order))
}
#[inline]
fn cmpxchg_weak(
&self,
old: Sync,
new: Sync,
success: Ordering,
failure: Ordering,
) -> Option<Sync> {
match self.0.compare_exchange_weak(old.0, new.0, success, failure) {
Ok(_) => None,
Err(cur) => Some(Sync(cur)),
}
}
#[inline]
fn fetch_or(&self, val: Sync, order: Ordering) -> Sync {
Sync(self.0.fetch_or(val.0, order))
}
#[inline]
fn fetch_sub(&self, val: Sync, order: Ordering) -> Sync {
Sync(self.0.fetch_sub(val.0, order))
}
}
pub struct ThreadPool {
pub sleep_on_idle_network_thread: bool,
pub needs_stack_bounds: bool,
pub stack_size: u32,
pub max_threads: u32,
sync: AtomicSync,
idle_event: Event,
join_event: Event,
run_queue: node::Queue,
threads: AtomicPtr<Thread>,
pub name: &'static [u8],
pub spawned_thread_count: AtomicU32,
stats: PoolStats,
}
#[derive(Clone, Copy)]
pub struct Config {
pub stack_size: u32,
pub max_threads: u32,
}
impl Default for Config {
fn default() -> Self {
Config {
stack_size: DEFAULT_THREAD_STACK_SIZE,
max_threads: 0,
}
}
}
impl ThreadPool {
pub fn init(config: Config) -> ThreadPool {
ThreadPool {
sleep_on_idle_network_thread: true,
needs_stack_bounds: true,
stack_size: 1.max(config.stack_size),
max_threads: 1.max(config.max_threads),
sync: AtomicSync::new(Sync::zero()),
idle_event: Event::default(),
join_event: Event::default(),
run_queue: node::Queue::default(),
threads: AtomicPtr::new(ptr::null_mut()),
name: b"",
spawned_thread_count: AtomicU32::new(0),
stats: PoolStats {
last_dump_ns: AtomicU64::new(if stats_enabled() { now_ns() } else { 0 }),
..PoolStats::default()
},
}
}
pub fn dump_stats(&self, label: &str) {
if !stats_enabled() {
return;
}
let now = now_ns();
let idle = self.stats.idle_ns.swap(0, Ordering::Relaxed);
let busy = self.stats.busy_ns.swap(0, Ordering::Relaxed);
let tasks = self.stats.tasks.swap(0, Ordering::Relaxed);
let sleeps = self.stats.sleeps.swap(0, Ordering::Relaxed);
let last = self.stats.last_dump_ns.swap(now, Ordering::Relaxed);
let spawned = self.sync.load(Ordering::Relaxed).spawned();
let total = idle + busy;
let util = if total > 0 {
(busy as f64 / total as f64) * 100.0
} else {
0.0
};
let wall = if last == 0 { 0 } else { now.wrapping_sub(last) };
let eff = if wall > 0 {
busy as f64 / wall as f64
} else {
0.0
};
Output::print_errorln(format_args!(
"[threadpool {}] workers={} tasks={} wall={:.3}s busy={:.3}s idle={:.3}s util={:.1}% eff_cpus={:.2} sleeps={}",
label,
spawned,
tasks,
wall as f64 / 1e9,
busy as f64 / 1e9,
idle as f64 / 1e9,
util,
eff,
sleeps,
));
}
pub fn wake_for_idle_events(&self) {
self.idle_event.wake(Event::NOTIFIED, u32::MAX);
}
}
impl Default for ThreadPool {
fn default() -> Self {
Self::init(Config::default())
}
}
impl Drop for ThreadPool {
fn drop(&mut self) {
self.shutdown();
self.join();
}
}
#[repr(C)]
pub struct Task {
pub node: Node,
pub callback: unsafe fn(*mut Task),
}
#[repr(C)]
pub struct CountedTask {
pub task: Task,
run: unsafe fn(*mut Task),
group: *const WaitGroup,
}
unsafe impl Send for CountedTask {}
impl CountedTask {
pub fn new(run: unsafe fn(*mut Task), group: &WaitGroup) -> Self {
Self {
task: Task {
node: Node::default(),
callback: Self::run_and_finish,
},
run,
group: core::ptr::from_ref(group),
}
}
unsafe fn run_and_finish(task: *mut Task) {
const _: () = assert!(core::mem::offset_of!(CountedTask, task) == 0);
let this = task.cast::<CountedTask>();
let (run, group) = unsafe { ((*this).run, (*this).group) };
unsafe { run(task) };
unsafe { WaitGroup::finish_raw(group) };
}
}
unsafe impl Send for Task {}
impl Default for Task {
#[inline]
fn default() -> Self {
fn unreachable_cb(_: *mut Task) {
unreachable!("ThreadPool.Task scheduled with default() callback");
}
Task {
node: Node::default(),
callback: unreachable_cb,
}
}
}
impl Task {
#[inline]
unsafe fn from_node(node: *mut Node) -> *mut Task {
unsafe { bun_core::from_field_ptr!(Task, node, node) }
}
#[inline]
fn node_of(task: NonNull<Task>) -> NonNull<Node> {
const _: () = assert!(core::mem::offset_of!(Task, node) == 0);
task.cast::<Node>()
}
}
#[derive(Default, Clone, Copy)]
pub struct Batch {
pub len: usize,
pub head: Option<NonNull<Task>>,
pub tail: Option<NonNull<Task>>,
}
impl Batch {
pub fn pop(&mut self) -> Option<NonNull<Task>> {
let len = unsafe { (*(&raw const self.len).cast::<AtomicUsize>()).load(Ordering::Relaxed) };
if len == 0 {
return None;
}
let task = self.head.unwrap();
let next = unsafe { (*Task::node_of(task).as_ptr()).next };
if !next.is_null() {
self.head = NonNull::new(unsafe { Task::from_node(next) });
} else {
if task != self.tail.unwrap() {
unreachable!();
}
self.tail = None;
self.head = None;
}
self.len -= 1;
if len == 0 {
self.tail = None;
}
Some(task)
}
pub fn from(task: *mut Task) -> Batch {
let task = NonNull::new(task);
Batch {
len: 1,
head: task,
tail: task,
}
}
pub fn push(&mut self, batch: Batch) {
if batch.len == 0 {
return;
}
if self.len == 0 {
*self = batch;
} else {
let tail_node = Task::node_of(self.tail.unwrap());
let new_next = batch
.head
.map_or(ptr::null_mut(), |h| Task::node_of(h).as_ptr());
unsafe { (*tail_node.as_ptr()).next = new_next };
self.tail = batch.tail;
self.len += batch.len;
}
}
}
trait EachCall<Ctx, V>: core::marker::Sync {
unsafe fn call(&self, ctx: &Ctx, value: *mut V, i: usize);
}
struct ByValue<F>(F);
impl<Ctx, V: Copy, F> EachCall<Ctx, V> for ByValue<F>
where
F: Fn(&Ctx, V, usize) + core::marker::Sync,
{
#[inline]
unsafe fn call(&self, ctx: &Ctx, value: *mut V, i: usize) {
(self.0)(ctx, unsafe { *value }, i);
}
}
struct ByPtr<F>(F);
impl<Ctx, V, F> EachCall<Ctx, V> for ByPtr<F>
where
F: Fn(&Ctx, *mut V, usize) + core::marker::Sync,
{
#[inline]
unsafe fn call(&self, ctx: &Ctx, value: *mut V, i: usize) {
(self.0)(ctx, value, i);
}
}
impl ThreadPool {
pub fn each<Ctx, V, F>(&self, ctx: Ctx, run_fn: F, values: &mut [V])
where
F: Fn(&Ctx, V, usize) + core::marker::Sync,
Ctx: core::marker::Sync,
V: Copy + core::marker::Sync + core::marker::Send,
{
self.each_impl(ctx, ByValue(run_fn), values);
}
pub fn each_ptr<Ctx, V, F>(&self, ctx: Ctx, run_fn: F, values: &mut [V])
where
F: Fn(&Ctx, *mut V, usize) + core::marker::Sync,
Ctx: core::marker::Sync,
V: core::marker::Sync + core::marker::Send,
{
self.each_impl(ctx, ByPtr(run_fn), values);
}
fn each_impl<Ctx, V, F>(&self, ctx: Ctx, run_fn: F, values: &mut [V])
where
F: EachCall<Ctx, V>,
Ctx: core::marker::Sync,
V: core::marker::Sync + core::marker::Send,
{
if values.is_empty() {
return;
}
struct WaitContext<Ctx, V, F> {
ctx: Ctx,
values: *mut [V],
run_fn: F,
}
#[repr(C)]
struct RunnerTask<Ctx, V, F> {
task: CountedTask,
ctx: bun_ptr::BackRef<WaitContext<Ctx, V, F>>,
i: usize,
}
unsafe fn call<Ctx, V, F: EachCall<Ctx, V>>(task: *mut Task) {
let runner_task =
unsafe { &mut *bun_core::from_field_ptr!(RunnerTask<Ctx, V, F>, task, task) };
let i = runner_task.i;
let wctx = runner_task.ctx.get();
let value: *mut V = unsafe { &raw mut (*wctx.values)[i] };
unsafe { wctx.run_fn.call(&wctx.ctx, value, i) };
}
let wait_context = WaitContext {
ctx,
values: std::ptr::from_mut::<[V]>(values),
run_fn,
};
let group = WaitGroup::init_with_count(values.len());
let mut tasks: Vec<RunnerTask<Ctx, V, F>> = Vec::with_capacity(values.len());
let mut batch = Batch::default();
let mut offset = values.len();
for _ in 0..values.len() {
offset -= 1;
tasks.push(RunnerTask {
i: offset,
task: CountedTask::new(call::<Ctx, V, F>, &group),
ctx: bun_ptr::BackRef::new(&wait_context),
});
}
for runner_task in tasks.iter_mut() {
batch.push(Batch::from(&raw mut runner_task.task.task));
}
self.schedule(batch);
group.wait();
}
fn schedule_impl(&self, batch: &Batch, try_current: bool) {
let Batch { len, head, tail } = *batch;
if len == 0 {
return;
}
let mut list = node::List {
head: Task::node_of(head.unwrap()),
tail: Task::node_of(tail.unwrap()),
};
let current: *mut Thread = 'blk: {
if !try_current {
break 'blk ptr::null_mut();
}
let Some(current) = NonNull::new(Thread::current()) else {
break 'blk ptr::null_mut();
};
if bun_ptr::BackRef::from(current)
.thread_pool
.as_ptr()
.cast_const()
== std::ptr::from_ref::<ThreadPool>(self)
{
current.as_ptr()
} else {
ptr::null_mut()
}
};
if !current.is_null() {
unsafe {
if (*current).run_buffer.push(&mut list).is_err() {
(*current).run_queue.push(&list);
}
}
} else {
self.run_queue.push(&list);
}
self.force_spawn();
}
pub fn schedule(&self, batch: Batch) {
self.schedule_impl(&batch, false);
}
pub fn schedule_inside_thread_pool(&self, batch: Batch) {
self.schedule_impl(&batch, true);
}
fn force_spawn(&self) {
let is_waking = false;
self.notify(is_waking);
}
#[inline(always)]
fn notify(&self, is_waking: bool) {
if !is_waking {
let sync = self.sync.fetch_or(Sync::zero(), Ordering::Release);
if sync.notified() {
return;
}
}
self.notify_slow(is_waking);
}
}
pub const DEFAULT_THREAD_STACK_SIZE: u32 = {
const DEFAULT: u32 = 4 * 1024 * 1024;
#[cfg(windows)]
{
let _ = DEFAULT;
0x1200000
}
#[cfg(all(not(target_os = "macos"), not(windows)))]
{
DEFAULT
}
#[cfg(target_os = "macos")]
{
const PAGE_SIZE_MAX: u32 = 16384;
let size = DEFAULT - (DEFAULT % PAGE_SIZE_MAX);
assert!(
size.is_multiple_of(PAGE_SIZE_MAX),
"Thread stack size is not a multiple of page size"
);
size
}
};
impl ThreadPool {
pub fn warm(&self, count: u16) {
let target = count.min((self.max_threads & Sync::IDLE_MASK) as u16);
let mut sync = self.sync.load(Ordering::Relaxed);
while sync.spawned() < target {
let mut new_sync = sync;
new_sync.set_spawned(new_sync.spawned() + 1);
if let Some(current) =
self.sync
.cmpxchg_weak(sync, new_sync, Ordering::Release, Ordering::Relaxed)
{
sync = current;
continue;
}
let stack_size = self.stack_size as usize;
let pool = bun_ptr::BackRef::new(self);
match std::thread::Builder::new()
.stack_size(stack_size)
.spawn(move || Thread::run(pool))
{
Ok(_handle) => {
}
Err(_) => {
return unsafe { Self::unregister(self, ptr::null_mut()) };
}
}
sync = new_sync;
}
}
#[inline(never)]
fn notify_slow(&self, is_waking: bool) {
let mut sync = self.sync.load(Ordering::Relaxed);
while sync.state() != SyncState::Shutdown {
let can_wake = is_waking || (sync.state() == SyncState::Pending);
if is_waking {
debug_assert!(sync.state() == SyncState::Waking);
}
let mut new_sync = sync;
new_sync.set_notified(true);
if can_wake && sync.idle() > 0 {
new_sync.set_state(SyncState::Signaled);
} else if can_wake && (sync.spawned() as u32) < self.max_threads {
new_sync.set_state(SyncState::Signaled);
new_sync.set_spawned(new_sync.spawned() + 1);
} else if is_waking {
new_sync.set_state(SyncState::Pending);
} else if sync.notified() {
return;
}
sync =
match self
.sync
.cmpxchg_weak(sync, new_sync, Ordering::Release, Ordering::Relaxed)
{
Some(cur) => cur,
None => {
if can_wake && sync.idle() > 0 {
return self.idle_event.notify();
}
if can_wake && (sync.spawned() as u32) < self.max_threads {
let stack_size = self.stack_size as usize;
let pool = bun_ptr::BackRef::new(self);
match std::thread::Builder::new()
.stack_size(stack_size)
.spawn(move || Thread::run(pool))
{
Ok(_handle) => {
}
Err(_) => {
return unsafe { Self::unregister(self, ptr::null_mut()) };
}
}
return;
}
return;
}
};
}
}
#[inline(never)]
fn wait(&self, _is_waking: bool) -> Result<bool, WaitError> {
let mut is_idle = false;
let mut is_waking = _is_waking;
let mut sync = self.sync.load(Ordering::Relaxed);
loop {
if sync.state() == SyncState::Shutdown {
return Err(WaitError::Shutdown);
}
if is_waking {
debug_assert!(sync.state() == SyncState::Waking);
}
if sync.notified() {
let mut new_sync = sync;
new_sync.set_notified(false);
if is_idle {
new_sync.set_idle(new_sync.idle() - 1);
}
if sync.state() == SyncState::Signaled {
new_sync.set_state(SyncState::Waking);
}
sync = match self.sync.cmpxchg_weak(
sync,
new_sync,
Ordering::Acquire,
Ordering::Relaxed,
) {
Some(cur) => cur,
None => {
return Ok(is_waking || (sync.state() == SyncState::Signaled));
}
};
} else if !is_idle {
let mut new_sync = sync;
new_sync.set_idle(new_sync.idle() + 1);
if is_waking {
new_sync.set_state(SyncState::Pending);
}
sync = match self.sync.cmpxchg_weak(
sync,
new_sync,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Some(cur) => cur,
None => {
is_waking = false;
is_idle = true;
continue;
}
};
} else {
if let Some(current) = NonNull::new(Thread::current()) {
bun_ptr::BackRef::from(current).drain_idle_events();
}
if stats_enabled() {
self.stats.sleeps.fetch_add(1, Ordering::Relaxed);
}
self.idle_event.wait();
sync = self.sync.load(Ordering::Relaxed);
}
}
}
#[inline(never)]
pub fn shutdown(&self) {
let mut sync = self.sync.load(Ordering::Relaxed);
while sync.state() != SyncState::Shutdown {
let mut new_sync = sync;
new_sync.set_notified(true);
new_sync.set_state(SyncState::Shutdown);
new_sync.set_idle(0);
sync = match self
.sync
.cmpxchg_weak(sync, new_sync, Ordering::AcqRel, Ordering::Relaxed)
{
Some(cur) => cur,
None => {
if sync.idle() > 0 {
self.idle_event.shutdown();
}
return;
}
};
}
}
fn register(&self, thread: *mut Thread) {
let mut threads = self.threads.load(Ordering::Relaxed);
loop {
unsafe { (*thread).next = threads };
match self.threads.compare_exchange_weak(
threads,
thread,
Ordering::Release,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(cur) => threads = cur,
}
}
}
unsafe fn unregister(pool: *const Self, maybe_thread: *mut Thread) {
let one_spawned = {
let mut s = Sync::zero();
s.set_spawned(1);
s
};
let sync = unsafe { (*pool).sync.fetch_sub(one_spawned, Ordering::Release) };
debug_assert!(sync.spawned() > 0);
if sync.state() == SyncState::Shutdown && sync.spawned() == 1 {
unsafe { (*pool).join_event.notify() };
}
let Some(thread) = NonNull::new(maybe_thread) else {
return;
};
let thread = bun_ptr::BackRef::from(thread);
thread.join_event.wait();
let Some(next_thread) = NonNull::new(thread.next) else {
return;
};
bun_ptr::BackRef::from(next_thread).join_event.notify();
}
fn join(&self) {
let mut sync = self.sync.load(Ordering::Relaxed);
if !(sync.state() == SyncState::Shutdown && sync.spawned() == 0) {
self.join_event.wait();
sync = self.sync.load(Ordering::Relaxed);
}
debug_assert!(sync.state() == SyncState::Shutdown);
debug_assert!(sync.spawned() == 0);
let Some(thread) = NonNull::new(self.threads.swap(ptr::null_mut(), Ordering::Acquire))
else {
return;
};
bun_ptr::BackRef::from(thread).join_event.notify();
}
}
#[derive(thiserror::Error, Debug, strum::IntoStaticStr)]
enum WaitError {
#[error("Shutdown")]
Shutdown,
}
#[repr(C)]
pub struct Thread {
next: *mut Thread,
target: *mut Thread,
join_event: Event,
run_queue: node::Queue,
idle_queue: node::Queue,
run_buffer: node::Buffer,
thread_pool: bun_ptr::BackRef<ThreadPool>,
}
thread_local! {
static CURRENT: Cell<*mut Thread> = const { Cell::new(ptr::null_mut()) };
}
struct ThreadRegistration {
pool: bun_ptr::BackRef<ThreadPool>,
thread: *mut Thread,
}
impl ThreadRegistration {
unsafe fn new(pool: &ThreadPool, thread: *mut Thread) -> Self {
CURRENT.with(|c| c.set(thread));
pool.register(thread);
Self {
pool: bun_ptr::BackRef::new(pool),
thread,
}
}
}
impl Drop for ThreadRegistration {
fn drop(&mut self) {
unsafe { ThreadPool::unregister(self.pool.as_ptr(), self.thread) };
CURRENT.with(|c| c.set(ptr::null_mut()));
}
}
static COUNTER: AtomicU32 = AtomicU32::new(0);
#[inline]
fn now_ns() -> u64 {
#[cfg(unix)]
{
unsafe extern "C" {
safe fn clock_gettime(
clk_id: libc::clockid_t,
tp: &mut libc::timespec,
) -> core::ffi::c_int;
}
let mut ts = libc::timespec {
tv_sec: 0,
tv_nsec: 0,
};
clock_gettime(libc::CLOCK_MONOTONIC, &mut ts);
(ts.tv_sec as u64)
.wrapping_mul(1_000_000_000)
.wrapping_add(ts.tv_nsec as u64)
}
#[cfg(not(unix))]
{
use std::time::Instant;
static START: std::sync::OnceLock<Instant> = std::sync::OnceLock::new();
START.get_or_init(Instant::now).elapsed().as_nanos() as u64
}
}
impl Thread {
#[inline]
pub fn current() -> *mut Thread {
CURRENT.with(|c| c.get())
}
pub fn push_idle_task(&self, task: *mut Task) {
let node_ptr = Task::node_of(NonNull::new(task).expect("non-null task"));
let list = node::List {
head: node_ptr,
tail: node_ptr,
};
self.idle_queue.push(&list);
}
fn run(thread_pool: bun_ptr::BackRef<ThreadPool>) {
unsafe extern "C" {
safe fn mi_thread_set_in_threadpool();
}
mi_thread_set_in_threadpool();
{
let mut counter_buf = [0u8; 100];
let int = COUNTER.fetch_add(1, Ordering::SeqCst);
use std::io::Write;
let len = {
let mut cur: &mut [u8] = &mut counter_buf[..99];
let before = cur.len();
match write!(&mut cur, "Bun Pool {}", int) {
Ok(()) => before - cur.len(),
Err(_) => 0,
}
};
let named: &bun_core::ZStr = unsafe {
if len > 0 {
counter_buf[len] = 0;
bun_core::ZStr::from_raw(counter_buf.as_ptr(), len)
} else {
bun_core::ZStr::from_raw(c"Bun Pool".as_ptr().cast(), 8)
}
};
if thread_pool.get().needs_stack_bounds {
Output::Source::configure_named_thread(named);
} else {
Output::Source::configure_named_thread_no_js(named);
}
}
let mut self_ = Thread {
next: ptr::null_mut(),
target: ptr::null_mut(),
join_event: Event::default(),
run_queue: node::Queue::default(),
idle_queue: node::Queue::default(),
run_buffer: node::Buffer::default(),
thread_pool,
};
let self_ptr: *mut Thread = &raw mut self_;
let pool: &ThreadPool = thread_pool.get();
let _registration = unsafe { ThreadRegistration::new(pool, self_ptr) };
let stats = stats_enabled();
let mut is_waking = false;
loop {
let wait_start = if stats { now_ns() } else { 0 };
is_waking = match pool.wait(is_waking) {
Ok(w) => w,
Err(_) => {
unsafe { (*self_ptr).drain_idle_events() };
return;
}
};
if stats {
pool.stats
.idle_ns
.fetch_add(now_ns().wrapping_sub(wait_start), Ordering::Relaxed);
}
while let Some(result) = unsafe { (*self_ptr).pop(pool) } {
if result.pushed || is_waking {
pool.notify(is_waking);
}
is_waking = false;
let task = unsafe { Task::from_node(result.node.as_ptr()) };
let task_start = if stats { now_ns() } else { 0 };
unsafe { ((*task).callback)(task) };
if stats {
pool.stats
.busy_ns
.fetch_add(now_ns().wrapping_sub(task_start), Ordering::Relaxed);
pool.stats.tasks.fetch_add(1, Ordering::Relaxed);
}
}
Output::flush();
unsafe { (*self_ptr).drain_idle_events() };
}
}
pub fn drain_idle_events(&self) {
let Ok(mut consumer) = self.idle_queue.try_acquire_consumer() else {
return;
};
while let Some(node) = consumer.pop() {
let task = unsafe { Task::from_node(node) };
unsafe { ((*task).callback)(task) };
}
}
pub fn pop(&mut self, thread_pool: &ThreadPool) -> Option<node::Stole> {
if let Some(node) = self.run_buffer.pop() {
return Some(node::Stole {
node,
pushed: false,
});
}
if let Some(stole) = self.run_buffer.consume(&self.run_queue) {
return Some(stole);
}
if let Some(stole) = self.run_buffer.consume(&thread_pool.run_queue) {
return Some(stole);
}
let mut num_threads = thread_pool.sync.load(Ordering::Relaxed).spawned();
while num_threads > 0 {
let target = if !self.target.is_null() {
self.target
} else {
let t = thread_pool.threads.load(Ordering::Acquire);
if t.is_null() {
unreachable!();
}
t
};
self.target = unsafe { (*target).next };
if let Some(stole) = self.run_buffer.consume(unsafe { &(*target).run_queue }) {
return Some(stole);
}
if target == std::ptr::from_mut::<Thread>(self) {
num_threads -= 1;
continue;
}
if let Some(stole) = self.run_buffer.steal(unsafe { &(*target).run_buffer }) {
return Some(stole);
}
num_threads -= 1;
}
None
}
}
struct Event {
state: AtomicU32,
}
impl Default for Event {
fn default() -> Self {
Event {
state: AtomicU32::new(Self::EMPTY),
}
}
}
impl Event {
const EMPTY: u32 = 0;
const WAITING: u32 = 1;
pub(crate) const NOTIFIED: u32 = 2;
const SHUTDOWN: u32 = 3;
#[inline(never)]
fn wait(&self) {
let mut acquire_with: u32 = Self::EMPTY;
let mut state = self.state.load(Ordering::Relaxed);
let mut has_shrunk_memory: bool = false;
loop {
if state == Self::SHUTDOWN {
return;
}
if state == Self::NOTIFIED {
match self.state.compare_exchange_weak(
state,
acquire_with,
Ordering::Acquire,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(cur) => state = cur,
}
continue;
}
if state != Self::WAITING {
match self.state.compare_exchange_weak(
state,
Self::WAITING,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => {
}
Err(cur) => {
state = cur;
continue;
}
}
}
let timeout_ns: Option<u64> = if !has_shrunk_memory {
Some(10_000_000_000) } else {
None
};
if Futex::wait(&self.state, Self::WAITING, timeout_ns).is_err() {
has_shrunk_memory = true;
bun_core::Global::mimalloc_cleanup(false);
bun_alloc::wtf::release_fast_malloc_free_memory_for_this_thread();
}
state = self.state.load(Ordering::Relaxed);
acquire_with = Self::WAITING;
}
}
fn notify(&self) {
self.wake(Self::NOTIFIED, 1);
}
fn shutdown(&self) {
self.wake(Self::SHUTDOWN, u32::MAX);
}
fn wake(&self, release_with: u32, wake_threads: u32) {
let state = self.state.swap(release_with, Ordering::Release);
if state == Self::WAITING {
Futex::wake(&self.state, wake_threads);
}
}
}
#[repr(C)]
#[derive(Default)]
pub struct Node {
pub next: *mut Node,
}
pub mod node {
use super::*;
pub struct List {
pub head: NonNull<Node>,
pub tail: NonNull<Node>,
}
#[derive(thiserror::Error, Debug, strum::IntoStaticStr)]
pub(crate) enum ConsumerError {
#[error("Empty")]
Empty,
#[error("Contended")]
Contended,
}
pub(crate) struct Queue {
stack: AtomicUsize,
cache: core::cell::Cell<*mut Node>,
}
unsafe impl core::marker::Sync for Queue {}
unsafe impl Send for Queue {}
impl Default for Queue {
fn default() -> Self {
Queue {
stack: AtomicUsize::new(0),
cache: core::cell::Cell::new(ptr::null_mut()),
}
}
}
impl Queue {
const HAS_CACHE: usize = 0b01;
const IS_CONSUMING: usize = 0b10;
const PTR_MASK: usize = !(Self::HAS_CACHE | Self::IS_CONSUMING);
const _ALIGN_CHECK: () =
assert!(core::mem::align_of::<Node>() >= ((Self::IS_CONSUMING | Self::HAS_CACHE) + 1));
pub(super) fn push(&self, list: &List) {
let List { head, tail } = *list;
let mut stack = self.stack.load(Ordering::Relaxed);
loop {
unsafe {
(*tail.as_ptr()).next = (stack & Self::PTR_MASK) as *mut Node;
}
let mut new_stack = head.as_ptr() as usize;
debug_assert!(new_stack & !Self::PTR_MASK == 0);
new_stack |= stack & !Self::PTR_MASK;
match self.stack.compare_exchange_weak(
stack,
new_stack,
Ordering::Release,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(cur) => stack = cur,
}
}
}
pub(super) fn try_acquire_consumer(&self) -> Result<Consumer<'_>, ConsumerError> {
let mut stack = self.stack.load(Ordering::Relaxed);
loop {
if stack & Self::IS_CONSUMING != 0 {
return Err(ConsumerError::Contended); }
if stack & (Self::HAS_CACHE | Self::PTR_MASK) == 0 {
return Err(ConsumerError::Empty); }
let mut new_stack = stack | Self::HAS_CACHE | Self::IS_CONSUMING;
if stack & Self::HAS_CACHE == 0 {
debug_assert!(stack & Self::PTR_MASK != 0);
new_stack &= !Self::PTR_MASK;
}
match self.stack.compare_exchange_weak(
stack,
new_stack,
Ordering::Acquire,
Ordering::Relaxed,
) {
Ok(_) => {
let cache = self.cache.get();
return Ok(Consumer {
queue: self,
cache: if !cache.is_null() {
cache
} else {
(stack & Self::PTR_MASK) as *mut Node
},
});
}
Err(cur) => stack = cur,
}
}
}
#[inline]
fn release_consumer(&self, consumer: *mut Node) {
let mut remove = Self::IS_CONSUMING;
if consumer.is_null() {
remove |= Self::HAS_CACHE;
}
self.cache.set(consumer);
let stack = self.stack.fetch_sub(remove, Ordering::Release);
debug_assert!(stack & remove != 0);
}
}
pub(super) struct Consumer<'a> {
queue: &'a Queue,
cache: *mut Node,
}
impl Consumer<'_> {
#[inline]
pub(super) fn pop(&mut self) -> Option<*mut Node> {
if !self.cache.is_null() {
let node = self.cache;
self.cache = unsafe { (*node).next };
return Some(node);
}
let mut stack = self.queue.stack.load(Ordering::Relaxed);
debug_assert!(stack & Queue::IS_CONSUMING != 0);
if stack & Queue::PTR_MASK == 0 {
return None;
}
stack = self
.queue
.stack
.swap(Queue::HAS_CACHE | Queue::IS_CONSUMING, Ordering::Acquire);
debug_assert!(stack & Queue::IS_CONSUMING != 0);
debug_assert!(stack & Queue::PTR_MASK != 0);
let node = (stack & Queue::PTR_MASK) as *mut Node;
self.cache = unsafe { (*node).next };
Some(node)
}
}
impl Drop for Consumer<'_> {
#[inline]
fn drop(&mut self) {
self.queue.release_consumer(self.cache);
}
}
#[derive(thiserror::Error, Debug, strum::IntoStaticStr)]
pub(crate) enum BufferPushError {
#[error("Overflow")]
Overflow,
}
type Index = u32;
pub(crate) const CAPACITY: usize = 256;
const _: () = assert!(Index::MAX as usize >= CAPACITY);
const _: () = assert!(CAPACITY.is_power_of_two());
#[repr(C)]
pub(crate) struct Buffer {
head: AtomicU32,
tail: AtomicU32,
array: [AtomicPtr<Node>; CAPACITY],
}
const _: fn() = || {
fn assert<T: Send + core::marker::Sync>() {}
assert::<Buffer>();
};
impl Default for Buffer {
fn default() -> Self {
Buffer {
head: AtomicU32::new(0),
tail: AtomicU32::new(0),
array: [const { AtomicPtr::new(ptr::null_mut()) }; CAPACITY],
}
}
}
pub struct Stole {
pub node: NonNull<Node>,
pub pushed: bool,
}
impl Buffer {
#[inline]
fn tail_raw(&self) -> Index {
self.tail.load(Ordering::Relaxed)
}
#[inline]
fn array_raw(&self, idx: usize) -> *mut Node {
self.array[idx].load(Ordering::Relaxed)
}
pub(super) fn push(&self, list: &mut List) -> Result<(), BufferPushError> {
let mut head = self.head.load(Ordering::Relaxed);
let mut tail = self.tail_raw();
loop {
let mut size = tail.wrapping_sub(head);
debug_assert!(size as usize <= CAPACITY);
if (size as usize) < CAPACITY {
let mut nodes: *mut Node = list.head.as_ptr();
while (size as usize) < CAPACITY {
if nodes.is_null() {
break;
}
let node = nodes;
nodes = unsafe { (*node).next };
self.array[(tail as usize) % CAPACITY].store(node, Ordering::Relaxed);
tail = tail.wrapping_add(1);
size += 1;
}
self.tail.store(tail, Ordering::Release);
match NonNull::new(nodes) {
None => return Ok(()),
Some(h) => list.head = h,
}
core::hint::spin_loop();
head = self.head.load(Ordering::Relaxed);
continue;
}
let mut migrate = size / 2;
match self.head.compare_exchange_weak(
head,
head.wrapping_add(migrate),
Ordering::Acquire,
Ordering::Relaxed,
) {
Err(cur) => head = cur,
Ok(_) => {
let first = self.array_raw((head as usize) % CAPACITY);
while migrate > 0 {
let prev = self.array_raw((head as usize) % CAPACITY);
head = head.wrapping_add(1);
unsafe {
(*prev).next = self.array_raw((head as usize) % CAPACITY);
}
migrate -= 1;
}
let last = self.array_raw((head.wrapping_sub(1) as usize) % CAPACITY);
unsafe {
(*last).next = list.head.as_ptr();
(*list.tail.as_ptr()).next = ptr::null_mut();
}
list.head = unsafe { NonNull::new_unchecked(first) };
return Err(BufferPushError::Overflow);
}
}
}
}
pub(super) fn pop(&self) -> Option<NonNull<Node>> {
let mut head = self.head.load(Ordering::Relaxed);
let tail = self.tail_raw();
loop {
let size = tail.wrapping_sub(head);
debug_assert!(size as usize <= CAPACITY);
if size == 0 {
return None;
}
match self.head.compare_exchange_weak(
head,
head.wrapping_add(1),
Ordering::Acquire,
Ordering::Relaxed,
) {
Err(cur) => head = cur,
Ok(_) => {
let node = self.array_raw((head as usize) % CAPACITY);
return Some(unsafe { NonNull::new_unchecked(node) });
}
}
}
}
pub(super) fn consume(&self, queue: &Queue) -> Option<Stole> {
let Ok(mut consumer) = queue.try_acquire_consumer() else {
return None;
};
let head = self.head.load(Ordering::Relaxed);
let tail = self.tail_raw();
let size = tail.wrapping_sub(head);
debug_assert!(size as usize <= CAPACITY);
debug_assert!(size == 0);
let mut pushed: Index = 0;
while (pushed as usize) < CAPACITY {
let Some(node) = consumer.pop() else {
break;
};
self.array[(tail.wrapping_add(pushed) as usize) % CAPACITY]
.store(node, Ordering::Relaxed);
pushed += 1;
}
let node = match consumer.pop() {
Some(n) => n,
None => 'blk: {
if pushed == 0 {
return None;
}
pushed -= 1;
break 'blk self.array_raw((tail.wrapping_add(pushed) as usize) % CAPACITY);
}
};
if pushed > 0 {
self.tail
.store(tail.wrapping_add(pushed), Ordering::Release);
}
Some(Stole {
node: unsafe { NonNull::new_unchecked(node) },
pushed: pushed > 0,
})
}
pub(super) fn steal(&self, buffer: &Buffer) -> Option<Stole> {
let head = self.head.load(Ordering::Relaxed);
let tail = self.tail_raw();
let size = tail.wrapping_sub(head);
debug_assert!(size as usize <= CAPACITY);
debug_assert!(size == 0);
loop {
let buffer_head = buffer.head.load(Ordering::Acquire);
let buffer_tail = buffer.tail.load(Ordering::Acquire);
let buffer_size = buffer_tail.wrapping_sub(buffer_head);
if buffer_size as usize > CAPACITY {
core::hint::spin_loop();
continue;
}
let steal_size = buffer_size - (buffer_size / 2);
if steal_size == 0 {
return None;
}
for i in 0..steal_size {
let node = buffer.array[(buffer_head.wrapping_add(i) as usize) % CAPACITY]
.load(Ordering::Relaxed);
self.array[(tail.wrapping_add(i) as usize) % CAPACITY]
.store(node, Ordering::Relaxed);
}
match buffer.head.compare_exchange(
buffer_head,
buffer_head.wrapping_add(steal_size),
Ordering::AcqRel,
Ordering::Relaxed,
) {
Err(_) => {
core::hint::spin_loop();
}
Ok(_) => {
let pushed = steal_size - 1;
let node = self.array_raw((tail.wrapping_add(pushed) as usize) % CAPACITY);
if pushed > 0 {
self.tail
.store(tail.wrapping_add(pushed), Ordering::Release);
}
return Some(Stole {
node: unsafe { NonNull::new_unchecked(node) },
pushed: pushed > 0,
});
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use core::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as StdOrdering};
use std::time::{Duration, Instant};
#[derive(Default)]
struct ParkFlags {
started: AtomicBool,
release: AtomicBool,
}
struct ParkedFsTask {
task: Task,
flags: Arc<ParkFlags>,
}
unsafe fn park_cb(task: *mut Task) {
let park: &ParkedFsTask = unsafe {
&*bun_core::from_field_ptr!(ParkedFsTask, task, task)
};
park.flags.started.store(true, StdOrdering::SeqCst);
let deadline = Instant::now() + Duration::from_secs(30);
while !park.flags.release.load(StdOrdering::Relaxed) {
if Instant::now() > deadline {
break;
}
core::hint::spin_loop();
}
}
fn wait_until_started(flags: &ParkFlags) {
let deadline = Instant::now() + Duration::from_secs(10);
while !flags.started.load(StdOrdering::SeqCst) {
assert!(Instant::now() < deadline, "park task never started on a worker");
core::hint::spin_loop();
}
}
fn init_output_for_pool_tests() {
bun_core::output::init_test();
}
#[test]
fn each_waits_for_its_batch_not_the_whole_pool() {
init_output_for_pool_tests();
let pool = ThreadPool::init(Config {
stack_size: 1 << 20,
max_threads: 2,
});
let flags = Arc::new(ParkFlags::default());
let mut park = ParkedFsTask {
task: Task {
node: Node::default(),
callback: park_cb,
},
flags: flags.clone(),
};
pool.schedule(Batch::from(&raw mut park.task));
wait_until_started(&flags);
let start = Instant::now();
let mut values = [0u32; 4];
pool.each_ptr(
(),
|_: &(), value: *mut u32, _: usize| {
unsafe { *value = 1 };
},
&mut values,
);
let elapsed = start.elapsed();
assert!(values.iter().all(|&v| v == 1), "each() did not run every task");
assert!(
elapsed < Duration::from_secs(5),
"each() waited for the parked task (pool-wide wait regression): {elapsed:?}"
);
assert!(
!flags.release.load(StdOrdering::SeqCst),
"test released the park before asserting; timing invalid"
);
flags.release.store(true, StdOrdering::SeqCst);
}
#[test]
fn counted_task_batch_waits_only_for_itself() {
struct BatchItem {
counted: CountedTask,
slot: *mut u32,
}
unsafe fn item_cb(task: *mut Task) {
let item: &mut BatchItem =
unsafe { &mut *bun_core::from_field_ptr!(BatchItem, counted, task) };
unsafe { *item.slot += 1 };
}
init_output_for_pool_tests();
let pool = ThreadPool::init(Config {
stack_size: 1 << 20,
max_threads: 2,
});
let flags = Arc::new(ParkFlags::default());
let mut park = ParkedFsTask {
task: Task {
node: Node::default(),
callback: park_cb,
},
flags: flags.clone(),
};
pool.schedule(Batch::from(&raw mut park.task));
wait_until_started(&flags);
let mut results = [0u32; 3];
let group = WaitGroup::init_with_count(results.len());
let mut items: Vec<BatchItem> = Vec::with_capacity(results.len());
let mut batch = Batch::default();
for slot in results.iter_mut() {
items.push(BatchItem {
counted: CountedTask::new(item_cb, &group),
slot: &raw mut *slot,
});
let last = items.last_mut().unwrap();
batch.push(Batch::from(&raw mut last.counted.task));
}
let start = Instant::now();
pool.schedule(batch);
group.wait();
let elapsed = start.elapsed();
assert!(
results.iter().all(|&r| r == 1),
"counted batch did not run every task: {results:?}"
);
assert!(
elapsed < Duration::from_secs(5),
"group.wait() waited for the parked task (pool-wide wait regression): {elapsed:?}"
);
flags.release.store(true, StdOrdering::SeqCst);
}
#[test]
fn counted_task_group_wait_happens_after_callbacks() {
init_output_for_pool_tests();
let pool = ThreadPool::init(Config {
stack_size: 1 << 20,
max_threads: 2,
});
let group = WaitGroup::init_with_count(1);
struct Solo {
counted: CountedTask,
counter: *const AtomicUsize,
}
unsafe fn solo_cb(task: *mut Task) {
let solo: &mut Solo =
unsafe { &mut *bun_core::from_field_ptr!(Solo, counted, task) };
unsafe {
(*solo.counter).fetch_add(1, StdOrdering::SeqCst);
}
}
let counter = AtomicUsize::new(0);
let mut solo = Solo {
counted: CountedTask::new(solo_cb, &group),
counter: &counter,
};
pool.schedule(Batch::from(&raw mut solo.counted.task));
group.wait();
assert_eq!(counter.load(StdOrdering::SeqCst), 1);
}
}