fix: one environment (#542)

Create exactly one (1) environment throughout the entire duration of the
process and release it when it is no longer used. Hopefully the last
time I'll ever have to deal with this.

The old `Weak<Environment>` solution was flawed because it would create
another environment after the old one was released, which is for some
reason not allowed.

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<Option<Weak<Environment>>> = Mutex::new(None);
+static G_ENV: Mutex<Option<Arc<Environment>>> = 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<Weak<Environment>>`, 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<EnvironmentBuilder> = 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<Arc<Environment>> {
 	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)
 }