From f1296d552ad2a183833d02922f5453029394a925 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=9B=A7=E5=9B=A7?= <liuc223@gmail.com>
Date: Fri, 16 Feb 2024 12:07:09 +0800
Subject: [PATCH] Fix Java extension loading on Linux (#2898)

---
 tools/java_api/src/jni/kuzu_java.cpp          | 22 +++++++++++++++++++
 .../src/main/java/com/kuzudb/KuzuNative.java  | 11 +++++++++-
 .../java/com/kuzudb/test/ExtensionTest.java   |  7 +++---
 tools/python_api/src_py/__init__.py           |  4 ++--
 4 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/tools/java_api/src/jni/kuzu_java.cpp b/tools/java_api/src/jni/kuzu_java.cpp
index 6add4a319b..159c677272 100644
--- 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
index 87ec6ca98d..9f37e9f0ed 100644
--- 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
index bab76e9c93..ec8d65321f 100644
--- 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
index b478c830d5..6a71049e29 100644
--- 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 *