diff --git a/src/environment.rs b/src/environment.rs index e4e31c5..63732e8 100644 --- a/src/environment.rs +++ b/src/environment.rs @@ -38,11 +38,7 @@ //! [global thread pool]: EnvironmentBuilder::with_global_thread_pool //! [custom logger]: EnvironmentBuilder::with_logger -use alloc::{ - boxed::Box, - string::String, - sync::{Arc, Weak} -}; +use alloc::{boxed::Box, string::String, sync::Arc}; use core::{ any::Any, ffi::c_void, @@ -66,11 +62,51 @@ use crate::{ util::{Mutex, OnceLock, STACK_EXECUTION_PROVIDERS, run_on_drop, with_cstr} }; -/// Hold onto a weak reference here so that the environment is dropped when all sessions under it are. Statics don't run -/// destructors; holding a strong reference to the environment here thus leads to issues, so instead of holding the -/// environment for the entire duration of the process, we keep `Arc` references of this weak `Environment` in sessions. -/// That way, the environment is actually dropped when it is no longer needed. -static G_ENV: Mutex>> = Mutex::new(None); +static G_ENV: Mutex>> = Mutex::new(None); + +// Rust doesn't run destructors for statics, but ONNX Runtime is *very* particular about `ReleaseEnv` being called +// before any C++ destructors are called. In order to drop the environment, we have to release the reference held in +// `G_ENV` at the end of the program, but before C++ destructors are called. On Linux & Windows (surprisingly), this is +// fairly simple: just put it in a custom linker section. +// +// `G_ENV` used to be `Mutex>`, which was much nicer, but apparently you can only ever call +// `CreateEnv` once throughout the lifetime of the process, *even if* the last env was `ReleaseEnv`'d. So once all +// `Session`s fell out of scope, if you ever tried to create another one, you'd crash. Grand. +#[cfg_attr(any(target_os = "linux", target_os = "android"), unsafe(link_section = ".text.exit"))] +unsafe extern "C" fn release_env_on_exit(#[cfg(target_vendor = "apple")] _: *const ()) { + G_ENV.lock().take(); +} + +#[used] +#[cfg(all(not(windows), not(target_vendor = "apple"), not(target_arch = "wasm32")))] +#[unsafe(link_section = ".fini_array")] +static _ON_EXIT: unsafe extern "C" fn() = release_env_on_exit; +#[used] +#[cfg(windows)] +#[unsafe(link_section = ".CRT$XLB")] +static _ON_EXIT: unsafe extern "system" fn(module: *mut (), reason: u32, reserved: *mut ()) = { + unsafe extern "system" fn on_exit(_h: *mut (), reason: u32, _pv: *mut ()) { + // XLB gets called on both init & exit (?). Also, XLA never gets called, and no one online ever mentions that. + // Only do destructor things if we're actually exiting the process (DLL_PROCESS_EXIT = 0) + if reason == 0 { + unsafe { release_env_on_exit() }; + } + } + on_exit +}; + +// macOS used to have the __mod_term_func section which worked similar to `.fini_array`, but one day they just decided +// to remove it I guess? So we have to set an atexit handler instead. But normal atexit doesn't work, we need to use +// __cxa_atexit. And if you register it too early in the program (i.e. in __mod_init_func), it'll fire *after* C++ +// destructors. So we call this after we create the environment instead. This shit took years off my life. +#[cfg(target_vendor = "apple")] +fn register_atexit() { + unsafe extern "C" { + static __dso_handle: *const (); + fn __cxa_atexit(cb: unsafe extern "C" fn(_: *const ()), arg: *const (), dso_handle: *const ()); + } + unsafe { __cxa_atexit(release_env_on_exit, core::ptr::null(), __dso_handle) }; +} static G_ENV_OPTIONS: OnceLock = OnceLock::new(); @@ -212,15 +248,17 @@ impl Drop for Environment { /// has fallen out of usage), a new environment will be created & committed. pub fn current() -> Result> { let mut env_lock = G_ENV.lock(); - if let Some(env) = env_lock.as_ref() - && let Some(upgraded) = Weak::upgrade(env) - { - return Ok(upgraded); + if let Some(env) = env_lock.as_ref() { + return Ok(env.clone()); } let options = G_ENV_OPTIONS.get_or_init(EnvironmentBuilder::new); let env = options.create_environment().map(Arc::new)?; - *env_lock = Some(Arc::downgrade(&env)); + *env_lock = Some(Arc::clone(&env)); + + #[cfg(target_vendor = "apple")] + register_atexit(); + Ok(env) }