1use std::cell::{Cell, RefCell};
29use std::collections::{BinaryHeap, HashMap, VecDeque};
30use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
31use std::sync::{Arc, LazyLock};
32use std::time::{Duration, Instant};
33
34use catallaxy_runtime::object::Wake;
35use catallaxy_runtime::{RtResult, Trap, Value};
36use catallaxy_vm::Vm;
37use corosensei::stack::DefaultStack;
38use corosensei::{Coroutine, CoroutineResult, Yielder};
39use parking_lot::{Condvar, Mutex};
40
41const FIBER_NATIVE_STACK: usize = 256 * 1024;
43const FIBER_REGISTERS: usize = 64 * 1024;
45
46#[cfg(catx_tsan)]
58mod tsan {
59 use std::cell::Cell;
60 use std::ffi::c_void;
61
62 extern "C" {
63 fn __tsan_get_current_fiber() -> *mut c_void;
64 fn __tsan_create_fiber(flags: u32) -> *mut c_void;
65 fn __tsan_destroy_fiber(fiber: *mut c_void);
66 fn __tsan_switch_to_fiber(fiber: *mut c_void, flags: u32);
67 }
68
69 thread_local! {
70 static MAIN: Cell<*mut c_void> = const { Cell::new(std::ptr::null_mut()) };
72 }
73
74 pub struct Handle(*mut c_void);
76
77 impl Handle {
78 pub fn new() -> Self {
79 Handle(unsafe { __tsan_create_fiber(0) })
80 }
81 }
82
83 impl Drop for Handle {
84 fn drop(&mut self) {
85 unsafe { __tsan_destroy_fiber(self.0) }
86 }
87 }
88
89 pub fn enter(h: &Handle) {
91 unsafe {
92 MAIN.with(|m| m.set(__tsan_get_current_fiber()));
93 __tsan_switch_to_fiber(h.0, 0);
94 }
95 }
96
97 pub fn leave() {
99 unsafe { __tsan_switch_to_fiber(MAIN.with(|m| m.get()), 0) }
100 }
101}
102
103#[cfg(not(catx_tsan))]
104mod tsan {
105 pub struct Handle;
106 impl Handle {
107 pub fn new() -> Self {
108 Handle
109 }
110 }
111 pub fn enter(_: &Handle) {}
112 pub fn leave() {}
113}
114
115struct LeaveOnExit;
117impl Drop for LeaveOnExit {
118 fn drop(&mut self) {
119 tsan::leave();
120 }
121}
122
123unsafe fn suspend(y: *const Yielder<(), ()>) {
128 tsan::leave();
129 (*y).suspend(());
130}
131
132static NEXT_ID: AtomicU64 = AtomicU64::new(1);
133
134const ROOT: u64 = 0;
136
137struct Shared {
139 woken: Mutex<Vec<u64>>,
140 cv: Condvar,
141}
142
143impl Shared {
144 fn wake(&self, id: u64) {
145 self.woken.lock().push(id);
146 self.cv.notify_one();
147 }
148}
149
150#[derive(Clone)]
151struct Waker {
152 shared: Arc<Shared>,
153 id: u64,
154 flag: Arc<AtomicBool>,
156}
157
158impl Waker {
159 fn wake(&self) {
160 self.flag.store(true, Ordering::Release);
161 self.shared.wake(self.id);
162 }
163 fn into_wake(self) -> Wake {
164 Arc::new(move || self.wake())
165 }
166}
167
168enum Outcome {
169 Running,
170 Done {
173 in_handle: bool,
174 },
175 Panicked(String),
176 Joined,
177}
178
179struct Ctl {
181 outcome: Mutex<Outcome>,
182 joiners: Mutex<Vec<Wake>>,
183}
184
185struct Fiber {
186 coro: Coroutine<(), (), (), DefaultStack>,
187 vm: *mut Vm,
188 flag: Arc<AtomicBool>,
189 yielder: Cell<*const Yielder<(), ()>>,
190 tsan: tsan::Handle,
192}
193
194impl Drop for Fiber {
195 fn drop(&mut self) {
196 if self.coro.started() && !self.coro.done() {
199 tsan::enter(&self.tsan);
200 }
201 }
202}
203
204struct Scheduler {
205 shared: Arc<Shared>,
206 root_flag: Arc<AtomicBool>,
207 fibers: HashMap<u64, Fiber>,
209 ready: VecDeque<u64>,
210 queued: std::collections::HashSet<u64>,
212 timers: BinaryHeap<std::cmp::Reverse<(Instant, u64)>>,
214 ctl: HashMap<u64, Arc<Ctl>>,
215}
216
217impl Scheduler {
218 fn enqueue(&mut self, id: u64) {
219 if self.queued.insert(id) {
220 self.ready.push_back(id);
221 }
222 }
223
224 fn dequeue(&mut self) -> Option<u64> {
225 let id = self.ready.pop_front()?;
226 self.queued.remove(&id);
227 Some(id)
228 }
229
230 fn new() -> Self {
231 Scheduler {
232 shared: Arc::new(Shared { woken: Mutex::new(Vec::new()), cv: Condvar::new() }),
233 root_flag: Arc::new(AtomicBool::new(false)),
234 fibers: HashMap::new(),
235 ready: VecDeque::new(),
236 queued: std::collections::HashSet::new(),
237 timers: BinaryHeap::new(),
238 ctl: HashMap::new(),
239 }
240 }
241}
242
243thread_local! {
244 static SCHED: RefCell<Option<Scheduler>> = const { RefCell::new(None) };
245 static CURRENT: Cell<u64> = const { Cell::new(ROOT) };
247 static YIELDER: Cell<*const Yielder<(), ()>> = const { Cell::new(std::ptr::null()) };
248}
249
250fn with_sched<R>(f: impl FnOnce(&mut Scheduler) -> R) -> R {
251 SCHED.with(|s| f(s.borrow_mut().get_or_insert_with(Scheduler::new)))
252}
253
254pub fn in_fiber() -> bool {
256 CURRENT.with(|c| c.get()) != ROOT
257}
258
259pub fn has_fibers() -> bool {
261 SCHED.with(|s| s.borrow().as_ref().is_some_and(|s| !s.fibers.is_empty()))
262}
263
264fn current_waker() -> Waker {
266 let id = CURRENT.with(|c| c.get());
267 let flag = if id == ROOT {
268 with_sched(|s| s.root_flag.clone())
269 } else {
270 RUNNING_FLAG.with(|f| f.borrow().clone()).expect("running fiber flag")
271 };
272 Waker { shared: with_sched(|s| s.shared.clone()), id, flag }
273}
274
275thread_local! {
276 static RUNNING_FLAG: RefCell<Option<Arc<AtomicBool>>> = const { RefCell::new(None) };
278}
279
280const ID_FIELD: usize = 0;
284const RESULT_FIELD: usize = 1;
285
286pub fn spawn(vm: &mut Vm, handle: Value, closure: Value) -> RtResult<u64> {
289 let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
290 let h = unsafe { &*handle.as_ptr() };
291 catallaxy_runtime::ownership::check_write(&vm.m, h)?;
292 unsafe { catallaxy_runtime::object::as_instance(handle) }.set(ID_FIELD, Value::small_int(id as i64));
293 let fvm = Box::into_raw(vm.new_fiber(FIBER_REGISTERS));
294 unsafe {
295 (*fvm).extra_roots().push(closure);
296 (*fvm).extra_roots().push(handle);
297 }
298 let ctl = Arc::new(Ctl { outcome: Mutex::new(Outcome::Running), joiners: Mutex::new(Vec::new()) });
299 let ctl2 = ctl.clone();
300 let vm_ptr = fvm as usize;
301 let stack = DefaultStack::new(FIBER_NATIVE_STACK).expect("fiber stack");
302 let coro = Coroutine::with_stack(stack, move |y: &Yielder<(), ()>, ()| {
303 let _leave = LeaveOnExit;
304 YIELDER.with(|c| c.set(y as *const _));
305 let vm = unsafe { &mut *(vm_ptr as *mut Vm) };
307 let closure = vm.extra_roots()[0];
308 let r = vm.call_value(closure, &[]);
309 let outcome = match r {
310 Ok(v) => {
311 let handle = vm.extra_roots()[1];
314 let h = unsafe { &*handle.as_ptr() };
315 let mine = h.state() == catallaxy_gc::ObjState::Owned && h.owner() == vm.tid();
316 if mine && catallaxy_runtime::ownership::check_write(&vm.m, h).is_ok() {
317 unsafe { catallaxy_runtime::object::as_instance(handle) }.set(RESULT_FIELD, v);
318 Outcome::Done { in_handle: true }
319 } else {
320 vm.pin(id, v);
321 Outcome::Done { in_handle: false }
322 }
323 }
324 Err(Trap::Panic(msg)) => {
325 let mut full = msg;
326 for line in vm.panic_trace.take().unwrap_or_default().iter().take(8) {
327 full.push_str(&format!("\n at {line}"));
328 }
329 Outcome::Panicked(full)
330 }
331 Err(Trap::Exit(code)) => {
332 crate::flush();
333 std::process::exit(code);
334 }
335 Err(Trap::Unwind { .. }) => Outcome::Panicked("break or return out of a fiber block".into()),
336 };
337 *ctl2.outcome.lock() = outcome;
338 for w in std::mem::take(&mut *ctl2.joiners.lock()) {
339 w();
340 }
341 });
342 with_sched(|s| {
343 s.fibers.insert(
344 id,
345 Fiber {
346 coro,
347 vm: fvm,
348 flag: Arc::new(AtomicBool::new(false)),
349 yielder: Cell::new(std::ptr::null()),
350 tsan: tsan::Handle::new(),
351 },
352 );
353 s.ctl.insert(id, ctl);
354 s.enqueue(id);
355 });
356 Ok(id)
357}
358
359fn resume(id: u64) {
361 let Some(fiber) = with_sched(|s| s.fibers.remove(&id)) else { return };
362 let prev_y = YIELDER.with(|c| c.replace(fiber.yielder.get()));
363 CURRENT.with(|c| c.set(id));
364 RUNNING_FLAG.with(|f| *f.borrow_mut() = Some(fiber.flag.clone()));
365 let mut fiber = fiber;
366 tsan::enter(&fiber.tsan);
367 let r = fiber.coro.resume(());
368 fiber.yielder.set(YIELDER.with(|c| c.get()));
369 RUNNING_FLAG.with(|f| *f.borrow_mut() = None);
370 CURRENT.with(|c| c.set(ROOT));
371 YIELDER.with(|c| c.set(prev_y));
372 match r {
373 CoroutineResult::Yield(()) => {
374 let woken = fiber.flag.load(Ordering::Acquire);
375 with_sched(|s| {
376 s.fibers.insert(id, fiber);
377 if woken {
378 s.enqueue(id);
379 }
380 });
381 }
382 CoroutineResult::Return(()) => {
383 let vm = fiber.vm;
384 drop(fiber);
385 Vm::release_fiber(unsafe { Box::from_raw(vm) });
387 }
388 }
389}
390
391fn absorb() {
393 let now = Instant::now();
394 with_sched(|s| {
395 let woken = std::mem::take(&mut *s.shared.woken.lock());
396 for id in woken {
397 if id != ROOT && s.fibers.contains_key(&id) {
398 s.enqueue(id);
399 }
400 }
401 while let Some(std::cmp::Reverse((t, id))) = s.timers.peek().copied() {
402 if t > now {
403 break;
404 }
405 s.timers.pop();
406 if id == ROOT {
407 s.root_flag.store(true, Ordering::Release);
408 } else if let Some(f) = s.fibers.get(&id) {
409 f.flag.store(true, Ordering::Release);
410 s.enqueue(id);
411 }
412 }
413 });
414}
415
416fn run_until(vm: &mut Vm, mut done: impl FnMut() -> bool) {
420 debug_assert!(!in_fiber());
421 loop {
422 absorb();
423 if done() {
424 return;
425 }
426 if let Some(id) = with_sched(|s| s.dequeue()) {
427 resume(id);
428 catallaxy_gc::safepoint(&vm.m.rec);
429 continue;
430 }
431 let (shared, deadline) = with_sched(|s| (s.shared.clone(), s.timers.peek().map(|r| r.0 .0)));
432 catallaxy_gc::blocking(&vm.m.rec, || {
433 let mut w = shared.woken.lock();
434 if w.is_empty() {
435 match deadline {
436 Some(t) => {
437 let now = Instant::now();
438 if t > now {
439 shared.cv.wait_for(&mut w, t - now);
440 }
441 }
442 None => {
443 shared.cv.wait_for(&mut w, Duration::from_millis(100));
444 }
445 }
446 }
447 });
448 }
449}
450
451fn park() {
453 let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("park outside a fiber");
454 if flag.swap(false, Ordering::AcqRel) {
455 return;
456 }
457 let y = YIELDER.with(|c| c.get());
458 unsafe { suspend(y) };
460 flag.store(false, Ordering::Release);
461}
462
463pub fn drain(vm: &mut Vm) {
465 if in_fiber() || !has_fibers() {
466 return;
467 }
468 run_until(vm, || !has_fibers());
469}
470
471pub fn wait_event(vm: &mut Vm, arm: impl FnOnce(Wake) -> bool) -> bool {
479 if !vm.park_allowed() {
480 return false;
481 }
482 if in_fiber() {
483 let w = current_waker();
484 if arm(w.into_wake()) {
485 park();
486 }
487 return true;
488 }
489 if !has_fibers() {
490 return false;
491 }
492 let w = current_waker();
493 let flag = w.flag.clone();
494 if arm(w.into_wake()) {
495 run_until(vm, || flag.swap(false, Ordering::AcqRel));
496 }
497 true
498}
499
500pub fn wait_event_until(vm: &mut Vm, deadline: Option<Instant>, arm: impl FnOnce(Wake) -> bool) -> Option<bool> {
504 let Some(deadline) = deadline else {
505 return wait_event(vm, arm).then_some(false);
506 };
507 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
508 return None;
509 }
510 let id = CURRENT.with(|c| c.get());
511 let w = current_waker();
512 let flag = w.flag.clone();
513 if !arm(w.into_wake()) {
514 return Some(false);
515 }
516 with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
517 if in_fiber() {
518 park();
519 } else {
520 run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
521 }
522 Some(Instant::now() >= deadline)
523}
524
525pub fn sleep(vm: &mut Vm, d: Duration) {
527 let deadline = Instant::now() + d;
528 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
529 crate::blocking(vm, || std::thread::sleep(d));
530 return;
531 }
532 let id = CURRENT.with(|c| c.get());
533 while Instant::now() < deadline {
534 with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
535 if in_fiber() {
536 park();
537 } else {
538 let flag = with_sched(|s| s.root_flag.clone());
539 run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
540 }
541 }
542}
543
544pub fn yield_now(vm: &mut Vm) {
546 if in_fiber() {
547 if !vm.park_allowed() {
548 return;
549 }
550 let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("running fiber flag");
552 flag.store(true, Ordering::Release);
553 let y = YIELDER.with(|c| c.get());
554 unsafe { suspend(y) };
555 flag.store(false, Ordering::Release);
556 } else if has_fibers() {
557 let n = with_sched(|s| s.ready.len());
558 for _ in 0..n {
559 absorb();
560 match with_sched(|s| s.dequeue()) {
561 Some(id) => resume(id),
562 None => break,
563 }
564 }
565 }
566}
567
568pub fn join(vm: &mut Vm, handle: Value) -> RtResult<Result<Value, String>> {
573 let id = unsafe { catallaxy_runtime::object::as_instance(handle) }.get(ID_FIELD).as_small_int() as u64;
574 let Some(ctl) = with_sched(|s| s.ctl.get(&id).cloned()) else {
575 return Err(Trap::Panic(
576 "FiberError: no such fiber on this thread (fibers are joined on the thread that spawned them)".into(),
577 ));
578 };
579 loop {
580 if !matches!(*ctl.outcome.lock(), Outcome::Running) {
581 break;
582 }
583 let ctl2 = ctl.clone();
584 let cooperative = wait_event(vm, move |wake| {
585 let o = ctl2.outcome.lock();
586 if !matches!(*o, Outcome::Running) {
587 return false;
588 }
589 ctl2.joiners.lock().push(wake);
590 true
591 });
592 if !cooperative {
593 return Err(Trap::Panic("FiberError: join while holding a lock would deadlock".into()));
596 }
597 }
598 with_sched(|s| s.ctl.remove(&id));
599 let outcome = std::mem::replace(&mut *ctl.outcome.lock(), Outcome::Joined);
600 Ok(match outcome {
601 Outcome::Done { in_handle: true } => {
602 Ok(unsafe { catallaxy_runtime::object::as_instance(handle) }.get(RESULT_FIELD))
603 }
604 Outcome::Done { in_handle: false } => Ok(vm.unpin(id).unwrap_or(Value::NIL)),
605 Outcome::Panicked(m) => Err(m),
606 Outcome::Joined => Err("fiber was already joined".into()),
607 Outcome::Running => unreachable!(),
608 })
609}
610
611pub fn is_done(id: u64) -> bool {
613 with_sched(|s| s.ctl.get(&id).is_none_or(|c| !matches!(*c.outcome.lock(), Outcome::Running)))
614}
615
616type Job = Box<dyn FnOnce() + Send + 'static>;
619
620struct IoPool {
623 queue: Mutex<VecDeque<Job>>,
624 cv: Condvar,
625}
626
627static IO: LazyLock<IoPool> = LazyLock::new(|| {
628 let n = std::env::var("CATX_IO_THREADS").ok().and_then(|v| v.parse::<usize>().ok()).unwrap_or(16).max(1);
629 for i in 0..n {
630 let _ = std::thread::Builder::new().name(format!("catx-io-{i}")).spawn(|| loop {
631 let job = {
632 let mut q = IO.queue.lock();
633 loop {
634 if let Some(j) = q.pop_front() {
635 break j;
636 }
637 IO.cv.wait(&mut q);
638 }
639 };
640 job();
641 });
642 }
643 IoPool { queue: Mutex::new(VecDeque::new()), cv: Condvar::new() }
644});
645
646pub fn blocking_io<R: Send>(vm: &mut Vm, f: impl FnOnce() -> R + Send) -> R {
650 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
651 return crate::blocking(vm, f);
652 }
653 let done = Arc::new(AtomicBool::new(false));
654 let mut result: Option<std::thread::Result<R>> = None;
655 let w = current_waker();
656 {
657 let slot = &mut result as *mut Option<std::thread::Result<R>> as usize;
658 let (done2, w2) = (done.clone(), w.clone());
659 let job: Box<dyn FnOnce() + Send + '_> = Box::new(move || {
660 let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
661 unsafe { *(slot as *mut Option<std::thread::Result<R>>) = Some(r) };
663 done2.store(true, Ordering::Release);
664 w2.wake();
665 });
666 let job: Job = unsafe { std::mem::transmute::<Box<dyn FnOnce() + Send + '_>, Job>(job) };
669 IO.queue.lock().push_back(job);
670 IO.cv.notify_one();
671 }
672 while !done.load(Ordering::Acquire) {
673 if in_fiber() {
674 park();
675 } else {
676 let flag = w.flag.clone();
677 let d = done.clone();
678 run_until(vm, || flag.swap(false, Ordering::AcqRel) || d.load(Ordering::Acquire));
679 }
680 }
681 match result.take().expect("I/O job result") {
682 Ok(r) => r,
683 Err(p) => std::panic::resume_unwind(p),
684 }
685}