Skip to content

Commit

Permalink
Allow sends from callback environments (#634)
Browse files Browse the repository at this point in the history
  • Loading branch information
filmor authored Jul 9, 2024
1 parent 00067e3 commit 51d0cc3
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 30 deletions.
56 changes: 30 additions & 26 deletions rustler/src/env.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,21 @@ use std::sync::{Arc, Weak};
/// auto-convert a `Env<'a>` to a `Env<'b>`.
type EnvId<'a> = PhantomData<*mut &'a u8>;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum EnvKind {
ProcessBound,
Callback,
Init,
ProcessIndependent,
}

/// On each NIF call, a Env is passed in. The Env is used for most operations that involve
/// communicating with the BEAM, like decoding and encoding terms.
///
/// There is no way to allocate a Env at the moment, but this may be possible in the future.
#[derive(Clone, Copy)]
pub struct Env<'a> {
pub(crate) init: bool,
pub(crate) kind: EnvKind,
env: NIF_ENV,
id: EnvId<'a>,
}
Expand All @@ -38,6 +46,19 @@ impl<'a, 'b> PartialEq<Env<'b>> for Env<'a> {
pub struct SendError;

impl<'a> Env<'a> {
#[doc(hidden)]
pub(crate) unsafe fn new_internal<T>(
_lifetime_marker: &'a T,
env: NIF_ENV,
kind: EnvKind,
) -> Env<'a> {
Env {
kind,
env,
id: PhantomData,
}
}

/// Create a new Env. For the `_lifetime_marker` argument, pass a
/// reference to any local variable that has its own lifetime, different
/// from all other `Env` values. The purpose of the argument is to make
Expand All @@ -48,18 +69,12 @@ impl<'a> Env<'a> {
/// # Unsafe
/// Don't create multiple `Env`s with the same lifetime.
pub unsafe fn new<T>(_lifetime_marker: &'a T, env: NIF_ENV) -> Env<'a> {
Env {
init: false,
env,
id: PhantomData,
}
Self::new_internal(_lifetime_marker, env, EnvKind::ProcessBound)
}

#[doc(hidden)]
pub unsafe fn new_init_env<T>(_lifetime_marker: &'a T, env: NIF_ENV) -> Env<'a> {
let mut res = Self::new(_lifetime_marker, env);
res.init = true;
res
Self::new_internal(_lifetime_marker, env, EnvKind::Init)
}

pub fn as_c_arg(self) -> NIF_ENV {
Expand All @@ -85,26 +100,15 @@ impl<'a> Env<'a> {
///
/// The result indicates whether the send was successful, see also
/// [enif\_send](https://www.erlang.org/doc/man/erl_nif.html#enif_send).
///
/// # Panics
///
/// Panics if the above rules are broken (by trying to send a message from
/// an `OwnedEnv` on a thread that's managed by the Erlang VM).
///
pub fn send(self, pid: &LocalPid, message: impl Encoder) -> Result<(), SendError> {
let thread_type = unsafe { rustler_sys::enif_thread_type() };
let env = if thread_type == rustler_sys::ERL_NIF_THR_UNDEFINED {
ptr::null_mut()
} else if thread_type == rustler_sys::ERL_NIF_THR_NORMAL_SCHEDULER
|| thread_type == rustler_sys::ERL_NIF_THR_DIRTY_CPU_SCHEDULER
|| thread_type == rustler_sys::ERL_NIF_THR_DIRTY_IO_SCHEDULER
{
// Panic if `self` is not the environment of the calling process.
self.pid();
let env = if is_scheduler_thread() {
if self.kind == EnvKind::ProcessIndependent {
return Err(SendError);
}

self.as_c_arg()
} else {
panic!("Env::send(): unrecognized calling thread type");
ptr::null_mut()
};

let message = message.encode(self);
Expand Down Expand Up @@ -213,7 +217,7 @@ impl OwnedEnv {
where
F: FnOnce(Env<'a>) -> R,
{
let env = unsafe { Env::new(&(), *self.env) };
let env = unsafe { Env::new_internal(&(), *self.env, EnvKind::ProcessIndependent) };
closure(env)
}

Expand Down
7 changes: 4 additions & 3 deletions rustler/src/resource/registration.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use super::traits;
use super::util::align_alloced_mem_for_struct;
use super::ResourceInitError;
use crate::env::EnvKind;
use crate::{Env, LocalPid, Monitor, Resource};
use rustler_sys::ErlNifResourceDtor;
use rustler_sys::{
Expand Down Expand Up @@ -98,7 +99,7 @@ impl Registration {
/// will only succeed when called from the `load` callback and if this type has not yet been
/// registered.
pub fn register(&self, env: Env) -> Result<(), ResourceInitError> {
if !env.init {
if env.kind != EnvKind::Init {
return Err(ResourceInitError);
}

Expand Down Expand Up @@ -127,7 +128,7 @@ unsafe extern "C" fn resource_destructor<T>(_env: *mut ErlNifEnv, handle: *mut c
where
T: Resource,
{
let env = Env::new(&_env, _env);
let env = Env::new_internal(&_env, _env, EnvKind::Callback);
let aligned = align_alloced_mem_for_struct::<T>(handle);
// Destructor takes ownership, thus the resource object will be dropped after the function has
// run.
Expand All @@ -143,7 +144,7 @@ unsafe extern "C" fn resource_down<T: Resource>(
pid: *const ErlNifPid,
mon: *const ErlNifMonitor,
) {
let env = Env::new(&env, env);
let env = Env::new_internal(&env, env, EnvKind::Callback);
let aligned = align_alloced_mem_for_struct::<T>(obj);
let res = &*(aligned as *const T);
let pid = LocalPid::from_c_arg(*pid);
Expand Down
1 change: 0 additions & 1 deletion rustler_codegen/src/init.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,6 @@ impl From<InitMacroInput> for proc_macro2::TokenStream {
) -> rustler::codegen_runtime::c_int {
unsafe {
let mut env = rustler::Env::new_init_env(&env, env);
// TODO: If an unwrap ever happens, we will unwind right into C! Fix this!
let load_info = rustler::Term::new(env, load_info);

if !rustler::codegen_runtime::ResourceRegistration::register_all_collected(env).is_ok() {
Expand Down

0 comments on commit 51d0cc3

Please sign in to comment.