Skip to main content

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(try_blocks)]
16
17use std::alloc::System;
18
19#[global_allocator]
20static GLOBAL: System = System;
21
22mod hummock_iterator;
23use std::ffi::c_void;
24use std::ops::Deref;
25
26use anyhow::anyhow;
27use cfg_or_panic::cfg_or_panic;
28use jni::objects::JByteArray;
29use jni::sys::{JNI_VERSION_1_2, jint};
30use jni::{JNIEnv, JavaVM};
31use prost::Message;
32use risingwave_common::error::AsReport;
33use risingwave_jni_core::jvm_runtime::{jvm_env, register_java_binding_native_methods};
34use risingwave_jni_core::{
35    EnvParam, JAVA_BINDING_ASYNC_RUNTIME, JavaBindingIterator, Pointer, execute_and_catch,
36    gen_class_name, to_guarded_slice,
37};
38
39use crate::hummock_iterator::new_hummock_java_binding_iter;
40
41fn register_hummock_java_binding_native_methods(
42    env: &mut JNIEnv<'_>,
43) -> Result<(), jni::errors::Error> {
44    let binding_class = env
45        .find_class(gen_class_name!(com.risingwave.java.binding.HummockIterator))
46        .inspect_err(|e| tracing::error!(error = ?e.as_report(), "jvm find class error"))?;
47    macro_rules! gen_native_method_array {
48        () => {{
49            risingwave_jni_core::split_extract_plain_native_methods! {{long iteratorNewHummock(byte[] readPlan);}, gen_native_method_array}
50        }};
51        ({$({ $func_name:ident, {$($ret:tt)+}, {$($args:tt)*} })*}) => {
52            [
53                $(
54                    risingwave_jni_core::gen_native_method_entry! {
55                        Java_com_risingwave_java_binding_HummockIterator_, $func_name, {$($ret)+}, {$($args)*}
56                    },
57                )*
58            ]
59        }
60    }
61    env.register_native_methods(binding_class, &gen_native_method_array!())
62        .inspect_err(
63            |e| tracing::error!(error = ?e.as_report(), "jvm register native methods error"),
64        )?;
65
66    tracing::info!("register native methods for jvm successfully");
67    Ok(())
68}
69
70#[unsafe(no_mangle)]
71#[allow(non_snake_case)]
72pub extern "system" fn JNI_OnLoad(jvm: JavaVM, _reserved: *mut c_void) -> jint {
73    let result: Result<(), jni::errors::Error> = try {
74        let mut env = jvm_env(&jvm)?;
75        register_java_binding_native_methods(&mut env)?;
76        register_hummock_java_binding_native_methods(&mut env)?;
77    };
78    let _ =
79        result.inspect_err(|e| eprintln!("unable to register native method: {:?}", e.as_report()));
80
81    JNI_VERSION_1_2
82}
83
84#[cfg_or_panic(not(madsim))]
85#[unsafe(no_mangle)]
86extern "system" fn Java_com_risingwave_java_binding_HummockIterator_iteratorNewHummock<'a>(
87    env: EnvParam<'a>,
88    read_plan: JByteArray<'a>,
89) -> Pointer<'static, JavaBindingIterator<'static>> {
90    execute_and_catch(env, move |env| {
91        let read_plan = Message::decode(to_guarded_slice(&read_plan, env)?.deref())?;
92        let iter = JAVA_BINDING_ASYNC_RUNTIME
93            .block_on(new_hummock_java_binding_iter(read_plan))
94            .map_err(|e| anyhow!(e))?;
95        let iter = JavaBindingIterator::new_hummock_iter(iter);
96        Ok(iter.into())
97    })
98}