risingwave_java_binding/
lib.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#![feature(type_alias_impl_trait)]
16#![feature(try_blocks)]
17
18use std::alloc::System;
19
20#[global_allocator]
21static GLOBAL: System = System;
22
23mod hummock_iterator;
24use std::ffi::c_void;
25use std::ops::Deref;
26
27use anyhow::anyhow;
28use cfg_or_panic::cfg_or_panic;
29use jni::objects::JByteArray;
30use jni::sys::{JNI_VERSION_1_2, jint};
31use jni::{JNIEnv, JavaVM};
32use prost::Message;
33use risingwave_common::error::AsReport;
34use risingwave_jni_core::jvm_runtime::{jvm_env, register_java_binding_native_methods};
35use risingwave_jni_core::{
36    EnvParam, JAVA_BINDING_ASYNC_RUNTIME, JavaBindingIterator, Pointer, execute_and_catch,
37    gen_class_name, to_guarded_slice,
38};
39
40use crate::hummock_iterator::new_hummock_java_binding_iter;
41
42fn register_hummock_java_binding_native_methods(
43    env: &mut JNIEnv<'_>,
44) -> Result<(), jni::errors::Error> {
45    let binding_class = env
46        .find_class(gen_class_name!(com.risingwave.java.binding.HummockIterator))
47        .inspect_err(|e| tracing::error!(error = ?e.as_report(), "jvm find class error"))?;
48    macro_rules! gen_native_method_array {
49        () => {{
50            risingwave_jni_core::split_extract_plain_native_methods! {{long iteratorNewHummock(byte[] readPlan);}, gen_native_method_array}
51        }};
52        ({$({ $func_name:ident, {$($ret:tt)+}, {$($args:tt)*} })*}) => {
53            [
54                $(
55                    risingwave_jni_core::gen_native_method_entry! {
56                        Java_com_risingwave_java_binding_HummockIterator_, $func_name, {$($ret)+}, {$($args)*}
57                    },
58                )*
59            ]
60        }
61    }
62    env.register_native_methods(binding_class, &gen_native_method_array!())
63        .inspect_err(
64            |e| tracing::error!(error = ?e.as_report(), "jvm register native methods error"),
65        )?;
66
67    tracing::info!("register native methods for jvm successfully");
68    Ok(())
69}
70
71#[unsafe(no_mangle)]
72#[allow(non_snake_case)]
73pub extern "system" fn JNI_OnLoad(jvm: JavaVM, _reserved: *mut c_void) -> jint {
74    let result: Result<(), jni::errors::Error> = try {
75        let mut env = jvm_env(&jvm)?;
76        register_java_binding_native_methods(&mut env)?;
77        register_hummock_java_binding_native_methods(&mut env)?;
78    };
79    let _ =
80        result.inspect_err(|e| eprintln!("unable to register native method: {:?}", e.as_report()));
81
82    JNI_VERSION_1_2
83}
84
85#[cfg_or_panic(not(madsim))]
86#[unsafe(no_mangle)]
87extern "system" fn Java_com_risingwave_java_binding_HummockIterator_iteratorNewHummock<'a>(
88    env: EnvParam<'a>,
89    read_plan: JByteArray<'a>,
90) -> Pointer<'static, JavaBindingIterator<'static>> {
91    execute_and_catch(env, move |env| {
92        let read_plan = Message::decode(to_guarded_slice(&read_plan, env)?.deref())?;
93        let iter = JAVA_BINDING_ASYNC_RUNTIME
94            .block_on(new_hummock_java_binding_iter(read_plan))
95            .map_err(|e| anyhow!(e))?;
96        let iter = JavaBindingIterator::new_hummock_iter(iter);
97        Ok(iter.into())
98    })
99}