add cdcJniChannelClose

risingwavelabs · Sep 8, 2023 · 67f0e72 · 67f0e72
1 parent 547941b
commit 67f0e72
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 6 deletions.
diff --git a/java/com_risingwave_java_binding_Binding.h b/java/com_risingwave_java_binding_Binding.h
diff --git a/...ector-service/src/main/java/com/risingwave/connector/source/core/JniDbzSourceHandler.java b/...ector-service/src/main/java/com/risingwave/connector/source/core/JniDbzSourceHandler.java
@@ -17,6 +17,7 @@
 import com.risingwave.connector.api.source.SourceTypeE;
 import com.risingwave.connector.source.common.DbzConnectorConfig;
 import com.risingwave.java.binding.Binding;
+import com.risingwave.java.binding.CdcJniChannel;
 import com.risingwave.metrics.ConnectorNodeMetrics;
 import com.risingwave.proto.ConnectorServiceProto;
 import java.util.HashMap;
@@ -64,7 +65,7 @@ public void start(long channelPtr) {
             return;
         }
 
-        try {
+        try (CdcJniChannel channel = new CdcJniChannel(channelPtr)) {
             // Start the engine
             runner.start();
             LOG.info("Start consuming events of table {}", config.getSourceId());
@@ -83,10 +84,12 @@ public void start(long channelPtr) {
                             "Engine#{}: emit one chunk {} events to network ",
                             config.getSourceId(),
                             resp.getEventsCount());
-                    success = Binding.sendCdcSourceMsgToChannel(channelPtr, resp.toByteArray());
+                    success =
+                            Binding.sendCdcSourceMsgToChannel(
+                                    channel.getPointer(), resp.toByteArray());
                 } else {
                     // If resp is null means just check whether channel is closed.
-                    success = Binding.sendCdcSourceMsgToChannel(channelPtr, null);
+                    success = Binding.sendCdcSourceMsgToChannel(channel.getPointer(), null);
                 }
                 if (!success) {
                     LOG.info(

diff --git a/java/java-binding/src/main/java/com/risingwave/java/binding/Binding.java b/java/java-binding/src/main/java/com/risingwave/java/binding/Binding.java
@@ -91,4 +91,6 @@ public class Binding {
     static native long streamChunkIteratorFromPretty(String str);
 
     public static native boolean sendCdcSourceMsgToChannel(long channelPtr, byte[] msg);
+
+    static native void cdcJniChannelClose(long pointer);
 }
diff --git a/java/java-binding/src/main/java/com/risingwave/java/binding/CdcJniChannel.java b/java/java-binding/src/main/java/com/risingwave/java/binding/CdcJniChannel.java
@@ -0,0 +1,23 @@
+package com.risingwave.java.binding;
+
+public class CdcJniChannel implements AutoCloseable {
+    private final long pointer;
+    private boolean isClosed;
+
+    public CdcJniChannel(long pointer) {
+        this.pointer = pointer;
+        this.isClosed = false;
+    }
+
+    public long getPointer() {
+        return this.pointer;
+    }
+
+    @Override
+    public void close() {
+        if (!isClosed) {
+            isClosed = true;
+            Binding.cdcJniChannelClose(pointer);
+        }
+    }
+}
diff --git a/src/jni_core/src/jvm_runtime.rs b/src/jni_core/src/jvm_runtime.rs
@@ -275,6 +275,12 @@ fn register_native_method_for_jvm(jvm: &JavaVM) {
                 fn_ptr: crate::Java_com_risingwave_java_binding_Binding_sendCdcSourceMsgToChannel
                     as *mut c_void,
             },
+            NativeMethod {
+                name: JNIString::from("cdcJniChannelClose"),
+                sig: JNIString::from("(J)V"),
+                fn_ptr: crate::Java_com_risingwave_java_binding_Binding_cdcJniChannelClose
+                    as *mut c_void,
+            },
         ],
     )
     .inspect_err(|e| tracing::error!("jvm register native methods error: {:?}", e))

diff --git a/src/jni_core/src/lib.rs b/src/jni_core/src/lib.rs
@@ -842,8 +842,6 @@ pub extern "system" fn Java_com_risingwave_java_binding_Binding_sendCdcSourceMsg
         // If msg is null means just check whether channel is closed.
         if msg.is_null() {
             if channel.as_ref().is_closed() {
-                // Drop channel as well.
-                channel.drop();
                 return Ok(JNI_FALSE);
             } else {
                 return Ok(JNI_TRUE);
@@ -860,14 +858,21 @@ pub extern "system" fn Java_com_risingwave_java_binding_Binding_sendCdcSourceMsg
                 Ok(JNI_TRUE)
             }
             Err(e) => {
-                channel.drop();
                 tracing::debug!("send error.  {:?}", e);
                 Ok(JNI_FALSE)
             }
         }
     })
 }
 
+#[no_mangle]
+pub extern "system" fn Java_com_risingwave_java_binding_Binding_cdcJniChannelClose<'a>(
+    _env: EnvParam<'a>,
+    pointer: Pointer<'a, GetEventStreamJniSender>,
+) {
+    pointer.drop()
+}
+
 #[cfg(test)]
 mod tests {
     use risingwave_common::types::{DataType, Timestamptz};