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, Stack};
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 stack_limit: usize,
192 tsan: tsan::Handle,
194}
195
196impl Drop for Fiber {
197 fn drop(&mut self) {
198 if self.coro.started() && !self.coro.done() {
201 tsan::enter(&self.tsan);
202 }
203 }
204}
205
206struct Scheduler {
207 shared: Arc<Shared>,
208 root_flag: Arc<AtomicBool>,
209 fibers: HashMap<u64, Fiber>,
211 ready: VecDeque<u64>,
212 queued: std::collections::HashSet<u64>,
214 timers: BinaryHeap<std::cmp::Reverse<(Instant, u64)>>,
216 ctl: HashMap<u64, Arc<Ctl>>,
217}
218
219impl Scheduler {
220 fn enqueue(&mut self, id: u64) {
221 if self.queued.insert(id) {
222 self.ready.push_back(id);
223 }
224 }
225
226 fn dequeue(&mut self) -> Option<u64> {
227 let id = self.ready.pop_front()?;
228 self.queued.remove(&id);
229 Some(id)
230 }
231
232 fn new() -> Self {
233 Scheduler {
234 shared: Arc::new(Shared { woken: Mutex::new(Vec::new()), cv: Condvar::new() }),
235 root_flag: Arc::new(AtomicBool::new(false)),
236 fibers: HashMap::new(),
237 ready: VecDeque::new(),
238 queued: std::collections::HashSet::new(),
239 timers: BinaryHeap::new(),
240 ctl: HashMap::new(),
241 }
242 }
243}
244
245thread_local! {
246 static SCHED: RefCell<Option<Scheduler>> = const { RefCell::new(None) };
247 static CURRENT: Cell<u64> = const { Cell::new(ROOT) };
249 static YIELDER: Cell<*const Yielder<(), ()>> = const { Cell::new(std::ptr::null()) };
250}
251
252fn with_sched<R>(f: impl FnOnce(&mut Scheduler) -> R) -> R {
253 SCHED.with(|s| f(s.borrow_mut().get_or_insert_with(Scheduler::new)))
254}
255
256pub fn in_fiber() -> bool {
258 CURRENT.with(|c| c.get()) != ROOT
259}
260
261pub fn has_fibers() -> bool {
263 SCHED.with(|s| s.borrow().as_ref().is_some_and(|s| !s.fibers.is_empty()))
264}
265
266fn current_waker() -> Waker {
268 let id = CURRENT.with(|c| c.get());
269 let flag = if id == ROOT {
270 with_sched(|s| s.root_flag.clone())
271 } else {
272 RUNNING_FLAG.with(|f| f.borrow().clone()).expect("running fiber flag")
273 };
274 Waker { shared: with_sched(|s| s.shared.clone()), id, flag }
275}
276
277thread_local! {
278 static RUNNING_FLAG: RefCell<Option<Arc<AtomicBool>>> = const { RefCell::new(None) };
280}
281
282const ID_FIELD: usize = 0;
286const RESULT_FIELD: usize = 1;
287
288fn check_handle(handle: Value) -> RtResult<()> {
291 let id = catallaxy_vm::program::class_named("Fiber").unwrap_or(u32::MAX);
292 catallaxy_runtime::object::instance_of(handle, id, "Fiber").map(|_| ())
293}
294
295pub fn spawn(vm: &mut Vm, handle: Value, closure: Value) -> RtResult<u64> {
298 check_handle(handle)?;
299 let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
300 let h = unsafe { &*handle.as_ptr() };
302 catallaxy_runtime::ownership::check_write(&vm.m, h)?;
303 unsafe { catallaxy_runtime::object::as_instance(handle) }.set(ID_FIELD, Value::small_int(id as i64));
304 let stack = match DefaultStack::new(FIBER_NATIVE_STACK) {
306 Ok(s) => s,
307 Err(e) => {
308 catallaxy_gc::limits::release_headroom();
309 return Err(Trap::Panic(format!(
310 "MemoryError: cannot allocate a {} KB fiber stack ({e})",
311 FIBER_NATIVE_STACK / 1024
312 )));
313 }
314 };
315 let stack_limit = catallaxy_runtime::stack::limit_for(stack.limit().get(), stack.base().get());
317 let fvm = Box::into_raw(vm.new_fiber(FIBER_REGISTERS)?);
318 unsafe {
319 (*fvm).extra_roots().push(closure);
320 (*fvm).extra_roots().push(handle);
321 }
322 let ctl = Arc::new(Ctl { outcome: Mutex::new(Outcome::Running), joiners: Mutex::new(Vec::new()) });
323 let ctl2 = ctl.clone();
324 let vm_ptr = fvm as usize;
325 let coro = Coroutine::with_stack(stack, move |y: &Yielder<(), ()>, ()| {
326 let _leave = LeaveOnExit;
327 YIELDER.with(|c| c.set(y as *const _));
328 let vm = unsafe { &mut *(vm_ptr as *mut Vm) };
330 let closure = vm.extra_roots()[0];
331 let r = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| vm.call_value(closure, &[]))) {
334 Ok(r) => r,
335 Err(payload) => {
336 let Some(msg) = crate::internal_error_message(&*payload) else { std::panic::resume_unwind(payload) };
338 let trace = vm.recover_from_internal_panic();
339 vm.panic_trace = Some(trace);
340 Err(Trap::Panic(msg))
341 }
342 };
343 let outcome = match r {
344 Ok(v) => {
345 let handle = vm.extra_roots()[1];
348 let h = unsafe { &*handle.as_ptr() };
349 let mine = h.state() == catallaxy_gc::ObjState::Owned && h.owner() == vm.tid();
350 if mine && catallaxy_runtime::ownership::check_write(&vm.m, h).is_ok() {
351 unsafe { catallaxy_runtime::object::as_instance(handle) }.set(RESULT_FIELD, v);
352 Outcome::Done { in_handle: true }
353 } else {
354 vm.pin(id, v);
355 Outcome::Done { in_handle: false }
356 }
357 }
358 Err(Trap::Panic(msg)) => {
359 let mut full = msg;
360 for line in vm.panic_trace.take().unwrap_or_default().iter().take(8) {
361 full.push_str(&format!("\n at {line}"));
362 }
363 Outcome::Panicked(full)
364 }
365 Err(Trap::Exit(code)) => {
366 crate::flush();
367 std::process::exit(code);
368 }
369 Err(Trap::Unwind { .. }) => Outcome::Panicked("break or return out of a fiber block".into()),
370 };
371 *ctl2.outcome.lock() = outcome;
372 for w in std::mem::take(&mut *ctl2.joiners.lock()) {
373 w();
374 }
375 });
376 with_sched(|s| {
377 s.fibers.insert(
378 id,
379 Fiber {
380 coro,
381 vm: fvm,
382 flag: Arc::new(AtomicBool::new(false)),
383 yielder: Cell::new(std::ptr::null()),
384 stack_limit,
385 tsan: tsan::Handle::new(),
386 },
387 );
388 s.ctl.insert(id, ctl);
389 s.enqueue(id);
390 });
391 Ok(id)
392}
393
394fn resume(id: u64) {
396 let Some(fiber) = with_sched(|s| s.fibers.remove(&id)) else { return };
397 let prev_y = YIELDER.with(|c| c.replace(fiber.yielder.get()));
398 CURRENT.with(|c| c.set(id));
399 RUNNING_FLAG.with(|f| *f.borrow_mut() = Some(fiber.flag.clone()));
400 let mut fiber = fiber;
401 let prev_limit = catallaxy_runtime::stack::set_limit(fiber.stack_limit);
402 tsan::enter(&fiber.tsan);
403 let r = fiber.coro.resume(());
404 catallaxy_runtime::stack::set_limit(prev_limit);
405 fiber.yielder.set(YIELDER.with(|c| c.get()));
406 RUNNING_FLAG.with(|f| *f.borrow_mut() = None);
407 CURRENT.with(|c| c.set(ROOT));
408 YIELDER.with(|c| c.set(prev_y));
409 match r {
410 CoroutineResult::Yield(()) => {
411 let woken = fiber.flag.load(Ordering::Acquire);
412 with_sched(|s| {
413 s.fibers.insert(id, fiber);
414 if woken {
415 s.enqueue(id);
416 }
417 });
418 }
419 CoroutineResult::Return(()) => {
420 let vm = fiber.vm;
421 drop(fiber);
422 Vm::release_fiber(unsafe { Box::from_raw(vm) });
424 }
425 }
426}
427
428fn absorb() {
430 let now = Instant::now();
431 with_sched(|s| {
432 let woken = std::mem::take(&mut *s.shared.woken.lock());
433 for id in woken {
434 if id != ROOT && s.fibers.contains_key(&id) {
435 s.enqueue(id);
436 }
437 }
438 while let Some(std::cmp::Reverse((t, id))) = s.timers.peek().copied() {
439 if t > now {
440 break;
441 }
442 s.timers.pop();
443 if id == ROOT {
444 s.root_flag.store(true, Ordering::Release);
445 } else if let Some(f) = s.fibers.get(&id) {
446 f.flag.store(true, Ordering::Release);
447 s.enqueue(id);
448 }
449 }
450 });
451}
452
453fn run_until(vm: &mut Vm, mut done: impl FnMut() -> bool) {
457 debug_assert!(!in_fiber());
458 loop {
459 absorb();
460 if done() {
461 return;
462 }
463 if let Some(id) = with_sched(|s| s.dequeue()) {
464 resume(id);
465 catallaxy_gc::safepoint(&vm.m.rec);
466 continue;
467 }
468 let (shared, deadline) = with_sched(|s| (s.shared.clone(), s.timers.peek().map(|r| r.0 .0)));
469 catallaxy_gc::blocking(&vm.m.rec, || {
470 let mut w = shared.woken.lock();
471 if w.is_empty() {
472 match deadline {
473 Some(t) => {
474 let now = Instant::now();
475 if t > now {
476 shared.cv.wait_for(&mut w, t - now);
477 }
478 }
479 None => {
480 shared.cv.wait_for(&mut w, Duration::from_millis(100));
481 }
482 }
483 }
484 });
485 }
486}
487
488fn park() {
490 let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("park outside a fiber");
491 if flag.swap(false, Ordering::AcqRel) {
492 return;
493 }
494 let y = YIELDER.with(|c| c.get());
495 unsafe { suspend(y) };
497 flag.store(false, Ordering::Release);
498}
499
500pub fn abandon() {
503 CURRENT.with(|c| c.set(ROOT));
504 catallaxy_runtime::stack::init_thread();
506 RUNNING_FLAG.with(|f| *f.borrow_mut() = None);
507 YIELDER.with(|c| c.set(std::ptr::null()));
508 let sched = SCHED.with(|s| s.borrow_mut().take());
510 drop(sched);
511}
512
513pub fn drain(vm: &mut Vm) {
515 if in_fiber() || !has_fibers() {
516 return;
517 }
518 run_until(vm, || !has_fibers());
519}
520
521pub fn wait_event(vm: &mut Vm, arm: impl FnOnce(Wake) -> bool) -> bool {
529 if !vm.park_allowed() {
530 return false;
531 }
532 if in_fiber() {
533 let w = current_waker();
534 if arm(w.into_wake()) {
535 park();
536 }
537 return true;
538 }
539 if !has_fibers() {
540 return false;
541 }
542 let w = current_waker();
543 let flag = w.flag.clone();
544 if arm(w.into_wake()) {
545 run_until(vm, || flag.swap(false, Ordering::AcqRel));
546 }
547 true
548}
549
550pub fn wait_event_until(vm: &mut Vm, deadline: Option<Instant>, arm: impl FnOnce(Wake) -> bool) -> Option<bool> {
554 let Some(deadline) = deadline else {
555 return wait_event(vm, arm).then_some(false);
556 };
557 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
558 return None;
559 }
560 let id = CURRENT.with(|c| c.get());
561 let w = current_waker();
562 let flag = w.flag.clone();
563 if !arm(w.into_wake()) {
564 return Some(false);
565 }
566 with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
567 if in_fiber() {
568 park();
569 } else {
570 run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
571 }
572 Some(Instant::now() >= deadline)
573}
574
575pub fn sleep(vm: &mut Vm, d: Duration) {
577 let deadline = Instant::now() + d;
578 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
579 crate::blocking(vm, || std::thread::sleep(d));
580 return;
581 }
582 let id = CURRENT.with(|c| c.get());
583 while Instant::now() < deadline {
584 with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
585 if in_fiber() {
586 park();
587 } else {
588 let flag = with_sched(|s| s.root_flag.clone());
589 run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
590 }
591 }
592}
593
594pub fn yield_now(vm: &mut Vm) {
596 if in_fiber() {
597 if !vm.park_allowed() {
598 return;
599 }
600 let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("running fiber flag");
602 flag.store(true, Ordering::Release);
603 let y = YIELDER.with(|c| c.get());
604 unsafe { suspend(y) };
605 flag.store(false, Ordering::Release);
606 } else if has_fibers() {
607 let n = with_sched(|s| s.ready.len());
608 for _ in 0..n {
609 absorb();
610 match with_sched(|s| s.dequeue()) {
611 Some(id) => resume(id),
612 None => break,
613 }
614 }
615 }
616}
617
618pub fn join(vm: &mut Vm, handle: Value) -> RtResult<Result<Value, String>> {
623 check_handle(handle)?;
624 let id = unsafe { catallaxy_runtime::object::as_instance(handle) }.get(ID_FIELD).as_small_int() as u64;
626 let Some(ctl) = with_sched(|s| s.ctl.get(&id).cloned()) else {
627 return Err(Trap::Panic(
628 "FiberError: no such fiber on this thread (fibers are joined on the thread that spawned them)".into(),
629 ));
630 };
631 loop {
632 if !matches!(*ctl.outcome.lock(), Outcome::Running) {
633 break;
634 }
635 let ctl2 = ctl.clone();
636 let cooperative = wait_event(vm, move |wake| {
637 let o = ctl2.outcome.lock();
638 if !matches!(*o, Outcome::Running) {
639 return false;
640 }
641 ctl2.joiners.lock().push(wake);
642 true
643 });
644 if !cooperative {
645 return Err(Trap::Panic("FiberError: join while holding a lock would deadlock".into()));
648 }
649 }
650 with_sched(|s| s.ctl.remove(&id));
651 let outcome = std::mem::replace(&mut *ctl.outcome.lock(), Outcome::Joined);
652 Ok(match outcome {
653 Outcome::Done { in_handle: true } => {
654 Ok(unsafe { catallaxy_runtime::object::as_instance(handle) }.get(RESULT_FIELD))
655 }
656 Outcome::Done { in_handle: false } => Ok(vm.unpin(id).unwrap_or(Value::NIL)),
657 Outcome::Panicked(m) => Err(m),
658 Outcome::Joined => Err("fiber was already joined".into()),
659 Outcome::Running => unreachable!(),
660 })
661}
662
663pub fn is_done(id: u64) -> bool {
665 with_sched(|s| s.ctl.get(&id).is_none_or(|c| !matches!(*c.outcome.lock(), Outcome::Running)))
666}
667
668type Job = Box<dyn FnOnce() + Send + 'static>;
671
672struct IoPool {
675 queue: Mutex<VecDeque<Job>>,
676 cv: Condvar,
677}
678
679static IO: LazyLock<IoPool> = LazyLock::new(|| {
680 let n = std::env::var("CATX_IO_THREADS").ok().and_then(|v| v.parse::<usize>().ok()).unwrap_or(16).max(1);
681 for i in 0..n {
682 let _ = std::thread::Builder::new().name(format!("catx-io-{i}")).spawn(|| loop {
683 let job = {
684 let mut q = IO.queue.lock();
685 loop {
686 if let Some(j) = q.pop_front() {
687 break j;
688 }
689 IO.cv.wait(&mut q);
690 }
691 };
692 job();
693 });
694 }
695 IoPool { queue: Mutex::new(VecDeque::new()), cv: Condvar::new() }
696});
697
698pub fn blocking_io<R: Send>(vm: &mut Vm, f: impl FnOnce() -> R + Send) -> R {
702 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
703 return crate::blocking(vm, f);
704 }
705 let done = Arc::new(AtomicBool::new(false));
706 let mut result: Option<std::thread::Result<R>> = None;
707 let w = current_waker();
708 {
709 let slot = &mut result as *mut Option<std::thread::Result<R>> as usize;
710 let (done2, w2) = (done.clone(), w.clone());
711 let job: Box<dyn FnOnce() + Send + '_> = Box::new(move || {
712 let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
713 unsafe { *(slot as *mut Option<std::thread::Result<R>>) = Some(r) };
715 done2.store(true, Ordering::Release);
716 w2.wake();
717 });
718 let job: Job = unsafe { std::mem::transmute::<Box<dyn FnOnce() + Send + '_>, Job>(job) };
721 IO.queue.lock().push_back(job);
722 IO.cv.notify_one();
723 }
724 while !done.load(Ordering::Acquire) {
725 if in_fiber() {
726 park();
727 } else {
728 let flag = w.flag.clone();
729 let d = done.clone();
730 run_until(vm, || flag.swap(false, Ordering::AcqRel) || d.load(Ordering::Acquire));
731 }
732 }
733 match result.take().expect("I/O job result") {
734 Ok(r) => r,
735 Err(p) => std::panic::resume_unwind(p),
736 }
737}