Rust implementation of fibers

Question

I am still new to rust, so any comment about what is the rust way of doing things is welcome. this is a hand-rolled implementation of fibers or user-land context-switching.

#![feature(allocator_api)]
#![feature(lazy_cell)]

use std::{
  alloc::{Allocator, Global, Layout},
  arch::global_asm,
  sync::{LazyLock, Mutex},
};

use std::cell::RefCell;
use std::thread::LocalKey;
use std::ptr::NonNull;

/// This might be not enough stack. so allocate more if we stack overflow
const STACK_SIZE: usize = 65536;
const STACK_ALIGN: usize = 16;

type FiberStarterFunc = fn(*mut u8, *mut u8, usize) -> ();

extern "sysv64" {
  fn switch_to_other_stack(new_stack: *mut u8) -> *mut u8;
  fn start_new_stack(new_stack: *mut u8, f_ptr: FiberStarterFunc, new_id: usize, data: *mut u8) -> *mut u8;
}

// f_ptr is function pointer with a rust call convention. so it cannot be called directly from inline assembly. so we use this jump-pad to call it.
extern "sysv64" fn jump_pad(old_stack: *mut u8, f_ptr: FiberStarterFunc, new_id: usize, data: *mut u8) -> *mut u8 {
  f_ptr(data, old_stack, new_id);
  std::unreachable!("this function should never exit");
}

// all context-switches are done inside a functions called with the sysv64 calling convention.
// so the only registers that need to be preserved are RBX, RSP, RBP, and R12, R13, R14, and R15
// all other registers are preserved by the caller, which is the compiler in this case.
global_asm!("
.global switch_to_other_stack
switch_to_other_stack:
push rbx
push r12
push r13
push r14
push r15
push rbp

mov rax, rsp
mov rsp, rdi

pop rbp
pop r15
pop r14
pop r13
pop r12
pop rbx
ret

.global start_new_stack
start_new_stack:

push rbx
push r12
push r13
push r14
push r15
push rbp

mov rbx, rsp
mov rsp, rdi
mov rdi, rbx
call {}
mov rsp, rax

pop rbp
pop r15
pop r14
pop r13
pop r12
pop rbx
ret
",
sym jump_pad,
);

trait Scheduler {
  fn on_new_id(&mut self, id: usize);
  fn on_yield(&mut self, id: usize) -> usize;
  fn on_exit(&mut self, id: usize) -> usize;
}

struct SchedElem {
  id: usize,
  next: usize,
  prev: usize,
}

struct LRUScheduler {
  elems: Vec<SchedElem>,
  head: usize,
  tail: usize,
}

impl LRUScheduler {
  // TODO all the linked-list logic should be factored out instead its own container.
  fn unlink(&mut self, idx: usize) {
    if self.head == idx {
      self.head = self.elems[idx].next;
    } else {
      let p = self.elems[idx].prev;
      self.elems[p].next = self.elems[idx].next;
    }
    if self.tail == idx {
      self.tail = self.elems[idx].prev;
    } else {
      let n = self.elems[idx].next;
      self.elems[n].prev = self.elems[idx].prev;
    }
    self.elems[idx].next = usize::MAX;
    self.elems[idx].prev = usize::MAX;
  }
  fn insert_head(&mut self, idx: usize) {
    assert_eq!(self.elems[idx].prev, usize::MAX);
    assert_eq!(self.elems[idx].next, usize::MAX);
    if self.head == usize::MAX {
      self.tail = idx;
    } else {
      self.elems[self.head].prev = idx;
    }
    self.elems[idx].next = self.head;
    self.head = idx;
  }
  fn insert_tail(&mut self, idx: usize) {
    assert_eq!(self.elems[idx].prev, usize::MAX);
    assert_eq!(self.elems[idx].next, usize::MAX);
    if self.tail == usize::MAX {
      self.head = idx;
    } else {
      self.elems[self.tail].next = idx;
    }
    self.elems[idx].prev = self.tail;
    self.tail = idx;
  }
}

impl Default for LRUScheduler {
  fn default() -> Self {
    Self {
      elems: Vec::new(),
      head: usize::MAX,
      tail: usize::MAX,
    }
  }
}

impl Scheduler for LRUScheduler {
  fn on_new_id(&mut self, id: usize) {
    if id >= self.elems.len() {
      self.elems.push(SchedElem {
        id,
        next: usize::MAX,
        prev: usize::MAX,
      });
    }
    self.insert_head(id);
  }
  fn on_yield(&mut self, id: usize) -> usize {
    self.unlink(id);
    self.insert_tail(id);
    self.elems[self.head].id
  }
  fn on_exit(&mut self, id: usize) -> usize {
    self.unlink(id);
    self.elems[self.head].id
  }
}

