Fix Java extension loading on Linux (#2898)

kuzudb · Feb 16, 2024 · f1296d5 · f1296d5
1 parent 6b739d9
commit f1296d5
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 7 deletions.
diff --git a/tools/java_api/src/jni/kuzu_java.cpp b/tools/java_api/src/jni/kuzu_java.cpp
@@ -1,5 +1,11 @@
 #include <unordered_map>
 
+#ifdef _WIN32
+// Do nothing on Windows
+#else
+#include <dlfcn.h>
+#endif
+
 // This header is generated at build time. See CMakeLists.txt.
 #include "com_kuzudb_KuzuNative.h"
 #include "common/exception/conversion.h"
@@ -149,6 +155,22 @@ std::unordered_map<std::string, std::unique_ptr<Value>> javaMapToCPPMap(
 /**
  * All Database native functions
  */
+//     protected static native void kuzu_native_reload_library(String lib_path);
+JNIEXPORT void JNICALL Java_com_kuzudb_KuzuNative_kuzu_1native_1reload_1library(
+    JNIEnv* env, jclass, jstring lib_path) {
+#ifdef _WIN32
+// Do nothing on Windows
+#else
+    const char* path = env->GetStringUTFChars(lib_path, JNI_FALSE);
+    void* handle = dlopen(path, RTLD_LAZY | RTLD_GLOBAL);
+    env->ReleaseStringUTFChars(lib_path, path);
+    if (handle == nullptr) {
+        jclass Exception = env->FindClass("java/lang/Exception");
+        auto error = dlerror(); // NOLINT(concurrency-mt-unsafe): load can only be executed in single thread.
+        env->ThrowNew(Exception, error);
+    }
+#endif
+}
 
 JNIEXPORT jlong JNICALL Java_com_kuzudb_KuzuNative_kuzu_1database_1init(JNIEnv* env, jclass,
     jstring database_path, jlong buffer_pool_size, jboolean enable_compression,

diff --git a/tools/java_api/src/main/java/com/kuzudb/KuzuNative.java b/tools/java_api/src/main/java/com/kuzudb/KuzuNative.java
@@ -54,12 +54,21 @@ public class KuzuNative {
             }
             Files.copy(lib_res.openStream(), lib_file, StandardCopyOption.REPLACE_EXISTING);
             new File(lib_file.toString()).deleteOnExit();
-            System.load(lib_file.toAbsolutePath().toString());
+            String lib_path = lib_file.toAbsolutePath().toString();
+            System.load(lib_path);
+            if(os_name.equals("linux")) {
+                kuzu_native_reload_library(lib_path);
+            }
         } catch (IOException e) {
             e.printStackTrace();
         }
     }
 
+    // Hack: Reload the native library again in JNI bindings to work around the 
+    // extension loading issue on Linux as System.load() does not set 
+    // `RTLD_GLOBAL` flag and there is no way to set it in Java.
+    protected static native void kuzu_native_reload_library(String lib_path);
+
     // Database
     protected static native long kuzu_database_init(String database_path, long buffer_pool_size, boolean enable_compression, boolean read_only);
 

diff --git a/tools/java_api/src/test/java/com/kuzudb/test/ExtensionTest.java b/tools/java_api/src/test/java/com/kuzudb/test/ExtensionTest.java
@@ -11,9 +11,8 @@ void HttpfsInstallAndLoad() throws KuzuObjectRefDestroyedException {
         KuzuQueryResult result = conn.query("INSTALL httpfs");
         assertTrue(result.isSuccess());
         result.destroy();
-        // Skip loading the extension for now until the fix is in place
-        // result = conn.query("LOAD EXTENSION httpfs");
-        // assertTrue(result.isSuccess());
-        // result.destroy();
+        result = conn.query("LOAD EXTENSION httpfs");
+        assertTrue(result.isSuccess());
+        result.destroy();
     }
 }
diff --git a/tools/python_api/src_py/__init__.py b/tools/python_api/src_py/__init__.py
@@ -40,11 +40,11 @@
 import sys
 import os
 
-# Set RTLD_GLOBAL and RTLD_NOW flags on Linux to fix the issue with loading
+# Set RTLD_GLOBAL and RTLD_LAZY flags on Linux to fix the issue with loading
 # extensions
 if sys.platform == "linux":
     original_dlopen_flags = sys.getdlopenflags()
-    sys.setdlopenflags(os.RTLD_GLOBAL | os.RTLD_NOW)
+    sys.setdlopenflags(os.RTLD_GLOBAL | os.RTLD_LAZY)
 
 from .database import *
 from .connection import *