Merge pull request #97 from Peefy/feat-java-plugin

Peefy · web-flow · commit 41cba02d7866 · 2024-06-28T16:45:59.000+08:00
feat: impl kcl java plugin
diff --git a/java/Makefile b/java/Makefile
@@ -7,10 +7,10 @@ fmt:
 	mvn formatter:format
 
 pkg:
-	mvn clean package
+	mvn clean package -Dcargo-build.profile=release
 
 deploy:
-	mvn clean deploy
+	mvn clean deploy -Dcargo-build.profile=release
 
 test:
-	mvn clean test
+	mvn clean test -Dcargo-build.profile=release
diff --git a/java/src/lib.rs b/java/src/lib.rs
@@ -8,20 +8,25 @@ extern crate once_cell;
 extern crate prost;
 
 use anyhow::Result;
-use jni::objects::{JByteArray, JClass, JObject};
+use jni::objects::{GlobalRef, JByteArray, JClass, JObject, JString};
 use jni::sys::jbyteArray;
 use jni::JNIEnv;
-use kclvm_api::call;
+use jni::JavaVM;
+use kclvm_api::call_with_plugin_agent;
 use kclvm_api::gpyrpc::LoadPackageArgs;
 use kclvm_api::service::KclvmServiceImpl;
 use kclvm_parser::KCLModuleCache;
 use kclvm_sema::resolver::scope::KCLScopeCache;
 use lazy_static::lazy_static;
 use once_cell::sync::OnceCell;
 use prost::Message;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
 use std::sync::Mutex;
 
 lazy_static! {
+    static ref JVM: Mutex<Option<JavaVM>> = Mutex::new(None);
+    static ref CALLBACK_OBJ: Mutex<Option<GlobalRef>> = Mutex::new(None);
     static ref MODULE_CACHE: Mutex<OnceCell<KCLModuleCache>> = Mutex::new(OnceCell::new());
     static ref SCOPE_CACHE: Mutex<OnceCell<KCLScopeCache>> = Mutex::new(OnceCell::new());
 }