struct FiberContext {
  stacks: Vec<*mut u8>,
  cleanups: Vec<*mut u8>,
  scheduler: Box<dyn Scheduler>,
  unused_ids: Vec<usize>,
}

impl FiberContext {
  fn get() -> &'static mut LazyLock<Mutex<FiberContext>> {
    static mut SINGLETON: LazyLock<Mutex<FiberContext>> = LazyLock::new(|| {
      Mutex::new(FiberContext {
        stacks: Vec::new(),
        cleanups: Vec::new(),
        scheduler: Box::new(LRUScheduler::default()),
        unused_ids: Vec::new(),
      })
    });
    unsafe { &mut SINGLETON }
  }
  fn fiber_id_impl() -> &'static LocalKey<RefCell<usize>> {
    thread_local! {
      static FIBER_ID: RefCell<usize> = RefCell::new(usize::MAX);
    }
    return &FIBER_ID;
  }
  pub fn fiber_id() -> usize {
    Self::fiber_id_impl().with(|fiber_id| *fiber_id.borrow())
  }
  fn set_fiber_id(id: usize) {
    Self::fiber_id_impl().with(|fiber_id| *fiber_id.borrow_mut() = id);
  }

  fn update_fiber_metadata(new_id: usize, old_stack: *mut u8) {
    let mut ctx = Self::get().lock().unwrap();
    ctx.stacks[Self::fiber_id()] = old_stack;
    Self::set_fiber_id(new_id);
  }

  pub fn yield_impl(is_exit: bool) {
    let fiber_id = Self::fiber_id();

    let new_stack = {
      let mut ctx = Self::get().lock().unwrap();
      let selected_id = if is_exit {
        ctx.scheduler.on_exit(fiber_id)
      } else {
        ctx.scheduler.on_yield(fiber_id)
      };

      // if the selected fiber is the current fiber, do nothing. note that the stack address is not valid anymore.
      if selected_id == fiber_id {
        return;
      }
      ctx.stacks[selected_id]
    };
    let old_stack = unsafe { switch_to_other_stack(new_stack) };
    Self::update_fiber_metadata(fiber_id, old_stack);
  }

  pub fn self_yield() {
    Self::yield_impl(false);
  }

    fn get_stack_layout() -> Layout {
    Layout::from_size_align(STACK_SIZE, STACK_ALIGN).expect("failed to create layout for stack")
  }

  fn get_new_id() -> usize {
    let mut ctx = Self::get().lock().unwrap();
    if let Some(id) = ctx.unused_ids.pop() {
      ctx.scheduler.on_new_id(id);
      return id;
    }
    ctx.stacks.push(std::ptr::null_mut());
    let id = ctx.stacks.len() - 1;
    ctx.scheduler.on_new_id(id);
    id
  }

  fn exit_fiber() {
    let id = Self::fiber_id();
    {
      let mut ctx = Self::get().lock().unwrap();
      ctx.stacks[id] = std::ptr::null_mut();
      ctx.unused_ids.push(id);
    };
    Self::yield_impl(true);
  }

  fn create_new_fiber() -> (usize, *mut u8) {
    let ptr: *mut u8 = Global
    .allocate(Self::get_stack_layout())
    .expect("failed to allocate new stack")
    .as_ptr() as *mut u8;
    Self::get().lock().unwrap().cleanups.push(ptr);
    (Self::get_new_id(), ptr)
  }

  fn get_or_create_id() -> usize {
    if Self::fiber_id() != usize::MAX {
      Self::fiber_id()
    } else {
      Self::get_new_id()
    }
  }

  // spawn a new fiber and move execution to it.
  pub fn spawn<T: FnOnce() -> ()>(func: T) {
    unsafe {
      let fiber_id = Self::get_or_create_id();
      let (new_id, ptr) = Self::create_new_fiber();

      fn specialization<T: FnOnce() -> ()>(data: *mut u8, old_stack: *mut u8, new_id: usize) {
        FiberContext::update_fiber_metadata(new_id, old_stack);
        unsafe {
          std::ptr::read(data as *mut T)();
        }
        FiberContext::exit_fiber();
        std::unreachable!("the scheduler should never re-schedule this fiber after a exit_fiber");
      }

      // stack should be aligned on 16 bytes
      let addr = ptr.add(STACK_SIZE - (std::mem::size_of::<T>() + 15 & !15));
      // write the closure at the top of the stack
      std::ptr::write(addr as *mut T, func);

      Self::set_fiber_id(fiber_id);
      let old_stack = start_new_stack(addr, specialization::<T>, new_id, addr);
      Self::update_fiber_metadata(fiber_id, old_stack);
    }
  }
  pub fn set_scheduler(scheduler: Box<dyn Scheduler>) {
    Self::get().lock().unwrap().scheduler = scheduler
  }
}

