Fix leak issues (#385)

Author: cretz

temporalio/Cargo.lock

@@ -26,3 +26,7 @@ tokio-util = "0.7"
 tonic = { workspace = true }
 tracing = "0.1"
 url = "2.5"
+dhat = { version = "0.3", optional = true }
+
+[features]
+dhat-heap = ["dhat"]

temporalio/ext/Cargo.toml

@@ -1,5 +1,12 @@
 use magnus::{Error, ExceptionClass, RModule, Ruby, prelude::*, value::Lazy};
 
+#[cfg(feature = "dhat-heap")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
+
+#[cfg(feature = "dhat-heap")]
+static DHAT_PROFILER: std::sync::Mutex<Option<dhat::Profiler>> = std::sync::Mutex::new(None);
+
 mod client;
 mod client_rpc_generated;
 mod envconfig;
@@ -58,12 +65,54 @@ macro_rules! lazy_id {
 fn init(ruby: &Ruby) -> Result<(), Error> {
     Lazy::force(&ROOT_ERR, ruby);
 
+    #[cfg(feature = "dhat-heap")]
+    {
+        let profiler = dhat::Profiler::builder()
+            .file_name("dhat-heap.json")
+            .build();
+        *DHAT_PROFILER.lock().unwrap() = Some(profiler);
+        eprintln!("[dhat] Heap profiling enabled.");
+    }
+
     client::init(ruby)?;
     envconfig::init(ruby)?;
     metric::init(ruby)?;
     runtime::init(ruby)?;
     testing::init(ruby)?;
     worker::init(ruby)?;
 
+    #[cfg(feature = "dhat-heap")]
+    {
+        let bridge_mod = ruby.get_inner(&ROOT_MOD);
+        bridge_mod
+            .define_module_function("dhat_heap_stats", magnus::function!(dhat_heap_stats, 0))?;
+        bridge_mod.define_module_function(
+            "dhat_dump_and_stop",
+            magnus::function!(dhat_dump_and_stop, 0),
+        )?;
+    }
+
     Ok(())
 }
+
+/// Print current dhat heap stats to stderr
+#[cfg(feature = "dhat-heap")]
+fn dhat_heap_stats() {
+    let stats = dhat::HeapStats::get();
+    eprintln!(
+        "[dhat] curr_bytes={} curr_blocks={} total_bytes={} total_blocks={}",
+        stats.curr_bytes, stats.curr_blocks, stats.total_bytes, stats.total_blocks
+    );
+}
+
+/// Drop the profiler to force writing the JSON profile
+#[cfg(feature = "dhat-heap")]
+fn dhat_dump_and_stop() {
+    let mut guard = DHAT_PROFILER.lock().unwrap();
+    if let Some(profiler) = guard.take() {
+        drop(profiler); // This triggers the profile write to dhat-heap.json
+        eprintln!("[dhat] Profile written to dhat-heap.json");
+    } else {
+        eprintln!("[dhat] Profiler already stopped or not initialized");
+    }
+}

temporalio/ext/src/lib.rs

@@ -88,22 +88,30 @@ where
     where
         U: FnMut(),
     {
-        let mut func: U = unsafe { *Box::from_raw(data as _) };
-
+        // Borrow rather than take ownership — the caller frees after
+        // rb_thread_call_without_gvl returns. This avoids leaking the
+        // unblock closure when Ruby never invokes it (the common case).
+        let func: &mut U = unsafe { &mut *(data as *mut U) };
         func();
     }
 
     let boxed_func = Box::new(func);
     let boxed_unblock = Box::new(unblock);
+    let unblock_ptr = Box::into_raw(boxed_unblock);
 
     unsafe {
         let result = rb_sys::rb_thread_call_without_gvl(
             Some(anon_func::<F, R>),
             Box::into_raw(boxed_func) as *mut _,
             Some(anon_unblock::<U>),
-            Box::into_raw(boxed_unblock) as *mut _,
+            unblock_ptr as *mut _,
         );
 
+        // Free the unblock closure. By the time rb_thread_call_without_gvl
+        // returns, anon_unblock (if called at all) has already completed,
+        // so this is safe.
+        drop(Box::from_raw(unblock_ptr));
+
         *Box::from_raw(result as _)
     }
 }