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 stack = match DefaultStack::new(FIBER_NATIVE_STACK) {
295 Ok(s) => s,
296 Err(e) => {
297 catallaxy_gc::limits::release_headroom();
298 return Err(Trap::Panic(format!(
299 "MemoryError: cannot allocate a {} KB fiber stack ({e})",
300 FIBER_NATIVE_STACK / 1024
301 )));
302 }
303 };
304 let fvm = Box::into_raw(vm.new_fiber(FIBER_REGISTERS)?);
305 unsafe {
306 (*fvm).extra_roots().push(closure);
307 (*fvm).extra_roots().push(handle);
308 }
309 let ctl = Arc::new(Ctl { outcome: Mutex::new(Outcome::Running), joiners: Mutex::new(Vec::new()) });
310 let ctl2 = ctl.clone();
311 let vm_ptr = fvm as usize;
312 let coro = Coroutine::with_stack(stack, move |y: &Yielder<(), ()>, ()| {
313 let _leave = LeaveOnExit;
314 YIELDER.with(|c| c.set(y as *const _));
315 let vm = unsafe { &mut *(vm_ptr as *mut Vm) };
317 let closure = vm.extra_roots()[0];
318 let r = vm.call_value(closure, &[]);
319 let outcome = match r {
320 Ok(v) => {
321 let handle = vm.extra_roots()[1];
324 let h = unsafe { &*handle.as_ptr() };
325 let mine = h.state() == catallaxy_gc::ObjState::Owned && h.owner() == vm.tid();
326 if mine && catallaxy_runtime::ownership::check_write(&vm.m, h).is_ok() {
327 unsafe { catallaxy_runtime::object::as_instance(handle) }.set(RESULT_FIELD, v);
328 Outcome::Done { in_handle: true }
329 } else {
330 vm.pin(id, v);
331 Outcome::Done { in_handle: false }
332 }
333 }
334 Err(Trap::Panic(msg)) => {
335 let mut full = msg;
336 for line in vm.panic_trace.take().unwrap_or_default().iter().take(8) {
337 full.push_str(&format!("\n at {line}"));
338 }
339 Outcome::Panicked(full)
340 }
341 Err(Trap::Exit(code)) => {
342 crate::flush();
343 std::process::exit(code);
344 }
345 Err(Trap::Unwind { .. }) => Outcome::Panicked("break or return out of a fiber block".into()),
346 };
347 *ctl2.outcome.lock() = outcome;
348 for w in std::mem::take(&mut *ctl2.joiners.lock()) {
349 w();
350 }
351 });
352 with_sched(|s| {
353 s.fibers.insert(
354 id,
355 Fiber {
356 coro,
357 vm: fvm,
358 flag: Arc::new(AtomicBool::new(false)),
359 yielder: Cell::new(std::ptr::null()),
360 tsan: tsan::Handle::new(),
361 },
362 );
363 s.ctl.insert(id, ctl);
364 s.enqueue(id);
365 });
366 Ok(id)
367}
368
369fn resume(id: u64) {
371 let Some(fiber) = with_sched(|s| s.fibers.remove(&id)) else { return };
372 let prev_y = YIELDER.with(|c| c.replace(fiber.yielder.get()));
373 CURRENT.with(|c| c.set(id));
374 RUNNING_FLAG.with(|f| *f.borrow_mut() = Some(fiber.flag.clone()));
375 let mut fiber = fiber;
376 tsan::enter(&fiber.tsan);
377 let r = fiber.coro.resume(());
378 fiber.yielder.set(YIELDER.with(|c| c.get()));
379 RUNNING_FLAG.with(|f| *f.borrow_mut() = None);
380 CURRENT.with(|c| c.set(ROOT));
381 YIELDER.with(|c| c.set(prev_y));
382 match r {
383 CoroutineResult::Yield(()) => {
384 let woken = fiber.flag.load(Ordering::Acquire);
385 with_sched(|s| {
386 s.fibers.insert(id, fiber);
387 if woken {
388 s.enqueue(id);
389 }
390 });
391 }
392 CoroutineResult::Return(()) => {
393 let vm = fiber.vm;
394 drop(fiber);
395 Vm::release_fiber(unsafe { Box::from_raw(vm) });
397 }
398 }
399}
400
401fn absorb() {
403 let now = Instant::now();
404 with_sched(|s| {
405 let woken = std::mem::take(&mut *s.shared.woken.lock());
406 for id in woken {
407 if id != ROOT && s.fibers.contains_key(&id) {
408 s.enqueue(id);
409 }
410 }
411 while let Some(std::cmp::Reverse((t, id))) = s.timers.peek().copied() {
412 if t > now {
413 break;
414 }
415 s.timers.pop();
416 if id == ROOT {
417 s.root_flag.store(true, Ordering::Release);
418 } else if let Some(f) = s.fibers.get(&id) {
419 f.flag.store(true, Ordering::Release);
420 s.enqueue(id);
421 }
422 }
423 });
424}
425
426fn run_until(vm: &mut Vm, mut done: impl FnMut() -> bool) {
430 debug_assert!(!in_fiber());
431 loop {
432 absorb();
433 if done() {
434 return;
435 }
436 if let Some(id) = with_sched(|s| s.dequeue()) {
437 resume(id);
438 catallaxy_gc::safepoint(&vm.m.rec);
439 continue;
440 }
441 let (shared, deadline) = with_sched(|s| (s.shared.clone(), s.timers.peek().map(|r| r.0 .0)));
442 catallaxy_gc::blocking(&vm.m.rec, || {
443 let mut w = shared.woken.lock();
444 if w.is_empty() {
445 match deadline {
446 Some(t) => {
447 let now = Instant::now();
448 if t > now {
449 shared.cv.wait_for(&mut w, t - now);
450 }
451 }
452 None => {
453 shared.cv.wait_for(&mut w, Duration::from_millis(100));
454 }
455 }
456 }
457 });
458 }
459}
460
461fn park() {
463 let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("park outside a fiber");
464 if flag.swap(false, Ordering::AcqRel) {
465 return;
466 }
467 let y = YIELDER.with(|c| c.get());
468 unsafe { suspend(y) };
470 flag.store(false, Ordering::Release);
471}
472
473pub fn drain(vm: &mut Vm) {
475 if in_fiber() || !has_fibers() {
476 return;
477 }
478 run_until(vm, || !has_fibers());
479}
480
481pub fn wait_event(vm: &mut Vm, arm: impl FnOnce(Wake) -> bool) -> bool {
489 if !vm.park_allowed() {
490 return false;
491 }
492 if in_fiber() {
493 let w = current_waker();
494 if arm(w.into_wake()) {
495 park();
496 }
497 return true;
498 }
499 if !has_fibers() {
500 return false;
501 }
502 let w = current_waker();
503 let flag = w.flag.clone();
504 if arm(w.into_wake()) {
505 run_until(vm, || flag.swap(false, Ordering::AcqRel));
506 }
507 true
508}
509
510pub fn wait_event_until(vm: &mut Vm, deadline: Option<Instant>, arm: impl FnOnce(Wake) -> bool) -> Option<bool> {
514 let Some(deadline) = deadline else {
515 return wait_event(vm, arm).then_some(false);
516 };
517 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
518 return None;
519 }
520 let id = CURRENT.with(|c| c.get());
521 let w = current_waker();
522 let flag = w.flag.clone();
523 if !arm(w.into_wake()) {
524 return Some(false);
525 }
526 with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
527 if in_fiber() {
528 park();
529 } else {
530 run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
531 }
532 Some(Instant::now() >= deadline)
533}
534
535pub fn sleep(vm: &mut Vm, d: Duration) {
537 let deadline = Instant::now() + d;
538 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
539 crate::blocking(vm, || std::thread::sleep(d));
540 return;
541 }
542 let id = CURRENT.with(|c| c.get());
543 while Instant::now() < deadline {
544 with_sched(|s| s.timers.push(std::cmp::Reverse((deadline, id))));
545 if in_fiber() {
546 park();
547 } else {
548 let flag = with_sched(|s| s.root_flag.clone());
549 run_until(vm, || flag.swap(false, Ordering::AcqRel) || Instant::now() >= deadline);
550 }
551 }
552}
553
554pub fn yield_now(vm: &mut Vm) {
556 if in_fiber() {
557 if !vm.park_allowed() {
558 return;
559 }
560 let flag = RUNNING_FLAG.with(|f| f.borrow().clone()).expect("running fiber flag");
562 flag.store(true, Ordering::Release);
563 let y = YIELDER.with(|c| c.get());
564 unsafe { suspend(y) };
565 flag.store(false, Ordering::Release);
566 } else if has_fibers() {
567 let n = with_sched(|s| s.ready.len());
568 for _ in 0..n {
569 absorb();
570 match with_sched(|s| s.dequeue()) {
571 Some(id) => resume(id),
572 None => break,
573 }
574 }
575 }
576}
577
578pub fn join(vm: &mut Vm, handle: Value) -> RtResult<Result<Value, String>> {
583 let id = unsafe { catallaxy_runtime::object::as_instance(handle) }.get(ID_FIELD).as_small_int() as u64;
584 let Some(ctl) = with_sched(|s| s.ctl.get(&id).cloned()) else {
585 return Err(Trap::Panic(
586 "FiberError: no such fiber on this thread (fibers are joined on the thread that spawned them)".into(),
587 ));
588 };
589 loop {
590 if !matches!(*ctl.outcome.lock(), Outcome::Running) {
591 break;
592 }
593 let ctl2 = ctl.clone();
594 let cooperative = wait_event(vm, move |wake| {
595 let o = ctl2.outcome.lock();
596 if !matches!(*o, Outcome::Running) {
597 return false;
598 }
599 ctl2.joiners.lock().push(wake);
600 true
601 });
602 if !cooperative {
603 return Err(Trap::Panic("FiberError: join while holding a lock would deadlock".into()));
606 }
607 }
608 with_sched(|s| s.ctl.remove(&id));
609 let outcome = std::mem::replace(&mut *ctl.outcome.lock(), Outcome::Joined);
610 Ok(match outcome {
611 Outcome::Done { in_handle: true } => {
612 Ok(unsafe { catallaxy_runtime::object::as_instance(handle) }.get(RESULT_FIELD))
613 }
614 Outcome::Done { in_handle: false } => Ok(vm.unpin(id).unwrap_or(Value::NIL)),
615 Outcome::Panicked(m) => Err(m),
616 Outcome::Joined => Err("fiber was already joined".into()),
617 Outcome::Running => unreachable!(),
618 })
619}
620
621pub fn is_done(id: u64) -> bool {
623 with_sched(|s| s.ctl.get(&id).is_none_or(|c| !matches!(*c.outcome.lock(), Outcome::Running)))
624}
625
626type Job = Box<dyn FnOnce() + Send + 'static>;
629
630struct IoPool {
633 queue: Mutex<VecDeque<Job>>,
634 cv: Condvar,
635}
636
637static IO: LazyLock<IoPool> = LazyLock::new(|| {
638 let n = std::env::var("CATX_IO_THREADS").ok().and_then(|v| v.parse::<usize>().ok()).unwrap_or(16).max(1);
639 for i in 0..n {
640 let _ = std::thread::Builder::new().name(format!("catx-io-{i}")).spawn(|| loop {
641 let job = {
642 let mut q = IO.queue.lock();
643 loop {
644 if let Some(j) = q.pop_front() {
645 break j;
646 }
647 IO.cv.wait(&mut q);
648 }
649 };
650 job();
651 });
652 }
653 IoPool { queue: Mutex::new(VecDeque::new()), cv: Condvar::new() }
654});
655
656pub fn blocking_io<R: Send>(vm: &mut Vm, f: impl FnOnce() -> R + Send) -> R {
660 if !vm.park_allowed() || (!in_fiber() && !has_fibers()) {
661 return crate::blocking(vm, f);
662 }
663 let done = Arc::new(AtomicBool::new(false));
664 let mut result: Option<std::thread::Result<R>> = None;
665 let w = current_waker();
666 {
667 let slot = &mut result as *mut Option<std::thread::Result<R>> as usize;
668 let (done2, w2) = (done.clone(), w.clone());
669 let job: Box<dyn FnOnce() + Send + '_> = Box::new(move || {
670 let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
671 unsafe { *(slot as *mut Option<std::thread::Result<R>>) = Some(r) };
673 done2.store(true, Ordering::Release);
674 w2.wake();
675 });
676 let job: Job = unsafe { std::mem::transmute::<Box<dyn FnOnce() + Send + '_>, Job>(job) };
679 IO.queue.lock().push_back(job);
680 IO.cv.notify_one();
681 }
682 while !done.load(Ordering::Acquire) {
683 if in_fiber() {
684 park();
685 } else {
686 let flag = w.flag.clone();
687 let d = done.clone();
688 run_until(vm, || flag.swap(false, Ordering::AcqRel) || d.load(Ordering::Acquire));
689 }
690 }
691 match result.take().expect("I/O job result") {
692 Ok(r) => r,
693 Err(p) => std::panic::resume_unwind(p),
694 }
695}