impl Drop for FiberContext {
  fn drop(&mut self) {
    let mut ctx = Self::get().lock().unwrap();
    for cleanup in self.cleanups.drain(..) {
      unsafe {
        Global.deallocate(
          NonNull::<u8>::new(cleanup).unwrap(),
          Self::get_stack_layout(),
        );
      }
    }
    ctx.cleanups.clear();
  }
}

fn stack_addr() -> *const u8 {
  let i : u8 = 0;
  &i as *const u8
}

fn main() {
  const COUNT: usize = 5;
  for i in 0..COUNT {
    FiberContext::spawn(move || {
      println!("stack({}) 1: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
      FiberContext::self_yield();
      println!("stack({}) 2: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
      FiberContext::self_yield();
      println!("stack({}) 3: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
    });
  }
  println!("main stack 2: {:#?}, fiber_id={}", stack_addr(), FiberContext::fiber_id());
  FiberContext::self_yield();
  println!("main stack 3: {:#?}, fiber_id={}", stack_addr(), FiberContext::fiber_id());
  FiberContext::self_yield();
  for i in COUNT..(COUNT * 2) {
    FiberContext::spawn(move || {
      println!("stack({}) 1: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
      FiberContext::self_yield();
      println!("stack({}) 2: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
      FiberContext::self_yield();
      println!("stack({}) 3: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
    });
  }
  println!("main stack 4: {:#?}, fiber_id={}", stack_addr(), FiberContext::fiber_id());
  FiberContext::self_yield();
  println!("main stack 5: {:#?}, fiber_id={}", stack_addr(), FiberContext::fiber_id());
  for i in COUNT..(COUNT * 2) {
    FiberContext::spawn(move || {
      println!("stack({}) 1: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
      FiberContext::self_yield();
      println!("stack({}) 2: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
      FiberContext::self_yield();
      println!("stack({}) 3: {:#?}, fiber_id={}", i, stack_addr(), FiberContext::fiber_id());
    });
  }
  FiberContext::self_yield();
  println!("main stack 6: {:#?}, fiber_id={}", stack_addr(), FiberContext::fiber_id());
}

I am aware the "linked-list" in the LRU scheduler should be a separate container. I there other things that should be improve ?

kmdreko · Accepted Answer · 2023-09-24 03:31:35Z

I feel unqualified to critique a lot of this, but I will address what I feel I can.

Your FiberContext::get API is unsound:

impl FiberContext {
  fn get() -> &'static mut LazyLock<Mutex<FiberContext>> {
    static mut SINGLETON: LazyLock<Mutex<FiberContext>> = LazyLock::new(|| {
      Mutex::new(FiberContext {
        stacks: Vec::new(),
        cleanups: Vec::new(),
        scheduler: Box::new(LRUScheduler::default()),
        unused_ids: Vec::new(),
      })
    });
    unsafe { &mut SINGLETON }
  }
  ...

This function could get called from multiple threads (or even multiple sites in the same thread). This would allow multiple mutable references to the same object at the same time which is immediate undefined behavior in Rust. In general stay away from static mut.

Fortunately, you don't even need mutable access since you've properly wrapped your context in a Mutex already. So just remove the mut everywhere; it is not needed.

However, doing so un-surfaces more problems; your FiberContext is not actually thread-safe!

error[E0277]: `(dyn Scheduler + 'static)` cannot be sent between threads safely
   --> src/main.rs:199:23
    |
199 |     static SINGLETON: LazyLock<Mutex<FiberContext>> = LazyLock::new(|| {
    |                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `(dyn Scheduler + 'static)` cannot be sent between threads safely
    |
    = help: the trait `Send` is not implemented for `(dyn Scheduler + 'static)`
error[E0277]: `*mut u8` cannot be sent between threads safely
   --> src/main.rs:199:23
    |
199 |     static SINGLETON: LazyLock<Mutex<FiberContext>> = LazyLock::new(|| {
    |                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `*mut u8` cannot be sent between threads safely
    |
    = help: the trait `Send` is not implemented for `*mut u8`

Note: I'll assume you actually want it to be thread-safe since you didn't use a thread-local like you did elsewhere.

To fix the one for dyn Scheduler, you should use dyn Scheduler + Send as the bound or if this is the primary use of the Scheduler trait, you may want to add it as a constraint on the trait itself:

trait Scheduler: Send { // <----------------
  fn on_new_id(&mut self, id: usize);
  fn on_yield(&mut self, id: usize) -> usize;
  fn on_exit(&mut self, id: usize) -> usize;
}

To fix the error for *mut u8, raw pointer types are not assumed to be thread-safe, so you will either have to declare the traits manually for your context:

unsafe impl Send for FiberContext {}
unsafe impl Sync for FiberContext {}

Or it might be better to make wrapper types for your "stack" and "cleanup" types around *mut u8 and just implement Send and Sync for those just so the behavior is well encapsulated.

You are already using the nightly toolchain for the allocator API, but you don't actually need to use LazyLock in this code. You can use OnceLock instead:

use std::sync::OnceLock;

impl FiberContext {
  fn get() -> &'static Mutex<FiberContext> {
    static SINGLETON: OnceLock<Mutex<FiberContext>> = OnceLock::new();
    SINGLETON.get_or_init(|| {
      Mutex::new(FiberContext {
        stacks: Vec::new(),
        cleanups: Vec::new(),
        scheduler: Box::new(LRUScheduler::default()),
        unused_ids: Vec::new(),
      })
    })
  }
  ...
}

It has a slightly different API but in this limited use-case it does the same thing. And even if nightly isn't a big deal for you, I would still make this change since the LazyLock API may end up changing and I wouldn't want experimental code to break yours needlessly.

Next, you should heed the compiler warnings:

warning: `extern` block uses type `fn(*mut u8, *mut u8, usize)`, which is not FFI-safe
  --> src/main.rs:21:49
   |
21 |   fn start_new_stack(new_stack: *mut u8, f_ptr: FiberStarterFunc, new_id: usize, data: *mut u8) -> *mut u8;
   |                                                 ^^^^^^^^^^^^^^^^ not FFI-safe
   |
   = help: consider using an `extern fn(...) -> ...` function pointer instead
   = note: this function pointer has Rust-specific calling convention
   = note: `#[warn(improper_ctypes)]` on by default

warning: `extern` fn uses type `fn(*mut u8, *mut u8, usize)`, which is not FFI-safe
  --> src/main.rs:25:56
   |
25 | extern "sysv64" fn jump_pad(old_stack: *mut u8, f_ptr: FiberStarterFunc, new_id: usize, data: *mut u8) -> *mut u8 {
   |                                                        ^^^^^^^^^^^^^^^^ not FFI-safe
   |
   = help: consider using an `extern fn(...) -> ...` function pointer instead
   = note: this function pointer has Rust-specific calling convention
   = note: `#[warn(improper_ctypes_definitions)]` on by default

I see you have comments that assert this is fine, which I don't have any particular issue with, but we should still address warnings and not let them linger. You could sprinkle in #[allow(improper_ctypes)] and #[allow(improper_ctypes_definitions)] to suppress the warnings but even in the smallest possible scope they are still too over-eager for my liking and may cover-up other problems down the line.

I would suggest using an opaque pointer type (*const ()) at the extern boundary and use casts at the input and output points:

extern "sysv64" {
  ...
  fn start_new_stack(new_stack: *mut u8, f_ptr: *const (), new_id: usize, data: *mut u8) -> *mut u8;
}

extern "sysv64" fn jump_pad(old_stack: *mut u8, f_ptr: *const (), new_id: usize, data: *mut u8) -> *mut u8 {
  let f_ptr: FiberStarterFunc = unsafe { std::mem::transmute(f_ptr) };
  ...
}

let old_stack = start_new_stack(addr, specialization::<T> as FiberStarterFunc as *const (), ...

Your unsafe block in spawn() casts a pretty wide net, whereas it is common practice that unsafe blocks should be as small as they can be to limit the code points that should be audited for undefined behavior.

You should also get in the habit of documenting the unsafe blocks with comments explaining why the operation is actually safe by addressing the unsafe call's safety requirements (though likely not needed on the extern calls).

Thank you. here is the updated code: godbolt.org/z/x9xa9nr7G — Ralender, Commented Sep 24, 2023 at 9:27

Stack Exchange Network

Rust implementation of fibers

1 Answer 1

Your Answer

Not the answer you're looking for? Browse other questions tagged
rust
assembly
or ask your own question.

Hot Network Questions

Rust implementation of fibers

1 Answer 1

Your Answer

Sign up or log in

Post as a guest

Not the answer you're looking for? Browse other questions tagged rustassembly or ask your own question.

Related

Hot Network Questions

Not the answer you're looking for? Browse other questions tagged
rust
assembly
or ask your own question.