@@ -33,12 +38,20 @@ pub extern "system" fn Java_com_kcl_api_API_callNative(
     name: JByteArray,
     args: JByteArray,
 ) -> jbyteArray {
-    intern_call_native(&mut env, name, args).unwrap_or_else(|e| {
+    intern_call_native_with_plugin(&mut env, name, args).unwrap_or_else(|e| {
         let _ = throw(&mut env, e);
         JObject::default().into_raw()
     })
 }
 
+#[no_mangle]
+pub extern "system" fn Java_com_kcl_api_API_registerPluginContext(env: JNIEnv, obj: JObject) {
+    let jvm = env.get_java_vm().unwrap();
+    *JVM.lock().unwrap() = Some(jvm);
+    let global_ref = env.new_global_ref(obj).unwrap();
+    *CALLBACK_OBJ.lock().unwrap() = Some(global_ref);
+}
+
 #[no_mangle]
 pub extern "system" fn Java_com_kcl_api_API_loadPackageWithCache(
     mut env: JNIEnv,
@@ -51,10 +64,14 @@ pub extern "system" fn Java_com_kcl_api_API_loadPackageWithCache(
     })
 }
 
-fn intern_call_native(env: &mut JNIEnv, name: JByteArray, args: JByteArray) -> Result<jbyteArray> {
+fn intern_call_native_with_plugin(
+    env: &mut JNIEnv,
+    name: JByteArray,
+    args: JByteArray,
+) -> Result<jbyteArray> {
     let name = env.convert_byte_array(name)?;
     let args = env.convert_byte_array(args)?;
-    let result = call(&name, &args)?;
+    let result = call_with_plugin_agent(&name, &args, plugin_agent as u64)?;
     let j_byte_array = env.byte_array_from_slice(&result)?;
     Ok(j_byte_array.into_raw())
 }
@@ -78,6 +95,47 @@ fn intern_load_package_with_cache(env: &mut JNIEnv, args: JByteArray) -> Result<
     Ok(j_byte_array.into_raw())
 }
 
+#[no_mangle]
+extern "C" fn plugin_agent(
+    method: *const c_char,
+    args: *const c_char,
+    kwargs: *const c_char,
+) -> *const c_char {
+    let jvm = JVM.lock().unwrap();
+    let jvm = jvm.as_ref().unwrap();
+    let mut env = jvm.attach_current_thread().unwrap();
+
+    let callback_obj = CALLBACK_OBJ.lock().unwrap();
+    let callback_obj = callback_obj.as_ref().unwrap();
+
+    let method = unsafe {
+        env.new_string(CStr::from_ptr(method).to_string_lossy().into_owned())
+            .expect("Failed to create Java string")
+    };
+    let args = unsafe {
+        env.new_string(CStr::from_ptr(args).to_string_lossy().into_owned())
+            .expect("Failed to create Java string")
+    };
+    let kwargs = unsafe {
+        env.new_string(CStr::from_ptr(kwargs).to_string_lossy().into_owned())
+            .expect("Failed to create Java string")
+    };
+    let params = &[(&method).into(), (&args).into(), (&kwargs).into()];
+    let result = env
+        .call_method(
+            callback_obj,
+            "callMethod",
+            "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;",
+            params,
+        )
+        .unwrap();
+    let result: JString = result.l().unwrap().into();
+    let result: String = env.get_string(&result).unwrap().into();
+    CString::new(result)
+        .expect("Failed to create CString")
+        .into_raw()
+}
+
 fn throw(env: &mut JNIEnv, error: anyhow::Error) -> jni::errors::Result<()> {
     env.throw(("java/lang/Exception", error.to_string()))
 }
diff --git a/java/src/main/java/com/kcl/api/API.java b/java/src/main/java/com/kcl/api/API.java
@@ -4,10 +4,13 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.UncheckedIOException;
+import java.util.Map;
 import java.nio.file.Files;
 import java.nio.file.StandardCopyOption;
 
 import com.kcl.api.Spec.*;
+import com.kcl.plugin.MethodFunction;
+import com.kcl.plugin.PluginContext;
 
 public class API implements Service {
     static String LIB_NAME = "kcl_lib_jni";
@@ -30,7 +33,7 @@ private static void doLoadLibrary() throws IOException {
             System.loadLibrary(LIB_NAME);
             return;
         } catch (UnsatisfiedLinkError ignore) {
-            // ignore - try to find native libraries from classpath
+            // ignore - try to find native libraries from class path
         }
         doLoadBundledLibrary();
     }
@@ -59,7 +62,21 @@ private static String bundledLibraryPath() {
 
     private native byte[] loadPackageWithCache(byte[] args);
 
+    private native void registerPluginContext(PluginContext ctx);
+
+    private static PluginContext pluginContext = new PluginContext();
+    private static byte[] buffer = new byte[1024];
+
+    public static void registerPlugin(String name, Map<String, MethodFunction> methodMap) {
+        pluginContext.registerPlugin(name, methodMap);
+    }
+
+    private String callMethod(String method, String argsJson, String kwArgsJson) {
+        return pluginContext.callMethod(method, argsJson, kwArgsJson);
+    }
+
     public API() {
+        registerPluginContext(pluginContext);
     }
 
     /**
diff --git a/java/src/main/java/com/kcl/plugin/MethodFunction.java b/java/src/main/java/com/kcl/plugin/MethodFunction.java
@@ -0,0 +1,8 @@
+package com.kcl.plugin;
+
+import java.util.Map;
+
+@FunctionalInterface
+public interface MethodFunction {
+    Object invoke(Object[] args, Map<String, Object> kwArgs);
+}
diff --git a/java/src/main/java/com/kcl/plugin/Plugin.java b/java/src/main/java/com/kcl/plugin/Plugin.java
@@ -0,0 +1,17 @@
+package com.kcl.plugin;
+
+import java.util.Map;
+
+public class Plugin {
+    public String name;
+    public Map<String, MethodFunction> methodMap;
+
+    public Map<String, MethodFunction> getMethodMap() {
+        return methodMap;
+    }
+
+    public Plugin(String name, Map<String, MethodFunction> methodMap) {
+        this.name = name;
+        this.methodMap = methodMap;
+    }
+}
diff --git a/java/src/main/java/com/kcl/plugin/PluginContext.java b/java/src/main/java/com/kcl/plugin/PluginContext.java
@@ -0,0 +1,71 @@
+package com.kcl.plugin;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import java.util.HashMap;
+import java.util.Map;
+
+public class PluginContext {
+    private Map<String, Plugin> pluginMap = new HashMap<>();
+    private ObjectMapper objectMapper = new ObjectMapper();
+
+    public String callMethod(String name, String argsJson, String kwArgsJson) {
+        return callJavaMethod(name, argsJson, kwArgsJson);
+    }
+
+    private String callJavaMethod(String name, String argsJson, String kwArgsJson) {
+        try {
+            return callJavaMethodUnsafe(name, argsJson, kwArgsJson);
+        } catch (Exception e) {
+            Map<String, Object> panicInfo = new HashMap<>();
+            panicInfo.put("__kcl_PanicInfo__", true);
+            panicInfo.put("message", e.getMessage());
+            return convertToJson(panicInfo);
+        }
+    }
+
+    private String callJavaMethodUnsafe(String name, String argsJson, String kwArgsJson) {
+        int dotIdx = name.lastIndexOf(".");
+        if (dotIdx < 0) {
+            return "";
+        }
+        String modulePath = name.substring(0, dotIdx);
+        String methodName = name.substring(dotIdx + 1);
+        String pluginName = modulePath.substring(modulePath.lastIndexOf(".") + 1);
+
+        MethodFunction methodFunc = this.pluginMap.getOrDefault(pluginName, new Plugin("", new HashMap<>()))
+                .getMethodMap().get(methodName);
+
+        Object[] args = convertFromJson(argsJson, Object[].class);
+        Map<String, Object> kwArgs = convertFromJson(kwArgsJson, HashMap.class);
+
+        Object result = null;
+        if (methodFunc != null) {
+            result = methodFunc.invoke(args, kwArgs);
+        }
+
+        return convertToJson(result);
+    }
+
+    public void registerPlugin(String name, Map<String, MethodFunction> methodMap) {
+        this.pluginMap.put(name, new Plugin(name, methodMap));
+    }
+
+    private String convertToJson(Object object) {
+        try {
+            return objectMapper.writeValueAsString(object);
+        } catch (JsonProcessingException e) {
+            e.printStackTrace();
+            return "";
+        }
+    }
+
+    private <T> T convertFromJson(String json, Class<T> type) {
+        try {
+            return objectMapper.readValue(json, type);
+        } catch (JsonProcessingException e) {
+            e.printStackTrace();
+            return null;
+        }
+    }
+}
diff --git a/java/src/test/java/com/kcl/PluginTest.java b/java/src/test/java/com/kcl/PluginTest.java
@@ -0,0 +1,26 @@
+package com.kcl;
+
+import com.kcl.api.API;
+import com.kcl.api.Spec.ExecProgram_Args;
+import com.kcl.api.Spec.ExecProgram_Result;
+
+import java.util.Collections;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class PluginTest {
+    @Test
+    public void testExecProgramWithPlugin() throws Exception {
+        // API instance
+        API.registerPlugin("my_plugin", Collections.singletonMap("add", (args, kwArgs) -> {
+            return (int) args[0] + (int) args[1];
+        }));
+        API api = new API();
+
+        ExecProgram_Result result = api
+                .execProgram(ExecProgram_Args.newBuilder().addKFilenameList("./src/test_data/plugin/plugin.k").build());
+
+        Assert.assertEquals(result.getYamlResult(), "result: 2");
+    }
+}
diff --git a/java/src/test_data/plugin/plugin.k b/java/src/test_data/plugin/plugin.k
@@ -0,0 +1,3 @@
+import kcl_plugin.my_plugin
+
+result = my_plugin.add(1, 1)
diff --git a/python/kcl_lib/plugin/plugin.py b/python/kcl_lib/plugin/plugin.py
@@ -21,7 +21,7 @@ def _call_py_method(self, name: str, args_json: str, kwargs_json: str) -> str:
         try:
             return self._call_py_method_unsafe(name, args_json, kwargs_json)
         except Exception as e:
-            return json.dumps({"__kcl_PanicInfo__": f"{e}"})
+            return json.dumps({"__kcl_PanicInfo__": True, "message": f"{e}"})
 
     def _call_py_method_unsafe(
         self, name: str, args_json: str, kwargs_json: str

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+import kcl_plugin.my_plugin`
	`2`	`+`
	`3`	`+result = my_plugin.add(1, 1)`