apacheGH-20379: [Java] Dataset Failed to update reservation while fre…

…eing bytes: JNIEnv was not attached to current thread
fb64 · Feb 16, 2024 · 9aab65c · 9aab65c
1 parent a03d957
commit 9aab65c
Showing 1 changed file with 38 additions and 8 deletions.
diff --git 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), 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:
+  bool should_detach_;
+  JavaVM* vm_;
+  JNIEnv* env_;
+};
+
 template <typename T>
 T JniGetOrThrow(arrow::Result<T> result) {
   const arrow::Status& status = result.status();
@@ -126,20 +160,16 @@ 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");
-    }
+    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();
   }
 
   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");
-    }
+    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();