[feature][connectors-v2][ActiveMQ] add blocking queue and thread pools. blocking until message consume and concurrent processing.

luops · luops · commit ec925cfc1f21 · 2025-03-20T10:05:31.000+08:00
diff --git a/seatunnel-connectors-v2/connector-activemq/src/main/java/org/apache/seatunnel/connectors/seatunnel/activemq/source/ActivemqSourceReader.java b/seatunnel-connectors-v2/connector-activemq/src/main/java/org/apache/seatunnel/connectors/seatunnel/activemq/source/ActivemqSourceReader.java
@@ -27,14 +27,12 @@
 import org.apache.seatunnel.connectors.seatunnel.activemq.exception.ActivemqConnectorErrorCode;
 import org.apache.seatunnel.connectors.seatunnel.activemq.exception.ActivemqConnectorException;
 import org.apache.seatunnel.connectors.seatunnel.activemq.split.ActivemqSplit;
-import org.apache.seatunnel.format.json.JsonDeserializationSchema;
 
 import lombok.extern.slf4j.Slf4j;
 
 import javax.jms.JMSException;
 import javax.jms.Message;
 import javax.jms.MessageConsumer;
-import javax.jms.MessageListener;
 import javax.jms.TextMessage;
 
 import java.io.IOException;
@@ -45,14 +43,20 @@
 import java.util.Set;
 import java.util.SortedMap;
 import java.util.TreeMap;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.seatunnel.connectors.seatunnel.activemq.config.ActivemqOptions.QUEUE_NAME;
 import static org.apache.seatunnel.connectors.seatunnel.activemq.config.ActivemqSourceOptions.USE_CORRELATION_ID;
 import static org.apache.seatunnel.connectors.seatunnel.activemq.exception.ActivemqConnectorErrorCode.HANDLE_SHUTDOWN_SIGNAL_FAILED;
