From b7fe9594556ed736962c23da744e6ad5f0533428 Mon Sep 17 00:00:00 2001 From: Gus Caplan Date: Wed, 28 Feb 2024 16:17:46 -0800 Subject: [PATCH 1/2] yields --- rustler/src/codegen_runtime.rs | 7 +-- rustler/src/lib.rs | 3 + rustler/src/nif.rs | 9 ++- rustler/src/yield.rs | 60 +++++++++++++++++++ rustler_tests/lib/rustler_test.ex | 1 + rustler_tests/native/rustler_test/src/lib.rs | 1 + .../native/rustler_test/src/test_dirty.rs | 14 ++++- rustler_tests/test/dirty_test.exs | 4 ++ 8 files changed, 89 insertions(+), 10 deletions(-) create mode 100644 rustler/src/yield.rs diff --git a/rustler/src/codegen_runtime.rs b/rustler/src/codegen_runtime.rs index d6b97a33..75357d19 100644 --- a/rustler/src/codegen_runtime.rs +++ b/rustler/src/codegen_runtime.rs @@ -1,6 +1,5 @@ //! Functions used by runtime generated code. Should not be used. -use std::ffi::CString; use std::fmt; use crate::{Encoder, Env, OwnedBinary, Term}; @@ -51,9 +50,9 @@ pub enum NifReturned { Raise(NIF_TERM), BadArg, Reschedule { - fun_name: CString, + fun_name: *const c_char, flags: crate::schedule::SchedulerFlags, - fun: unsafe extern "C" fn(NIF_ENV, i32, *const NIF_TERM) -> NIF_TERM, + fun: crate::nif::RawFunc, args: Vec, }, } @@ -73,7 +72,7 @@ impl NifReturned { args, } => rustler_sys::enif_schedule_nif( env.as_c_arg(), - fun_name.as_ptr() as *const c_char, + fun_name, flags as i32, fun, args.len() as i32, diff --git a/rustler/src/lib.rs b/rustler/src/lib.rs index ce700ef6..2b44e165 100644 --- a/rustler/src/lib.rs +++ b/rustler/src/lib.rs @@ -62,6 +62,9 @@ pub use crate::error::Error; pub mod r#return; pub use crate::r#return::Return; +pub mod r#yield; +pub use crate::r#yield::Yield; + #[doc(hidden)] mod nif; pub use nif::Nif; diff --git a/rustler/src/nif.rs b/rustler/src/nif.rs index 7ee07ede..9e38615b 100644 --- a/rustler/src/nif.rs +++ b/rustler/src/nif.rs @@ -1,13 +1,12 @@ use crate::codegen_runtime::{c_char, c_int, c_uint, DEF_NIF_FUNC, NIF_ENV, NIF_TERM}; +pub(crate) type RawFunc = + unsafe extern "C" fn(nif_env: NIF_ENV, argc: c_int, argv: *const NIF_TERM) -> NIF_TERM; + pub trait Nif { const NAME: *const c_char; const ARITY: c_uint; const FLAGS: c_uint; const FUNC: DEF_NIF_FUNC; - const RAW_FUNC: unsafe extern "C" fn( - nif_env: NIF_ENV, - argc: c_int, - argv: *const NIF_TERM, - ) -> NIF_TERM; + const RAW_FUNC: RawFunc; } diff --git a/rustler/src/yield.rs b/rustler/src/yield.rs new file mode 100644 index 00000000..5684ac3d --- /dev/null +++ b/rustler/src/yield.rs @@ -0,0 +1,60 @@ +use crate::{ + codegen_runtime::{NifReturnable, NifReturned}, + nif::RawFunc, + schedule::SchedulerFlags, + Env, Nif, Term, +}; + +/// If a function needs to break up long-running work, it can return +/// a `Yield`. The runtime will resume execution by calling the provided +/// NIF. +/// +/// ```rust +/// use rustler::{Encoder, Env, Yield}; +/// +/// #[rustler::nif] +/// fn start_expensive_work(env: Env) -> Yield { +/// let args = vec![42_i32.encode(env)]; +/// Yield::to(perform_expensive_work, args) +/// } +/// +/// #[rustler::nif] +/// fn perform_expensive_work(env: Env, value: i32) { +/// assert_eq!(value, 42); +/// } +/// ``` +pub struct Yield<'a> { + fun_name: *const i8, + fun: RawFunc, + flags: SchedulerFlags, + args: Vec>, +} + +impl<'a> Yield<'a> { + /// Yield back to the runtime with the provided nif. + pub fn to(nif: N, args: Vec>) -> Self { + Self::to_with_flags(nif, SchedulerFlags::Normal, args) + } + + /// Yield back to the runtime with the provided function and scheduler flags. + pub fn to_with_flags(nif: N, flags: SchedulerFlags, args: Vec>) -> Self { + let _ = nif; + Self { + fun_name: N::NAME, + fun: N::RAW_FUNC, + flags, + args, + } + } +} + +unsafe impl<'a> NifReturnable for Yield<'a> { + unsafe fn into_returned(self, _env: Env) -> NifReturned { + NifReturned::Reschedule { + fun_name: self.fun_name, + fun: self.fun, + flags: self.flags, + args: self.args.into_iter().map(|term| term.as_c_arg()).collect(), + } + } +} diff --git a/rustler_tests/lib/rustler_test.ex b/rustler_tests/lib/rustler_test.ex index ce424d40..dcfadafd 100644 --- a/rustler_tests/lib/rustler_test.ex +++ b/rustler_tests/lib/rustler_test.ex @@ -111,6 +111,7 @@ defmodule RustlerTest do def dirty_io(), do: err() def dirty_cpu(), do: err() + def yields(), do: err() def sum_range(_), do: err() diff --git a/rustler_tests/native/rustler_test/src/lib.rs b/rustler_tests/native/rustler_test/src/lib.rs index efe2372e..1bdbfa98 100644 --- a/rustler_tests/native/rustler_test/src/lib.rs +++ b/rustler_tests/native/rustler_test/src/lib.rs @@ -82,6 +82,7 @@ rustler::init!( test_codegen::tuplestruct_record_echo, test_dirty::dirty_cpu, test_dirty::dirty_io, + test_dirty::yields, test_range::sum_range, test_error::bad_arg_error, test_error::atom_str_error, diff --git a/rustler_tests/native/rustler_test/src/test_dirty.rs b/rustler_tests/native/rustler_test/src/test_dirty.rs index 65082718..111dc0d3 100644 --- a/rustler_tests/native/rustler_test/src/test_dirty.rs +++ b/rustler_tests/native/rustler_test/src/test_dirty.rs @@ -1,4 +1,4 @@ -use rustler::Atom; +use rustler::{Atom, Encoder, Env, Yield}; use std::time::Duration; mod atoms { @@ -22,3 +22,15 @@ pub fn dirty_io() -> Atom { atoms::ok() } + +#[rustler::nif] +fn yield_resume(input: i32) -> Atom { + assert_eq!(input, 42); + atoms::ok() +} + +#[rustler::nif] +pub fn yields(env: Env) -> Yield { + let term = 42.encode(env); + Yield::to(yield_resume, vec![term]) +} diff --git a/rustler_tests/test/dirty_test.exs b/rustler_tests/test/dirty_test.exs index e48ec15a..980a04d3 100644 --- a/rustler_tests/test/dirty_test.exs +++ b/rustler_tests/test/dirty_test.exs @@ -8,4 +8,8 @@ defmodule RustlerTest.DirtyTest do test "dirty cpu" do RustlerTest.dirty_cpu() end + + test "yields" do + RustlerTest.yields() + end end From 63aba42aced7b99b60caf3cfd3b120204c85830a Mon Sep 17 00:00:00 2001 From: Gus Caplan Date: Wed, 28 Feb 2024 18:34:05 -0800 Subject: [PATCH 2/2] wip yield --- rustler/src/lib.rs | 2 +- rustler/src/yield.rs | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/rustler/src/lib.rs b/rustler/src/lib.rs index 2b44e165..126e3a12 100644 --- a/rustler/src/lib.rs +++ b/rustler/src/lib.rs @@ -63,7 +63,7 @@ pub mod r#return; pub use crate::r#return::Return; pub mod r#yield; -pub use crate::r#yield::Yield; +pub use crate::r#yield::{Branch, Yield}; #[doc(hidden)] mod nif; diff --git a/rustler/src/yield.rs b/rustler/src/yield.rs index 5684ac3d..5b297fb7 100644 --- a/rustler/src/yield.rs +++ b/rustler/src/yield.rs @@ -58,3 +58,20 @@ unsafe impl<'a> NifReturnable for Yield<'a> { } } } + +pub enum Branch<'a, T> { + Yield(Yield<'a>), + Stop(T), +} + +unsafe impl<'a, T> NifReturnable for Branch<'a, T> +where + T: NifReturnable, +{ + unsafe fn into_returned(self, env: Env) -> NifReturned { + match self { + Self::Yield(inner) => inner.into_returned(env), + Self::Stop(inner) => inner.into_returned(env), + } + } +}