[env] Introduce JvmUtils to support global JNIEnv

(cherry picked from commit 44debe7)

CMakeLists.txt

@@ -1020,7 +1020,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)
@@ -1165,6 +1166,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 ROCKSDB_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 ROCKSDB_NAMESPACE

env/flink/jvm_util.h

@@ -0,0 +1,74 @@
+/*
+ * 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"
+#include "rocksdb/env.h"
+
+namespace ROCKSDB_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 ROCKSDB_NAMESPACE

src.mk

@@ -114,6 +114,7 @@ LIB_SOURCES =                                                   \
   env/mock_env.cc                                               \
   env/unique_id_gen.cc                                          \
   env/flink/env_flink.cc										\
+  env/flink/jvm_util.cc											\
   file/delete_scheduler.cc                                      \
   file/file_prefetch_buffer.cc                                  \
   file/file_util.cc                                             \