[env] Introduce JvmUtils to support global JNIEnv

CMakeLists.txt

@@ -1017,7 +1017,8 @@ else()
     env/env_posix.cc
     env/fs_posix.cc
     env/io_posix.cc
-    env/flink/env_flink.cc)
+    env/flink/env_flink.cc
+    env/flink/jvm_util.cc)
 endif()
 
 if(USE_FOLLY_LITE)
@@ -1151,6 +1152,10 @@ endif()
 if(WITH_JNI OR JNI)
   message(STATUS "JNI library is enabled")
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/java)
+  include_directories(${JNI_INCLUDE_DIRS})
+  if(${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
+    include_directories(${JNI_INCLUDE_DIRS}/linux)
+  endif ()
 else()
   message(STATUS "JNI library is disabled")
 endif()

env/flink/jvm_util.cc

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "env/flink/jvm_util.h"
+
+namespace FLINK_NAMESPACE {
+
+std::atomic<JavaVM*> jvm_ = std::atomic<JavaVM*>(nullptr);
+
+JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void* reserved) {
+  JNIEnv* env = nullptr;
+  if (vm->GetEnv((void**)&env, JNI_VERSION_1_8) != JNI_OK) {
+    return -1;
+  }
+
+  jvm_.store(vm);
+  return JNI_VERSION_1_8;
+}
+
+JNIEXPORT void JNICALL JNI_OnUnload(JavaVM* vm, void* reserved) {
+  jvm_.store(nullptr);
+}
+
+void setJVM(JavaVM* jvm) { jvm_.store(jvm); }
+
+JNIEnv* getJNIEnv(bool attach) {
+  JavaVM* jvm = jvm_.load();
+  if (jvm == nullptr) {
+    return nullptr;
+  }
+
+  thread_local JavaEnv env;
+  if (env.getEnv() == nullptr) {
+    auto status = jvm->GetEnv((void**)&(env.getEnv()), JNI_VERSION_1_8);
+    if (attach && (status == JNI_EDETACHED || env.getEnv() == nullptr)) {
+      if (jvm->AttachCurrentThread((void**)&(env.getEnv()), nullptr) ==
+          JNI_OK) {
+        env.setNeedDetach();
+      }
+    }
+  }
+  return env.getEnv();
+}
+}  // namespace FLINK_NAMESPACE

env/flink/jvm_util.h

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <atomic>
+#include <cstdbool>
+#include <cstddef>
+#include <string>
+
+#include "jni.h"
+
+namespace FLINK_NAMESPACE {
+
+extern std::atomic<JavaVM*> jvm_;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void* reserved);
+JNIEXPORT void JNICALL JNI_OnUnload(JavaVM* vm, void* reserved);
+
+#ifdef __cplusplus
+}
+#endif
+
+void setJVM(JavaVM* jvm);
+
+JNIEnv* getJNIEnv(bool attach = true);
+
+static inline std::string parseJavaString(JNIEnv* jni_env,
+                                          jstring java_string) {
+  const char* chars = jni_env->GetStringUTFChars(java_string, nullptr);
+  auto length = jni_env->GetStringUTFLength(java_string);
+  std::string native_string = std::string(chars, length);
+  jni_env->ReleaseStringUTFChars(java_string, chars);
+  return native_string;
+}
+
+class JavaEnv {
+ public:
+  ~JavaEnv() {
+    if (env_ != nullptr && need_detach_) {
+      jvm_.load()->DetachCurrentThread();
+      need_detach_ = false;
+    }
+  }
+
+  JNIEnv*& getEnv() { return env_; }
+
+  void setNeedDetach() { need_detach_ = true; }
+
+ private:
+  JNIEnv* env_ = nullptr;
+  bool need_detach_ = false;
+};
+}  // namespace FLINK_NAMESPACE