-import static org.apache.seatunnel.connectors.seatunnel.activemq.exception.ActivemqConnectorErrorCode.MESSAGE_ACK_FAILED;
 
 @Slf4j
 public class ActivemqSourceReader<T> implements SourceReader<T, ActivemqSplit> {
+    private static final long POLL_TIMEOUT_MILLIS = 1000L;
+
     protected final Context context;
     protected final MessageConsumer consumer;
     protected transient Set<String> correlationIdsProcessedButNotAcknowledged;
@@ -62,6 +66,8 @@ public class ActivemqSourceReader<T> implements SourceReader<T, ActivemqSplit> {
     private final DeserializationSchema<SeaTunnelRow> deserializationSchema;
     private ActivemqClient activemqClient;
     private final ReadonlyConfig config;
+    private final ExecutorService executorService;
+    private final LinkedBlockingQueue<Message> messageQueue;
 
     public ActivemqSourceReader(
             DeserializationSchema<SeaTunnelRow> deserializationSchema,
@@ -74,12 +80,26 @@ public ActivemqSourceReader(
         this.config = config;
         this.activemqClient = new ActivemqClient(config);
         this.consumer = activemqClient.getConsumer();
+        this.executorService =
+                Executors.newCachedThreadPool(r -> new Thread(r, "ActiveMQ Source Data Consumer"));
+        this.messageQueue = new LinkedBlockingQueue<>();
     }
 
     @Override
     public void open() throws Exception {
         this.correlationIdsProcessedButNotAcknowledged = new HashSet<>();
         this.massageIdsProcessedForCurrentSnapshot = new ArrayList<>();
+        // start consumer listening and put messages in a queue
+        consumer.setMessageListener(
+                message -> {
+                    try {
+                        messageQueue.put(message);
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        throw new ActivemqConnectorException(
+                                ActivemqConnectorErrorCode.HANDLE_SHUTDOWN_SIGNAL_FAILED, e);
+                    }
+                });
     }
 
     @Override
@@ -97,16 +117,26 @@ public void close() throws IOException {
         if (activemqClient != null) {
             activemqClient.close();
         }
+        if (executorService != null) {
+            executorService.shutdownNow();
+        }
     }
 
     @Override
     public void pollNext(Collector output) throws Exception {
-        consumer.setMessageListener(
-                new MessageListener() {
-                    @Override
-                    public void onMessage(Message message) {
+        while (true) {
+            Message message = messageQueue.poll(POLL_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
+            if (message == null) {
+                if (Boundedness.BOUNDED.equals(context.getBoundedness())) {
+                    break;
+                }
+                continue;
+            }
+            CompletableFuture<Void> completableFuture = new CompletableFuture<>();
+            executorService.submit(
+                    () -> {
                         try {
-                            if (message != null && message instanceof TextMessage) {
+                            if (message instanceof TextMessage) {
                                 TextMessage textMessage = (TextMessage) message;
                                 String correlationId = textMessage.getJMSCorrelationID();
                                 byte[] body = textMessage.getText().getBytes();
@@ -117,29 +147,21 @@ public void onMessage(Message message) {
                                     }
                                     massageIdsProcessedForCurrentSnapshot.add(
                                             message.getJMSMessageID());
-                                    try {
-                                        if (deserializationSchema
-                                                instanceof JsonDeserializationSchema) {
-                                            ((JsonDeserializationSchema) deserializationSchema)
-                                                    .collect(body, output);
-                                        } else {
-                                            deserializationSchema.deserialize(body, output);
-                                        }
-                                    } catch (IOException e) {
-                                        log.error("Failed to deserialize message", e);
-                                    }
+                                    deserializationSchema.deserialize(body, output);
+
                                     // verify that message processing is complete
                                     textMessage.acknowledge();
                                 }
                             }
-                        } catch (JMSException e) {
-                            throw new ActivemqConnectorException(MESSAGE_ACK_FAILED, e);
                         } catch (Exception e) {
+                            completableFuture.join();
                             throw new ActivemqConnectorException(HANDLE_SHUTDOWN_SIGNAL_FAILED, e);
                         }
-                    }
-                });
-        // source support streaming mode, this is for test
+                        completableFuture.complete(null);
+                    });
+            completableFuture.join();
+        }
+
         if (Boundedness.BOUNDED.equals(context.getBoundedness())) {
             // signal to the source that we have reached the end of the data.
             context.signalNoMoreElement();
diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-activemq-e2e/pom.xml b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-activemq-e2e/pom.xml
@@ -56,5 +56,12 @@
             <artifactId>connector-activemq</artifactId>
             <version>${project.version}</version>
         </dependency>
+
+        <dependency>
+            <groupId>org.apache.seatunnel</groupId>
+            <artifactId>connector-assert</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>
diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-activemq-e2e/src/test/java/org/apache/seatunnel/e2e/connector/activemq/ActivemqIT.java b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-activemq-e2e/src/test/java/org/apache/seatunnel/e2e/connector/activemq/ActivemqIT.java
@@ -190,6 +190,13 @@ public void testActiveMQ(TestContainer container) throws Exception {
         }
     }
 
+    @TestTemplate
+    public void testSourceActiveMQJsonToConsole(TestContainer container) throws Exception {
+        initSourceData();
+        Container.ExecResult execResult = container.executeJob("/activemq-source_to_console.conf");
+        Assertions.assertEquals(0, execResult.getExitCode(), execResult.getStderr());
+    }
+
     private static Pair<SeaTunnelRowType, List<SeaTunnelRow>> generateTestDataSet() {
 
         SeaTunnelRowType rowType =
diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-activemq-e2e/src/test/resources/activemq-source_to_console.conf b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-activemq-e2e/src/test/resources/activemq-source_to_console.conf
@@ -0,0 +1,87 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+######
+###### This config file is a demonstration of batch processing in SeaTunnel config
+######
+
+env {
+  parallelism = 1
+  job.mode = "BATCH"
+}
+
+source {
+  ActiveMQ {
+    uri="tcp://activemq-host:61616"
+    username = "admin"
+    password = "admin"
+    queue_name = "sourceQueue"
+    format = json
+    plugin_output = "activemq_table"
+    schema = {
+      fields {
+        id = bigint
+        c_map = "map<string, smallint>"
+        c_array = "array<tinyint>"
+        c_string = string
+        c_boolean = boolean
+        c_tinyint = tinyint
+        c_smallint = smallint
+        c_int = int
+        c_bigint = bigint
+        c_float = float
+        c_double = double
+        c_decimal = "decimal(2, 1)"
+        c_bytes = bytes
+        c_date = date
+        c_timestamp = timestamp
+      }
+    }
+  }
+}
+
+transform {
+}
+
+sink {
+  Assert {
+    plugin_input = "activemq_table"
+    rules =
+      {
+        field_rules = [
+          {
+            field_name = id
+            field_type = bigint
+            field_value = [
+              {
+                rule_type = NOT_NULL
+              },
+              {
+                rule_type = MIN
+                rule_value = 0
+              },
+              {
+                rule_type = MAX
+                rule_value = 9
+              }
+            ]
+          }
+        ]
+      }
+
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -190,6 +190,13 @@ public void testActiveMQ(TestContainer container) throws Exception {`
`190`	`190`	`}`
`191`	`191`	`}`
`192`	`192`
	`193`	`+ @TestTemplate`
	`194`	`+ public void testSourceActiveMQJsonToConsole(TestContainer container) throws Exception {`
	`195`	`+ initSourceData();`
	`196`	`+ Container.ExecResult execResult = container.executeJob("/activemq-source_to_console.conf");`
	`197`	`+ Assertions.assertEquals(0, execResult.getExitCode(), execResult.getStderr());`
	`198`	`+ }`
	`199`	`+`
`193`	`200`	`private static Pair<SeaTunnelRowType, List<SeaTunnelRow>> generateTestDataSet() {`
`194`	`201`
`195`	`202`	`SeaTunnelRowType rowType =`