(arrow) branch main updated: GH-20379: [Java] Dataset Failed to update reservation while freeing bytes (#40101)

[lidavidm]

This is an automated email from the ASF dual-hosted git repository.

lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new df9e0c1264 GH-20379: [Java] Dataset Failed to update reservation while 
freeing bytes (#40101)
df9e0c1264 is described below

commit df9e0c1264e3d7b83f913bccaec2c9a85fe6777e
Author: Florian Bernard <flor...@flob.fr>
AuthorDate: Tue Feb 20 14:15:44 2024 +0100

    GH-20379: [Java] Dataset Failed to update reservation while freeing bytes 
(#40101)
    
    ### Rationale for this change
    Better controls JNI Thread management in java dataset module to fix #20379
    Re-use the same code found in the java arrow-c-data module : 
https://github.com/apache/arrow/blob/main/java/c/src/main/cpp/jni_wrapper.cc#L107
    
    May JNIEnvGuard class code can be put in a common place for both libraries 
...
    
    ### What changes are included in this PR?
    N/A
    
    ### Are these changes tested?
    These changes has been tested with :  
https://gist.github.com/fb64/71880cde297bc5234b02b68b785670fd
    on Linux X86_64 architecture
    
    ### Are there any user-facing changes?
    N/A
    
    * Closes: #20379
    
    Authored-by: Florian Bernard <florian.bern...@openairlines.com>
    Signed-off-by: David Li <li.david...@gmail.com>
---
 java/dataset/src/main/cpp/jni_wrapper.cc | 62 +++++++++++++++++++++++++-------
 1 file changed, 50 insertions(+), 12 deletions(-)

diff --git a/java/dataset/src/main/cpp/jni_wrapper.cc 
b/java/dataset/src/main/cpp/jni_wrapper.cc
index d2d976677b..19a43c8d2f 100644
--- a/java/dataset/src/main/cpp/jni_wrapper.cc
+++ b/java/dataset/src/main/cpp/jni_wrapper.cc
@@ -83,6 +83,40 @@ void ThrowIfError(const arrow::Status& status) {
   }
 }
 
+class JNIEnvGuard {
+ public:
+  explicit JNIEnvGuard(JavaVM* vm) : vm_(vm), env_(nullptr), 
should_detach_(false) {
+    JNIEnv* env;
+    jint code = vm->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION);
+    if (code == JNI_EDETACHED) {
+      JavaVMAttachArgs args;
+      args.version = JNI_VERSION;
+      args.name = NULL;
+      args.group = NULL;
+      code = vm->AttachCurrentThread(reinterpret_cast<void**>(&env), &args);
+      should_detach_ = (code == JNI_OK);
+    }
+    if (code != JNI_OK) {
+      ThrowPendingException("Failed to attach the current thread to a Java 
VM");
+    }
+    env_ = env;
+  }
+
+  JNIEnv* env() { return env_; }
+
+  ~JNIEnvGuard() {
+    if (should_detach_) {
+      vm_->DetachCurrentThread();
+      should_detach_ = false;
+    }
+  }
+
+ private:
+  JavaVM* vm_;
+  JNIEnv* env_;
+  bool should_detach_;
+};
+
 template <typename T>
 T JniGetOrThrow(arrow::Result<T> result) {
   const arrow::Status& status = result.status();
@@ -126,23 +160,27 @@ class ReserveFromJava : public 
arrow::dataset::jni::ReservationListener {
       : vm_(vm), java_reservation_listener_(java_reservation_listener) {}
 
   arrow::Status OnReservation(int64_t size) override {
-    JNIEnv* env;
-    if (vm_->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION) != JNI_OK) {
-      return arrow::Status::Invalid("JNIEnv was not attached to current 
thread");
+    try {
+      JNIEnvGuard guard(vm_);
+      JNIEnv* env = guard.env();
+      env->CallObjectMethod(java_reservation_listener_, reserve_memory_method, 
size);
+      RETURN_NOT_OK(arrow::dataset::jni::CheckException(env));
+      return arrow::Status::OK();
+    } catch (const JniPendingException& e) {
+      return arrow::Status::Invalid(e.what());
     }
-    env->CallObjectMethod(java_reservation_listener_, reserve_memory_method, 
size);
-    RETURN_NOT_OK(arrow::dataset::jni::CheckException(env));
-    return arrow::Status::OK();
   }
 
   arrow::Status OnRelease(int64_t size) override {
-    JNIEnv* env;
-    if (vm_->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION) != JNI_OK) {
-      return arrow::Status::Invalid("JNIEnv was not attached to current 
thread");
+    try {
+      JNIEnvGuard guard(vm_);
+      JNIEnv* env = guard.env();
+      env->CallObjectMethod(java_reservation_listener_, 
unreserve_memory_method, size);
+      RETURN_NOT_OK(arrow::dataset::jni::CheckException(env));
+      return arrow::Status::OK();
+    } catch (const JniPendingException& e) {
+      return arrow::Status::Invalid(e.what());
     }
-    env->CallObjectMethod(java_reservation_listener_, unreserve_memory_method, 
size);
-    RETURN_NOT_OK(arrow::dataset::jni::CheckException(env));
-    return arrow::Status::OK();
   }
 
   jobject GetJavaReservationListener() { return java_reservation_listener_; }