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)
}