use std::cell::{Cell, RefCell};
use std::collections::{BinaryHeap, HashMap, VecDeque};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, LazyLock};
use std::time::{Duration, Instant};
use catallaxy_runtime::object::Wake;
use catallaxy_runtime::{RtResult, Trap, Value};
use catallaxy_vm::Vm;
use corosensei::stack::DefaultStack;
use corosensei::{Coroutine, CoroutineResult, Yielder};
use parking_lot::{Condvar, Mutex};
const FIBER_NATIVE_STACK: usize = 256 * 1024;
const FIBER_REGISTERS: usize = 64 * 1024;
#[cfg(catx_tsan)]
mod tsan {
use std::cell::Cell;
use std::ffi::c_void;
extern "C" {
fn __tsan_get_current_fiber() -> *mut c_void;
fn __tsan_create_fiber(flags: u32) -> *mut c_void;
fn __tsan_destroy_fiber(fiber: *mut c_void);
fn __tsan_switch_to_fiber(fiber: *mut c_void, flags: u32);
}
thread_local! {
static MAIN: Cell<*mut c_void> = const { Cell::new(std::ptr::null_mut()) };
}
pub struct Handle(*mut c_void);
impl Handle {
pub fn new() -> Self {
Handle(unsafe { __tsan_create_fiber(0) })
}
}
impl Drop for Handle {
fn drop(&mut self) {
unsafe { __tsan_destroy_fiber(self.0) }
}
}
pub fn enter(h: &Handle) {
unsafe {
MAIN.with(|m| m.set(__tsan_get_current_fiber()));
__tsan_switch_to_fiber(h.0, 0);
}
}
pub fn leave() {
unsafe { __tsan_switch_to_fiber(MAIN.with(|m| m.get()), 0) }
}
}
#[cfg(not(catx_tsan))]
mod tsan {
pub struct Handle;
impl Handle {
pub fn new() -> Self {
Handle
}
}
pub fn enter(_: &Handle) {}
pub fn leave() {}
}
struct LeaveOnExit;
impl Drop for LeaveOnExit {
fn drop(&mut self) {
tsan::leave();
}
}
unsafe fn suspend(y: *const Yielder<(), ()>) {
tsan::leave();
(*y).suspend(());
}
static NEXT_ID: AtomicU64 = AtomicU64::new(1);
const ROOT: u64 = 0;
struct Shared {
woken: Mutex<Vec<u64>>,
cv: Condvar,
}
impl Shared {
fn wake(&self, id: u64) {
self.woken.lock().push(id);
self.cv.notify_one();
}
}
#[derive(Clone)]
struct Waker {
shared: Arc<Shared>,
id: u64,
flag: Arc<AtomicBool>,
}
impl Waker {
fn wake(&self) {
self.flag.store(true, Ordering::Release);
self.shared.wake(self.id);
}
fn into_wake(self) -> Wake {
Arc::new(move || self.wake())
}
}
enum Outcome {
Running,
Done {
in_handle: bool,
},
Panicked(String),
Joined,
}
struct Ctl {
outcome: Mutex<Outcome>,
joiners: Mutex<Vec<Wake>>,
}
struct Fiber {
coro: Coroutine<(), (), (), DefaultStack>,
vm: *mut Vm,
flag: Arc<AtomicBool>,
yielder: Cell<*const Yielder<(), ()>>,
tsan: tsan::Handle,
}
impl Drop for Fiber {
fn drop(&mut self) {
if self.coro.started() && !self.coro.done() {
tsan::enter(&self.tsan);
}
}
}
struct Scheduler {
shared: Arc<Shared>,
root_flag: Arc<AtomicBool>,
fibers: HashMap<u64, Fiber>,
ready: VecDeque<u64>,
queued: std::collections::HashSet<u64>,
timers: BinaryHeap<std::cmp::Reverse<(Instant, u64)>>,
ctl: HashMap<u64, Arc<Ctl>>,
}
impl Scheduler {
fn enqueue(&mut self, id: u64) {
if self.queued.insert(id) {
self.ready.push_back(id);
}
}
fn dequeue(&mut self) -> Option<u64> {
let id = self.ready.pop_front()?;
self.queued.remove(&id);
Some(id)
}
fn new() -> Self {
Scheduler {
shared: Arc::new(Shared { woken: Mutex::new(Vec::new()), cv: Condvar::new() }),
root_flag: Arc::new(AtomicBool::new(false)),
fibers: HashMap::new(),
ready: VecDeque::new(),
queued: std::collections::HashSet::new(),
timers: BinaryHeap::new(),
ctl: HashMap::new(),
}
}
}
thread_local! {
static SCHED: RefCell<Option<Scheduler>> = const { RefCell::new(None) };
static CURRENT: Cell<u64> = const { Cell::new(ROOT) };
static YIELDER: Cell<*const Yielder<(), ()>> = const { Cell::new(std::ptr::null()) };
}
fn with_sched<R>(f: impl FnOnce(&mut Scheduler) -> R) -> R {
SCHED.with(|s| f(s.borrow_mut().get_or_insert_with(Scheduler::new)))
}
pub fn in_fiber() -> bool {
CURRENT.with(|c| c.get()) != ROOT
}
pub fn has_fibers() -> bool {
SCHED.with(|s| s.borrow().as_ref().is_some_and(|s| !s.fibers.is_empty()))
}
fn current_waker() -> Waker {
let id = CURRENT.with(|c| c.get());
let flag = if id == ROOT {
with_sched(|s| s.root_flag.clone())
} else {
RUNNING_FLAG.with(|f| f.borrow().clone()).expect("running fiber flag")
};
Waker { shared: with_sched(|s| s.shared.clone()), id, flag }
}
thread_local! {
static RUNNING_FLAG: RefCell<Option<Arc<AtomicBool>>> = const { RefCell::new(None) };
}
const ID_FIELD: usize = 0;
const RESULT_FIELD: usize = 1;
pub fn spawn(vm: &mut Vm, handle: Value, closure: Value) -> RtResult<u64> {
let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
let h = unsafe { &*handle.as_ptr() };
catallaxy_runtime::ownership::check_write(&vm.m, h)?;
unsafe { catallaxy_runtime::object::as_instance(handle) }.set(ID_FIELD, Value::small_int(id as i64));
let fvm = Box::into_raw(vm.new_fiber(FIBER_REGISTERS));
unsafe {
(*fvm).extra_roots().push(closure);
(*fvm).extra_roots().push(handle);
}
let ctl = Arc::new(Ctl { outcome: Mutex::new(Outcome::Running), joiners: Mutex::new(Vec::new()) });
let ctl2 = ctl.clone();
let vm_ptr = fvm as usize;
let stack = DefaultStack::new(FIBER_NATIVE_STACK).expect("fiber stack");
let coro = Coroutine::with_stack(stack, move |y: &Yielder<(), ()>, ()| {
let _leave = LeaveOnExit;
YIELDER.with(|c| c.set(y as *const _));
let vm = unsafe { &mut *(vm_ptr as *mut Vm) };
let closure = vm.extra_roots()[0];
let r = vm.call_value(closure, &[]);
let outcome = match r {
Ok(v) => {
let handle = vm.extra_roots()[1];
let h = unsafe { &*handle.as_ptr() };
let mine = h.state() == catallaxy_gc::ObjState::Owned && h.owner() == vm.tid();
if mine && catallaxy_runtime::ownership::check_write(&vm.m, h).is_ok() {
unsafe { catallaxy_runtime::object::as_instance(handle) }.set(RESULT_FIELD, v);
Outcome::Done { in_handle: true }
} else {
vm.pin(id, v);
Outcome::Done { in_handle: false }
}
}
Err(Trap::Panic(msg)) => {
let mut full = msg;
for line in vm.panic_trace.take().unwrap_or_default().iter().take(8) {
full.push_str(&format!("\n at {line}"));
}
Outcome::Panicked(full)
}
Err(Trap::Exit(code)) => {
crate::flush();
std::process::exit(code);
}
Err(Trap::Unwind { .. }) => Outcome::Panicked("break or return out of a fiber block".into()),
};
*ctl2.outcome.lock() = outcome;
for w in std::mem::take(&mut *ctl2.joiners.lock()) {
w();
}
});
with_sched(|s| {
s.fibers.insert(
id,
Fiber {
coro,
vm: fvm,
flag: Arc::new(AtomicBool::new(false)),
yielder: Cell::new(std::ptr::null()),
tsan: tsan::Handle::new(),
},
);
s.ctl.insert(id, ctl);
s.enqueue(id);
});
Ok(id)
}
fn resume(id: u64) {
let Some(fiber) = with_sched(|s| s.fibers.remove(&id)) else { return };
let prev_y = YIELDER.with(|c| c.replace(fiber.yielder.get()));
CURRENT.with(|c| c.set(id));
RUNNING_FLAG.with(|f| *f.borrow_mut() = Some(fiber.flag.clone()));
let mut fiber = fiber;
tsan::enter(&fiber.tsan);
let r = fiber.coro.resume(());
fiber.yielder.set(YIELDER.with(|c| c.get()));
RUNNING_FLAG.with(|f| *f.borrow_mut() = None);
CURRENT.with(|c| c.set(ROOT));
YIELDER.with(|c| c.set(prev_y));
match r {
CoroutineResult::Yield(()) => {
let woken = fiber.flag.load(Ordering::Acquire);
with_sched(|s| {
s.fibers.insert(id, fiber);
if woken {
s.enqueue(id);
}
});
}
CoroutineResult::Return(()) => {
let vm = fiber.vm;
drop(fiber);
Vm::release_fiber(unsafe { Box::from_raw(vm) });
}
}
}
fn absorb() {
let now = Instant::now();
with_sched(|s| {
let woken = std::mem::take(&mut *s.shared.woken.lock());
for id in woken {
if id != ROOT && s.fibers.contains_key(&id) {
s.enqueue(id);
}
}
while let Some(std::cmp::Reverse((t, id))) = s.timers.peek().copied() {
if t > now {
break;
}
s.timers.pop();
if id == ROOT {
s.root_flag.store(true, Ordering::Release);
} else if let Some(f) = s.fibers.get(&id) {
f.flag.store(true, Ordering::Release);
s.enqueue(id);
}
}
});
}
fn run_until(vm: &mut Vm, mut done: impl FnMut() -> bool) {
debug_assert!(!in_fiber());
loop {
absorb();
if done() {
return;
}
if let Some(id) = with_sched(|s| s.dequeue()) {
resume(id);
catallaxy_gc::safepoint(&vm.m.rec);
continue;
}
let (shared, deadline) = with_sched(|s| (s.shared.clone(), s.timers.peek().map(|r| r.0 .0)));
catallaxy_gc::blocking(&vm.m.rec, || {
let mut w = shared.woken.lock();
if w.is_empty() {
match deadline {
Some(t) => {
let now = Instant::now();
if t > now {
shared.cv.wait_for(&mut w, t - now);
}
}
None => {
shared.cv.wait_for(&mut w, Duration::from_millis(100));
}
}
}
});
}
}
fn park() {
let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("park outside a fiber");
if flag.swap(false, Ordering::AcqRel) {
return;
}
let y = YIELDER.with(|c| c.get());
unsafe { suspend(y) };
flag.store(false, Ordering::Release);
}
pub fn drain(vm: &mut Vm) {
if in_fiber() || !has_fibers() {
return;
}
run_until(vm, || !has_fibers());
}
pub fn wait_event(vm: &mut Vm, arm: impl FnOnce(Wake) -> bool) -> bool {
if !vm.park_allowed() {
return false;
}
if in_fiber() {
let w = current_waker();
if arm(w.into_wake()) {
park();
}
return true;
}
if !has_fibers() {
return false;
}
let w = current_waker();
let flag = w.flag.clone();
if arm(w.into_wake()) {
run_until(vm, || flag.swap(false, Ordering::AcqRel));
}
true
}
pub fn wait_event_until(vm: &mut Vm, deadline: Option<Instant>, arm: impl FnOnce(Wake) -> bool) -> Option<bool> {
let Some(deadline) = deadline else {
return wait_event(vm, arm).then_some(false);
};
if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
return None;
}
let id = CURRENT.with(|c| c.get());
let w = current_waker();
let flag = w.flag.clone();
if !arm(w.into_wake()) {
return Some(false);
}
with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
if in_fiber() {
park();
} else {
run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
}
Some(Instant::now() >= deadline)
}
pub fn sleep(vm: &mut Vm, d: Duration) {
let deadline = Instant::now() + d;
if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
crate::blocking(vm, || std::thread::sleep(d));
return;
}
let id = CURRENT.with(|c| c.get());
while Instant::now() < deadline {
with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
if in_fiber() {
park();
} else {
let flag = with_sched(|s| s.root_flag.clone());
run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
}
}
}
pub fn yield_now(vm: &mut Vm) {
if in_fiber() {
if !vm.park_allowed() {
return;
}
let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("running fiber flag");
flag.store(true, Ordering::Release);
let y = YIELDER.with(|c| c.get());
unsafe { suspend(y) };
flag.store(false, Ordering::Release);
} else if has_fibers() {
let n = with_sched(|s| s.ready.len());
for _ in 0..n {
absorb();
match with_sched(|s| s.dequeue()) {
Some(id) => resume(id),
None => break,
}
}
}
}
pub fn join(vm: &mut Vm, handle: Value) -> RtResult<Result<Value, String>> {
let id = unsafe { catallaxy_runtime::object::as_instance(handle) }.get(ID_FIELD).as_small_int() as u64;
let Some(ctl) = with_sched(|s| s.ctl.get(&id).cloned()) else {
return Err(Trap::Panic(
"FiberError: no such fiber on this thread (fibers are joined on the thread that spawned them)".into(),
));
};
loop {
if !matches!(*ctl.outcome.lock(), Outcome::Running) {
break;
}
let ctl2 = ctl.clone();
let cooperative = wait_event(vm, move |wake| {
let o = ctl2.outcome.lock();
if !matches!(*o, Outcome::Running) {
return false;
}
ctl2.joiners.lock().push(wake);
true
});
if !cooperative {
return Err(Trap::Panic("FiberError: join while holding a lock would deadlock".into()));
}
}
with_sched(|s| s.ctl.remove(&id));
let outcome = std::mem::replace(&mut *ctl.outcome.lock(), Outcome::Joined);
Ok(match outcome {
Outcome::Done { in_handle: true } => {
Ok(unsafe { catallaxy_runtime::object::as_instance(handle) }.get(RESULT_FIELD))
}
Outcome::Done { in_handle: false } => Ok(vm.unpin(id).unwrap_or(Value::NIL)),
Outcome::Panicked(m) => Err(m),
Outcome::Joined => Err("fiber was already joined".into()),
Outcome::Running => unreachable!(),
})
}
pub fn is_done(id: u64) -> bool {
with_sched(|s| s.ctl.get(&id).is_none_or(|c| !matches!(*c.outcome.lock(), Outcome::Running)))
}
type Job = Box<dyn FnOnce() + Send + 'static>;
struct IoPool {
queue: Mutex<VecDeque<Job>>,
cv: Condvar,
}
static IO: LazyLock<IoPool> = LazyLock::new(|| {
let n = std::env::var("CATX_IO_THREADS").ok().and_then(|v| v.parse::<usize>().ok()).unwrap_or(16).max(1);
for i in 0..n {
let _ = std::thread::Builder::new().name(format!("catx-io-{i}")).spawn(|| loop {
let job = {
let mut q = IO.queue.lock();
loop {
if let Some(j) = q.pop_front() {
break j;
}
IO.cv.wait(&mut q);
}
};
job();
});
}
IoPool { queue: Mutex::new(VecDeque::new()), cv: Condvar::new() }
});
pub fn blocking_io<R: Send>(vm: &mut Vm, f: impl FnOnce() -> R + Send) -> R {
if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
return crate::blocking(vm, f);
}
let done = Arc::new(AtomicBool::new(false));
let mut result: Option<std::thread::Result<R>> = None;
let w = current_waker();
{
let slot = &mut result as *mut Option<std::thread::Result<R>> as usize;
let (done2, w2) = (done.clone(), w.clone());
let job: Box<dyn FnOnce() + Send + '_> = Box::new(move || {
let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
unsafe { *(slot as *mut Option<std::thread::Result<R>>) = Some(r) };
done2.store(true, Ordering::Release);
w2.wake();
});
let job: Job = unsafe { std::mem::transmute::<Box<dyn FnOnce() + Send + '_>, Job>(job) };
IO.queue.lock().push_back(job);
IO.cv.notify_one();
}
while !done.load(Ordering::Acquire) {
if in_fiber() {
park();
} else {
let flag = w.flag.clone();
let d = done.clone();
run_until(vm, || flag.swap(false, Ordering::AcqRel) || d.load(Ordering::Acquire));
}
}
match result.take().expect("I/O job result") {
Ok(r) => r,
Err(p) => std::panic::resume_unwind(p),
}
}