From 4714473758a9dd3769b01d989e496ee059a9fcf4 Mon Sep 17 00:00:00 2001
From: 0lai0 <yuchenlai87@gmail.com>
Date: Tue, 5 May 2026 16:45:02 +0800
Subject: [PATCH] enhance JNI local frame management with error handling in
 with_env function

---
 native/jni-bridge/src/lib.rs | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/native/jni-bridge/src/lib.rs b/native/jni-bridge/src/lib.rs
index 5b0c0a4a56..21c647135b 100644
--- a/native/jni-bridge/src/lib.rs
+++ b/native/jni-bridge/src/lib.rs
@@ -26,7 +26,7 @@ use jni::{
     errors::Error,
     objects::{JMethodID, JObject, JString, JThrowable, JValueOwned},
     signature::ReturnType,
-    Env, JavaVM,
+    Env, JavaVM, DEFAULT_LOCAL_FRAME_CAPACITY,
 };
 use once_cell::sync::OnceCell;
 
@@ -34,6 +34,17 @@ use errors::{CometError, CometResult};
 
 pub mod errors;
 
+enum LocalFrameError<E> {
+    Closure(E),
+    Jni(jni::errors::Error),
+}
+
+impl<E> From<jni::errors::Error> for LocalFrameError<E> {
+    fn from(source: jni::errors::Error) -> Self {
+        Self::Jni(source)
+    }
+}
+
 /// Global reference to the Java VM, initialized during native library setup.
 pub static JAVA_VM: OnceCell<JavaVM> = OnceCell::new();
 
@@ -300,6 +311,13 @@ impl JVMClasses<'_> {
     }
 
     /// Runs a closure with an attached JNI environment for the current thread.
+    ///
+    /// The closure executes inside a JNI local frame. All JNI local references
+    /// created inside the closure are freed when the frame is popped on return.
+    /// Callers must return only Rust values (e.g. `Vec`, `ArrayRef`, scalars);
+    /// returning a raw JNI local reference will produce a dangling reference
+    /// after the frame is popped. Panic safety is guaranteed by
+    /// `jni::Env::with_local_frame`, which pops the frame even on unwinding.
     pub fn with_env<T, E, F>(f: F) -> Result<T, E>
     where
         F: FnOnce(&mut Env) -> Result<T, E>,
@@ -316,7 +334,15 @@ impl JVMClasses<'_> {
                 .attach_current_thread_guard(Default::default, &mut scope)
                 .map_err(CometError::from)
                 .map_err(E::from)?;
-            f(guard.borrow_env_mut())
+            guard
+                .borrow_env_mut()
+                .with_local_frame(DEFAULT_LOCAL_FRAME_CAPACITY, |env| {
+                    f(env).map_err(LocalFrameError::Closure)
+                })
+                .map_err(|err| match err {
+                    LocalFrameError::Closure(err) => err,
+                    LocalFrameError::Jni(err) => E::from(CometError::from(err)),
+                })
         }
     }
 }