Skip to content
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
adds Schedule type
  • Loading branch information
sunny-g authored and filmor committed Feb 13, 2024
commit bb6d59cc9b1a64a62c9bc52278cd82f8509e0d1d
91 changes: 91 additions & 0 deletions rustler/src/schedule.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
use rustler_sys::c_char;

use crate::codegen_runtime::{NifReturnable, NifReturned};
use crate::wrapper::ErlNifTaskFlags;
use crate::Env;

Expand All @@ -7,7 +10,95 @@ pub enum SchedulerFlags {
DirtyIo = ErlNifTaskFlags::ERL_NIF_DIRTY_JOB_IO_BOUND as isize,
}

impl SchedulerFlags {
fn from(n: isize) -> Self {
match n {
_ if n == Self::Normal as isize => Self::Normal,
_ if n == Self::DirtyCpu as isize => Self::DirtyCpu,
_ if n == Self::DirtyIo as isize => Self::DirtyIo,
_ => unreachable!(),
}
}
}

pub fn consume_timeslice(env: Env, percent: i32) -> bool {
let success = unsafe { rustler_sys::enif_consume_timeslice(env.as_c_arg(), percent) };
success == 1
}

/// Convenience type for scheduling a future invokation of a NIF.
///
/// ## Usage:
///
/// The first generic type should be the NIF that will be scheduled, with a
/// current limitation being that it must be same throughout the lifetime of the
/// NIF.
///
/// The second generic type defined should be the type of the return value.
///
/// Every other generic type is optional, but should reflect the non-`rustler`
/// arguments provided to the NIF, in the same order.
///
/// ## Example:
/// ```rust,ignore
/// #[nif]
/// fn factorial(input: u32, result: Option<u32>) -> Schedule<factorial, u32, u32> {
/// let result = result.unwrap_or(1);
/// if input == 0 {
/// Schedule::Result(result)
/// } else {
/// Schedule::Next2(factorial, input - 1, result * input)
/// }
/// }
/// ```
pub enum Schedule<N: crate::Nif, T, A = (), B = (), C = (), D = (), E = (), F = (), G = ()> {
/// The final result type to return back to the BEAM.
Result(T),
/// Single- and multiple-argument variants that should reflect the scheduled
/// NIF's function signature.
Next(N, A),
Next2(N, A, B),
Next3(N, A, B, C),
Next4(N, A, B, C, D),
Next5(N, A, B, C, D, E),
Next6(N, A, B, C, D, E, F),
Next7(N, A, B, C, D, E, F, G),
}

unsafe impl<N, T, A, B, C, D, E, F, G> NifReturnable for Schedule<N, T, A, B, C, D, E, F, G>
where
N: crate::Nif,
T: crate::Encoder,
A: crate::Encoder,
B: crate::Encoder,
C: crate::Encoder,
D: crate::Encoder,
E: crate::Encoder,
F: crate::Encoder,
G: crate::Encoder,
{
#[inline]
unsafe fn into_returned(self, env: Env) -> NifReturned {
macro_rules! branch {
($($arg:tt),*) => (
NifReturned::Reschedule {
fun_name: std::ffi::CStr::from_ptr(N::NAME as *const c_char).into(),
flags: SchedulerFlags::from(N::FLAGS as isize),
fun: N::RAW_FUNC,
args: vec![$($arg.encode(env).as_c_arg()),*],
}
)
}

match self {
Self::Result(res) => NifReturned::Term(res.encode(env).as_c_arg()),
Self::Next(_, a) => branch!(a),
Self::Next2(_, a, b) => branch!(a, b),
Self::Next3(_, a, b, c) => branch!(a, b, c),
Self::Next4(_, a, b, c, d) => branch!(a, b, c, d),
Self::Next5(_, a, b, c, d, e) => branch!(a, b, c, d, e),
Self::Next6(_, a, b, c, d, e, f) => branch!(a, b, c, d, e, f),
Self::Next7(_, a, b, c, d, e, f, g) => branch!(a, b, c, d, e, f, g),
}
}
}