From 66bd2c90f862e5069d9117cb264958b40b449a34 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Wed, 8 Jul 2026 15:16:07 -0400 Subject: [PATCH 1/4] Support ACKNOWLEDGEMENT_MODE for JmsIO * CheckpointMark behavior in alignment with different ACKNOWLEDGEMENT_MODE * Ref count active checkpoint for quicker onClose that releases session * Fix hanging checkpoint when no incoming data in direct runner. This allows us to do an exact assert * Optimize long running unit test usign a short retry * Re-enable AMQP integration test after stuck unack messages resolved --- .../direct/UnboundedReadEvaluatorFactory.java | 13 +- .../beam/sdk/io/jms/JmsCheckpointMark.java | 110 ++++++++-- .../org/apache/beam/sdk/io/jms/JmsIO.java | 189 +++++++++++++--- .../org/apache/beam/sdk/io/jms/CommonJms.java | 27 ++- .../org/apache/beam/sdk/io/jms/JmsIOIT.java | 205 ++++++++++++------ .../org/apache/beam/sdk/io/jms/JmsIOTest.java | 188 +++++++++++++--- 6 files changed, 573 insertions(+), 159 deletions(-) diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java index e45990a760ea..c4a4fe86bf90 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java @@ -173,16 +173,17 @@ public void processElement( } else { Instant watermark = reader.getWatermark(); if (watermark.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { - // If the reader had no elements available, but the shard is not done, reuse it later - // Might be better to finalize old checkpoint. + // If the reader had no elements available, but the shard is not done, reuse it later. + // Finalize old checkpoint now. + final CheckpointMarkT checkpoint = shard.getCheckpoint(); + if (checkpoint != null) { + checkpoint.finalizeCheckpoint(); + } resultBuilder.addUnprocessedElements( Collections.>singleton( WindowedValues.timestampedValueInGlobalWindow( UnboundedSourceShard.of( - shard.getSource(), - shard.getDeduplicator(), - reader, - shard.getCheckpoint()), + shard.getSource(), shard.getDeduplicator(), reader, null), watermark))); } else { // End of input. Close the reader after finalizing old checkpoint. diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java index 276d61ae5ef5..c7bbb9730ece 100644 --- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java @@ -19,13 +19,17 @@ import java.io.IOException; import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; import java.util.Objects; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.locks.ReentrantReadWriteLock; import javax.jms.JMSException; import javax.jms.Message; import javax.jms.MessageConsumer; import javax.jms.Session; import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.io.jms.JmsIO.AcknowledgeMode; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; @@ -41,29 +45,32 @@ class JmsCheckpointMark implements UnboundedSource.CheckpointMark, Serializable private static final Logger LOG = LoggerFactory.getLogger(JmsCheckpointMark.class); private Instant oldestMessageTimestamp; - private transient @Nullable Message lastMessage; + private transient @Nullable List messages; private transient @Nullable MessageConsumer consumer; private transient @Nullable Session session; + private transient @Nullable AtomicInteger activeCheckpoints; private JmsCheckpointMark( Instant oldestMessageTimestamp, - @Nullable Message lastMessage, + @Nullable List messages, @Nullable MessageConsumer consumer, - @Nullable Session session) { + @Nullable Session session, + @Nullable AtomicInteger activeCheckpoints) { this.oldestMessageTimestamp = oldestMessageTimestamp; - this.lastMessage = lastMessage; + this.messages = messages; this.consumer = consumer; this.session = session; + this.activeCheckpoints = activeCheckpoints; } /** Acknowledge all outstanding message. */ @Override public void finalizeCheckpoint() { try { - // Jms spec will implicitly acknowledge _all_ messaged already received by the same - // session if one message in this session is being acknowledged. - if (lastMessage != null) { - lastMessage.acknowledge(); + if (messages != null) { + for (Message message : messages) { + message.acknowledge(); + } } } catch (JMSException e) { // The effect of this is message not get acknowledged and thus will be redelivered. It is @@ -93,14 +100,37 @@ public void finalizeCheckpoint() { LOG.info("Error closing JMS session. It may have already been closed."); } } + + if (activeCheckpoints != null) { + activeCheckpoints.decrementAndGet(); + } + } + + @VisibleForTesting + @Nullable + List getMessages() { + return messages; + } + + @VisibleForTesting + @Nullable + Session getSession() { + return session; + } + + @VisibleForTesting + @Nullable + MessageConsumer getConsumer() { + return consumer; } // set an empty list to messages when deserialize private void readObject(java.io.ObjectInputStream stream) throws IOException, ClassNotFoundException { stream.defaultReadObject(); - lastMessage = null; + messages = null; session = null; + consumer = null; } @Override @@ -120,8 +150,8 @@ public int hashCode() { return Objects.hash(oldestMessageTimestamp); } - static Preparer newPreparer() { - return new Preparer(); + static Preparer newPreparer(AcknowledgeMode acknowledgeMode) { + return new Preparer(acknowledgeMode); } /** @@ -129,15 +159,18 @@ static Preparer newPreparer() { */ static class Preparer { private Instant oldestMessageTimestamp = Instant.now(); - private transient @Nullable Message lastMessage = null; + private transient List messages = new ArrayList<>(); + private final AcknowledgeMode acknowledgeMode; @VisibleForTesting transient boolean discarded = false; @VisibleForTesting final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); - private Preparer() {} + private Preparer(AcknowledgeMode acknowledgeMode) { + this.acknowledgeMode = acknowledgeMode; + } - void add(Message message) throws Exception { + void add(Message message) throws JMSException { lock.writeLock().lock(); try { if (discarded) { @@ -149,7 +182,18 @@ void add(Message message) throws Exception { if (currentMessageTimestamp.isBefore(oldestMessageTimestamp)) { oldestMessageTimestamp = currentMessageTimestamp; } - lastMessage = message; + if (acknowledgeMode == AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE) { + messages.add(message); + } else { + // Jms spec will implicitly acknowledge _all_ messaged already received by the same + // session if one message in this session is being acknowledged. Only need to ack + // last seen one. + if (messages.isEmpty()) { + messages.add(message); + } else { + messages.set(0, message); + } + } } finally { lock.writeLock().unlock(); } @@ -167,6 +211,7 @@ Instant getOldestMessageTimestamp() { void discard() { lock.writeLock().lock(); try { + messages.clear(); this.discarded = true; } finally { lock.writeLock().unlock(); @@ -175,20 +220,39 @@ void discard() { /** * Create a new checkpoint mark based on the current preparer. This will reset the messages held - * by the preparer, and the owner of the preparer is responsible to create a new Jms session - * after this call. + * by the preparer. If AcknowledgeMode is CLIENT_ACKNOWLEDGE, the owner of the preparer is + * responsible to create a new Jms session after this call. */ - JmsCheckpointMark newCheckpoint(@Nullable MessageConsumer consumer, @Nullable Session session) { + JmsCheckpointMark newCheckpoint( + @Nullable MessageConsumer consumer, + @Nullable Session session, + @Nullable AcknowledgeMode acknowledgeMode, + @Nullable AtomicInteger activeCheckpoints) { JmsCheckpointMark checkpointMark; lock.writeLock().lock(); try { if (discarded) { - lastMessage = null; + messages.clear(); checkpointMark = this.emptyCheckpoint(); } else { + List messagesCopy = null; + MessageConsumer consumerToPass = null; + Session sessionToPass = null; + if (!messages.isEmpty()) { + messagesCopy = new ArrayList<>(messages); + } + if (acknowledgeMode == AcknowledgeMode.CLIENT_ACKNOWLEDGE) { + consumerToPass = consumer; + sessionToPass = session; + } checkpointMark = - new JmsCheckpointMark(oldestMessageTimestamp, lastMessage, consumer, session); - lastMessage = null; + new JmsCheckpointMark( + oldestMessageTimestamp, + messagesCopy, + consumerToPass, + sessionToPass, + activeCheckpoints); + messages.clear(); oldestMessageTimestamp = Instant.now(); } } finally { @@ -198,11 +262,11 @@ JmsCheckpointMark newCheckpoint(@Nullable MessageConsumer consumer, @Nullable Se } JmsCheckpointMark emptyCheckpoint() { - return new JmsCheckpointMark(oldestMessageTimestamp, null, null, null); + return new JmsCheckpointMark(oldestMessageTimestamp, null, null, null, null); } boolean isEmpty() { - return lastMessage == null; + return messages.isEmpty(); } } } diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java index 5deb7a2b62c7..2a37285a5778 100644 --- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java @@ -34,7 +34,7 @@ import java.util.Optional; import java.util.UUID; import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import javax.jms.Connection; import javax.jms.ConnectionFactory; @@ -118,6 +118,22 @@ * * } * + *

Acknowledgment Modes and Client Prefetch Configuration

+ * + *

By default, {@link JmsIO} consumes messages using {@link AcknowledgeMode#CLIENT_ACKNOWLEDGE} + * where a new {@link javax.jms.Session} is created for each checkpoint to prevent premature + * acknowledgments across bundles. When using {@link AcknowledgeMode#CLIENT_ACKNOWLEDGE}, if your + * JMS broker or client library utilizes client-side message prefetch buffers (such as Apache + * ActiveMQ), you should configure {@code prefetch=0} on your {@link javax.jms.ConnectionFactory} + * (e.g., via {@code ?jms.prefetchPolicy.all=0} in the broker URL or {@code + * ActiveMQPrefetchPolicy.setAll(0)}). Otherwise, unconsumed messages could be held inside old + * consumers in low throughput scenario and could lead to message backlog. + * + *

Alternatively, if your JMS broker supports individual message acknowledgment (such as ActiveMQ + * or Amazon MQ {@code ActiveMQSession.INDIVIDUAL_ACKNOWLEDGE = 4}), you can specify {@link + * Read#withAcknowledgeMode(AcknowledgeMode)} with {@link AcknowledgeMode#INDIVIDUAL_ACKNOWLEDGE}. + * In this mode, a single shared session and consumer are reused across all checkpoints. + * *

Writing to a JMS destination

* *

JmsIO sink supports writing text messages to a JMS destination on a broker. To configure a JMS @@ -146,10 +162,11 @@ public static Read read() { .setCoder(SerializableCoder.of(JmsRecord.class)) .setCloseTimeout(DEFAULT_CLOSE_TIMEOUT) .setRequiresDeduping(false) + .setAcknowledgeMode(AcknowledgeMode.CLIENT_ACKNOWLEDGE) .setMessageMapper( new MessageMapper() { @Override - public JmsRecord mapMessage(Message message) throws Exception { + public JmsRecord mapMessage(Message message) throws JMSException { TextMessage textMessage = (TextMessage) message; Map properties = new HashMap<>(); @SuppressWarnings("rawtypes") @@ -182,6 +199,7 @@ public static Read readMessage() { .setMaxNumRecords(Long.MAX_VALUE) .setCloseTimeout(DEFAULT_CLOSE_TIMEOUT) .setRequiresDeduping(false) + .setAcknowledgeMode(AcknowledgeMode.CLIENT_ACKNOWLEDGE) .build(); } @@ -261,6 +279,10 @@ ConnectionFactory getConnectionFactory() { abstract boolean isRequiresDeduping(); + abstract AcknowledgeMode getAcknowledgeMode(); + + abstract @Nullable Integer getIndividualAcknowledgeModeCode(); + abstract Builder builder(); @AutoValue.Builder @@ -292,6 +314,10 @@ abstract Builder setConnectionFactoryProviderFn( abstract Builder setRequiresDeduping(boolean requiresDeduping); + abstract Builder setAcknowledgeMode(AcknowledgeMode acknowledgeMode); + + abstract Builder setIndividualAcknowledgeModeCode(Integer individualAcknowledgeModeCode); + abstract Read build(); } @@ -484,6 +510,23 @@ public Read withRequiresDeduping() { return builder().setRequiresDeduping(true).build(); } + /** Specify the {@link AcknowledgeMode} used for consuming and acknowledging JMS messages. */ + public Read withAcknowledgeMode(AcknowledgeMode acknowledgeMode) { + checkArgument(acknowledgeMode != null, "acknowledgeMode can not be null"); + return builder().setAcknowledgeMode(acknowledgeMode).build(); + } + + /** + * Specify the custom integer code for individual message acknowledgment when using {@link + * AcknowledgeMode#INDIVIDUAL_ACKNOWLEDGE}. + * + *

Different JMS providers use different proprietary integer constants for individual + * acknowledgment (e.g., ActiveMQ uses 4, Qpid JMS / ActiveMQ Artemis / IBM MQ use 101). + */ + public Read withIndividualAcknowledgeModeCode(int individualAcknowledgeModeCode) { + return builder().setIndividualAcknowledgeModeCode(individualAcknowledgeModeCode).build(); + } + @Override public PCollection expand(PBegin input) { checkArgument( @@ -515,6 +558,7 @@ public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); builder.addIfNotNull(DisplayData.item("queue", getQueue())); builder.addIfNotNull(DisplayData.item("topic", getTopic())); + builder.add(DisplayData.item("acknowledgeMode", getAcknowledgeMode().name())); } /////////////////////////////////////////////////////////////////////////////////////// @@ -530,6 +574,35 @@ UnboundedSource createSource() { private JmsIO() {} + public enum AcknowledgeMode { + /** + * Acknowledge on checkpoint finalization with session isolation across checkpoints. Due to that + * runners may or may not finalize checkpoint in timely minor, unacked messages in unclosed + * consumer could be stuck when source throughput is zero. When using this mode in low + * throughput use cases, client-side prefetch buffers should be disabled (e.g. set {@code + * prefetch=0} for ActiveMQ) in connection factory properties. + */ + CLIENT_ACKNOWLEDGE, + + /** + * CLIENT_ACKNOWLEDGE but without session isolation across checkpoints. Acknowledging a message + * implicitly acknowledge all messages received up to that point per JMS spec. Best for + * performance but not safe on worker crash or scaling down. Good when best effort delivery is + * acceptable. + */ + CLIENT_ACKNOWLEDGE_UNSAFE, + + /** + * Acknowledge messages individually on checkpoint finalization. Recommended for JMS providers + * that support individual message acknowledgment (e.g., ActiveMQ, Amazon MQ, Artemis, Qpid + * JMS). + */ + INDIVIDUAL_ACKNOWLEDGE + + // When adding new AcknowledgeMode enum, update getAckModeCode(AcknowledgeMode mode) to handle + // the new mode. + } + /** * An interface used by {@link JmsIO.Read} for converting each jms {@link Message} into an element * of the resulting {@link PCollection}. @@ -601,10 +674,13 @@ static class UnboundedJmsReader extends UnboundedReader { private byte[] currentID; private long receiveTimeoutMillis; private PipelineOptions options; + // Acknowlging messages need open consumer. Tracking active checkpoints allows delayed close of + // session and consumer. + private final AtomicInteger activeCheckpoints = new AtomicInteger(0); public UnboundedJmsReader(UnboundedJmsSource source, PipelineOptions options) { this.source = source; - this.checkpointMarkPreparer = JmsCheckpointMark.newPreparer(); + this.checkpointMarkPreparer = JmsCheckpointMark.newPreparer(source.spec.getAcknowledgeMode()); this.currentMessage = null; this.currentID = EMPTY; this.options = options; @@ -613,7 +689,8 @@ public UnboundedJmsReader(UnboundedJmsSource source, PipelineOptions options) /** recreate session and consumer. */ private synchronized void recreateSession() throws IOException { try { - this.session = this.connection.createSession(false, Session.CLIENT_ACKNOWLEDGE); + int ackMode = getAckModeCode(source.spec.getAcknowledgeMode()); + this.session = this.connection.createSession(false, ackMode); } catch (Exception e) { throw new IOException("Error creating JMS session", e); } @@ -634,6 +711,34 @@ private synchronized void recreateSession() throws IOException { } } + private int getAckModeCode(AcknowledgeMode mode) { + if (mode == AcknowledgeMode.CLIENT_ACKNOWLEDGE + || mode == AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE) { + return Session.CLIENT_ACKNOWLEDGE; + } else if (mode == AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE) { + Integer configuredCode = source.spec.getIndividualAcknowledgeModeCode(); + if (configuredCode != null) { + return configuredCode; + } + String connectionClassName = this.connection.getClass().getName(); + if (connectionClassName.contains("org.apache.activemq.ActiveMQConnection")) { + return 4; + } else if (connectionClassName.contains("org.apache.qpid.jms")) { + return 101; + } else if (connectionClassName.contains("org.apache.activemq.artemis")) { + return 101; + } else { + throw new IllegalArgumentException( + String.format( + "Unknown JMS provider '%s' for INDIVIDUAL_ACKNOWLEDGE. " + + "Please specify the code explicitly via Read#withIndividualAcknowledgeModeCode(int).", + connectionClassName)); + } + } else { + throw new IllegalArgumentException(String.format("Unknown AcknowledgeMode: %s", mode)); + } + } + @Override public boolean start() throws IOException { Read spec = source.spec; @@ -752,16 +857,24 @@ public CheckpointMark getCheckpointMark() { MessageConsumer consumerToClose; Session sessionTofinalize; + AcknowledgeMode mode = source.spec.getAcknowledgeMode(); synchronized (this) { - consumerToClose = consumer; - sessionTofinalize = session; - } - try { - recreateSession(); - } catch (IOException e) { - throw new RuntimeException(e); + if (mode == AcknowledgeMode.CLIENT_ACKNOWLEDGE) { + consumerToClose = consumer; + sessionTofinalize = session; + try { + recreateSession(); + } catch (IOException e) { + throw new RuntimeException(e); + } + } else { + consumerToClose = null; + sessionTofinalize = null; + } } - return checkpointMarkPreparer.newCheckpoint(consumerToClose, sessionTofinalize); + activeCheckpoints.incrementAndGet(); + return checkpointMarkPreparer.newCheckpoint( + consumerToClose, sessionTofinalize, mode, activeCheckpoints); } @Override @@ -783,21 +896,43 @@ public void close() { private void doClose() { try { closeAutoscaler(); - closeConsumer(); - ScheduledExecutorService executorService = - options.as(ExecutorOptions.class).getScheduledExecutorService(); - executorService.schedule( - () -> { - LOG.debug("Closing connection after delay {}", source.spec.getCloseTimeout()); - // Discard the checkpoints and set the reader as inactive - checkpointMarkPreparer.discard(); - closeSession(); - closeConnection(); - }, - source.spec.getCloseTimeout().getMillis(), - TimeUnit.MILLISECONDS); + // Discard the checkpoints and set the reader as inactive + checkpointMarkPreparer.discard(); + if (source.spec.getAcknowledgeMode() == AcknowledgeMode.CLIENT_ACKNOWLEDGE) { + // checkpointMark holds session in CLIENT_ACKNOWLEDGE mode. Therefore + // we can close consumer and session immediately. + closeConsumer(); + closeSession(); + } + if (activeCheckpoints.get() == 0) { + closeConsumer(); + closeSession(); + closeConnection(); + } else { + ScheduledExecutorService executorService = + options.as(ExecutorOptions.class).getScheduledExecutorService(); + executorService.submit( + () -> { + long startTime = System.currentTimeMillis(); + long timeoutMillis = source.spec.getCloseTimeout().getMillis(); + while (activeCheckpoints.get() > 0 + && System.currentTimeMillis() - startTime < timeoutMillis) { + try { + Thread.sleep(1_000); // poll in 1 sec interval + } catch (InterruptedException ignored) { + break; + } + } + LOG.debug( + "Closing connection after checkpoints finalized or timeout: {}", + source.spec.getCloseTimeout()); + closeConsumer(); + closeSession(); + closeConnection(); + }); + } } catch (Exception e) { - LOG.error("Error closing reader", e); + LOG.warn("Error closing reader", e); } } @@ -809,7 +944,7 @@ private void closeConnection() { connection = null; } } catch (Exception e) { - LOG.error("Error closing connection", e); + LOG.warn("Error closing connection", e); } } diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java index 1d1245e6877d..16173a36248c 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java @@ -26,7 +26,6 @@ import javax.jms.BytesMessage; import javax.jms.ConnectionFactory; import javax.jms.Message; -import org.apache.activemq.ActiveMQConnectionFactory; import org.apache.activemq.broker.BrokerPlugin; import org.apache.activemq.broker.BrokerService; import org.apache.activemq.security.AuthenticationUser; @@ -91,12 +90,16 @@ public CommonJms( this.connectionFactoryClass = connectionFactoryClass; } + private boolean shouldAppendPort() { + return brokerPort != null && brokerPort != 0; + } + void startBroker() throws Exception { broker = new BrokerService(); broker.setUseJmx(false); broker.setPersistenceAdapter(new MemoryPersistenceAdapter()); TransportFactory.registerTransportFactory("amqp", new AmqpTransportFactory()); - if (connectionFactoryClass != ActiveMQConnectionFactory.class) { + if (shouldAppendPort()) { broker.addConnector(String.format("%s:%d?transport.transformer=jms", brokerUrl, brokerPort)); } else { broker.addConnector(brokerUrl); @@ -119,10 +122,18 @@ void startBroker() throws Exception { broker.waitUntilStarted(); } + private String getBrokerUrlWithPort() { + if (shouldAppendPort()) { + return String.format("%s:%d", brokerUrl, brokerPort); + } else { + return brokerUrl; + } + } + ConnectionFactory createConnectionFactory() throws NoSuchMethodException, InvocationTargetException, InstantiationException, IllegalAccessException { - return connectionFactoryClass.getConstructor(String.class).newInstance(brokerUrl); + return connectionFactoryClass.getConstructor(String.class).newInstance(getBrokerUrlWithPort()); } ConnectionFactory createConnectionFactoryWithSyncAcksAndWithoutPrefetch() @@ -130,13 +141,15 @@ ConnectionFactory createConnectionFactoryWithSyncAcksAndWithoutPrefetch() IllegalAccessException { return connectionFactoryClass .getConstructor(String.class) - .newInstance(brokerUrl + BROKER_WITHOUT_PREFETCH_PARAM + forceAsyncAcksParam); + .newInstance(getBrokerUrlWithPort() + BROKER_WITHOUT_PREFETCH_PARAM + forceAsyncAcksParam); } void stopBroker() throws Exception { - broker.stop(); - broker.waitUntilStopped(); - broker = null; + if (broker != null) { + broker.stop(); + broker.waitUntilStopped(); + broker = null; + } } Class getConnectionFactoryClass() { diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java index 266d04342d1f..83c3906c49e9 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java @@ -20,8 +20,8 @@ import static org.apache.beam.sdk.io.jms.CommonJms.PASSWORD; import static org.apache.beam.sdk.io.jms.CommonJms.QUEUE; import static org.apache.beam.sdk.io.jms.CommonJms.USERNAME; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; -import static org.junit.Assert.assertTrue; import java.io.IOException; import java.io.Serializable; @@ -29,8 +29,10 @@ import java.util.Collection; import java.util.Enumeration; import java.util.HashSet; +import java.util.Map; import java.util.Set; import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Function; import javax.jms.Connection; import javax.jms.ConnectionFactory; @@ -46,6 +48,7 @@ import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.common.IOITHelper; import org.apache.beam.sdk.io.common.IOTestPipelineOptions; +import org.apache.beam.sdk.io.common.NetworkTestHelper; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.Default; @@ -63,6 +66,7 @@ import org.apache.qpid.jms.JmsConnectionFactory; import org.joda.time.Duration; import org.junit.After; +import org.junit.AfterClass; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -145,84 +149,145 @@ public interface JmsIOITOptions extends IOTestPipelineOptions, StreamingOptions @Parameterized.Parameters(name = "with client class {3}") public static Collection connectionFactories() { return ImmutableList.of( - new Object[] { - "vm://localhost", 5672, "jms.sendAcksAsync=false", ActiveMQConnectionFactory.class - }); - // TODO(https://github.com/apache/beam/issues/26175) Test failure on direct runner due to - // JmsIO read on amqp slow on CI (passed locally) - // new Object[] { - // "amqp://localhost", 5672, "jms.forceAsyncAcks=false", JmsConnectionFactory.class - // }); + new Object[] {"vm://localhost", "jms.sendAcksAsync=false", ActiveMQConnectionFactory.class}, + new Object[] {"amqp://localhost", "jms.forceAsyncAcks=false", JmsConnectionFactory.class}); } - private final CommonJms commonJms; + private static final Map BROKERS = new ConcurrentHashMap<>(); + + private final String brokerUrl; + private final Integer brokerPort; + private final String forceAsyncAcksParam; + private final Class connectionFactoryClassParam; + private CommonJms commonJms; private ConnectionFactory connectionFactory; private Class connectionFactoryClass; public JmsIOIT( String brokerUrl, - Integer brokerPort, String forceAsyncAcksParam, Class connectionFactoryClass) { - this.commonJms = - new CommonJms( - OPTIONS.isLocalJmsBrokerEnabled() ? brokerUrl : OPTIONS.getJmsBrokerHost(), - OPTIONS.isLocalJmsBrokerEnabled() ? brokerPort : OPTIONS.getJmsBrokerPort(), - forceAsyncAcksParam, - connectionFactoryClass); + this.brokerUrl = brokerUrl; + if (OPTIONS.isLocalJmsBrokerEnabled()) { + try { + this.brokerPort = NetworkTestHelper.getAvailableLocalPort(); + } catch (IOException e) { + throw new RuntimeException("Failed to find available port", e); + } + } else { + this.brokerPort = OPTIONS.getJmsBrokerPort(); + } + this.forceAsyncAcksParam = forceAsyncAcksParam; + this.connectionFactoryClassParam = connectionFactoryClass; } @Before public void setup() throws Exception { if (OPTIONS.isLocalJmsBrokerEnabled()) { - this.commonJms.startBroker(); - connectionFactory = this.commonJms.createConnectionFactory(); - connectionFactoryClass = this.commonJms.getConnectionFactoryClass(); - // use a small number of record for local integration test + String key = brokerUrl + ":" + connectionFactoryClassParam.getName(); + commonJms = + BROKERS.computeIfAbsent( + key, + k -> { + CommonJms broker = + new CommonJms( + brokerUrl, brokerPort, forceAsyncAcksParam, connectionFactoryClassParam); + try { + broker.startBroker(); + } catch (Exception e) { + throw new RuntimeException(e); + } + return broker; + }); OPTIONS.setNumberOfRecords(10000); + } else { + commonJms = + new CommonJms( + OPTIONS.getJmsBrokerHost(), + OPTIONS.getJmsBrokerPort(), + forceAsyncAcksParam, + connectionFactoryClassParam); } } + @AfterClass + public static void afterClass() throws Exception { + for (CommonJms broker : BROKERS.values()) { + try { + broker.stopBroker(); + } catch (Exception e) { + // ignore errors on shutdown + } + } + BROKERS.clear(); + } + + private void setupConnection(JmsIO.AcknowledgeMode acknowledgeMode) throws Exception { + if (acknowledgeMode == JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE) { + connectionFactory = this.commonJms.createConnectionFactoryWithSyncAcksAndWithoutPrefetch(); + } else { + connectionFactory = this.commonJms.createConnectionFactory(); + } + connectionFactoryClass = this.commonJms.getConnectionFactoryClass(); + } + @After public void tearDown() throws Exception { - if (OPTIONS.isLocalJmsBrokerEnabled()) { - this.commonJms.stopBroker(); - connectionFactory = null; - connectionFactoryClass = null; - } + connectionFactory = null; + connectionFactoryClass = null; + } + + @Test + public void testPublishingThenReadingAll() throws Exception { + runPublishingThenReadingAll(JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE); + } + + @Test + public void testPublishingThenReadingAllIndividualAcknowledge() throws Exception { + runPublishingThenReadingAll(JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE); } @Test - public void testPublishingThenReadingAll() throws IOException, JMSException { - PipelineResult writeResult = publishingMessages(); + public void testPublishingThenReadingAllClientAcknowledgeUnsafe() throws Exception { + runPublishingThenReadingAll(JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE); + } + + private void runPublishingThenReadingAll(JmsIO.AcknowledgeMode acknowledgeMode) throws Exception { + setupConnection(acknowledgeMode); + String queue = QUEUE + "_" + acknowledgeMode.name(); + PipelineResult writeResult = publishingMessages(queue); PipelineResult.State writeState = writeResult.waitUntilFinish(); assertNotEquals(PipelineResult.State.FAILED, writeState); - PipelineResult readResult = readMessages(); - PipelineResult.State readState = - readResult.waitUntilFinish(Duration.standardSeconds(OPTIONS.getReadTimeout())); - // A workaround to stop the pipeline for waiting for too long + PipelineResult readResult = readMessages(acknowledgeMode, queue); + MetricsReader metricsReader = new MetricsReader(readResult, NAMESPACE); + long startTime = System.currentTimeMillis(); + long timeoutMillis = OPTIONS.getReadTimeout() * 1000L; + PipelineResult.State readState = readResult.getState(); + while (System.currentTimeMillis() - startTime < timeoutMillis + && (readState == null || !readState.isTerminal())) { + Thread.sleep(500); + readState = readResult.getState(); + if (readState != null && readState.isTerminal()) { + break; + } + long actualRecords = metricsReader.getCounterMetric(READ_ELEMENT_METRIC_NAME); + if (actualRecords >= OPTIONS.getNumberOfRecords()) { + int unackRecords = countRemain(queue); + if (unackRecords == 0) { + readResult.cancel(); + readState = readResult.getState(); + break; + } + } + } cancelIfTimeouted(readResult, readState); assertNotEquals(PipelineResult.State.FAILED, readState); - MetricsReader metricsReader = new MetricsReader(readResult, NAMESPACE); long actualRecords = metricsReader.getCounterMetric(READ_ELEMENT_METRIC_NAME); - - // TODO(yathu) resolve pending messages with direct runner then we can simply assert - // actual-records == total-records. - // Due to direct runner only finalize checkpoint at very end, there are open consumers (may - // with buffer) and O(open_consumer) message won't get delivered to other session. - int unackRecords = countRemain(QUEUE); - assertTrue( - String.format("Too many unacknowledged messages: %d", unackRecords), - unackRecords < OPTIONS.getNumberOfRecords() * 0.003); - - // acknowledged records - int ackRecords = OPTIONS.getNumberOfRecords() - unackRecords; - assertTrue( - String.format( - "actual number of records %d smaller than expected: %d.", actualRecords, ackRecords), - ackRecords <= actualRecords); + int unackRecords = countRemain(queue); + assertEquals("All messages should be acknowledged", 0, unackRecords); + assertEquals("All records should be read", (long) OPTIONS.getNumberOfRecords(), actualRecords); collectAndPublishMetrics(writeResult, readResult); } @@ -233,14 +298,22 @@ private void cancelIfTimeouted(PipelineResult readResult, PipelineResult.State r } } - private PipelineResult readMessages() { + private PipelineResult readMessages(JmsIO.AcknowledgeMode acknowledgeMode, String queue) { pipelineRead.getOptions().as(JmsIOITOptions.class).setStreaming(true); pipelineRead.getOptions().as(JmsIOITOptions.class).setBlockOnRun(false); - JmsIO.Read jmsIORead = JmsIO.readMessage(); + JmsIO.Read jmsIORead = + JmsIO.readMessage() + .withAcknowledgeMode(acknowledgeMode) + // Decrease withCloseTimeout to be smaller than pipeline runtime. Direct runner randomly + // closes reader causing cached pending consumer hanging until closeTimeout + .withCloseTimeout(Duration.standardSeconds(5)); if (pipelineRead.getOptions().as(JmsIOITOptions.class).getUseConnectionFactoryProviderFn()) { jmsIORead = jmsIORead.withConnectionFactoryProviderFn( - CommonJms.toSerializableFunction(commonJms::createConnectionFactory)); + CommonJms.toSerializableFunction( + acknowledgeMode == JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE + ? commonJms::createConnectionFactoryWithSyncAcksAndWithoutPrefetch + : commonJms::createConnectionFactory)); } else { jmsIORead = jmsIORead.withConnectionFactory(connectionFactory); } @@ -248,7 +321,7 @@ private PipelineResult readMessages() { .apply( "Read Messages", jmsIORead - .withQueue(QUEUE) + .withQueue(queue) .withUsername(USERNAME) .withPassword(PASSWORD) .withCoder(SerializableCoder.of(String.class)) @@ -258,7 +331,7 @@ private PipelineResult readMessages() { return pipelineRead.run(); } - private PipelineResult publishingMessages() { + private PipelineResult publishingMessages(String queue) { JmsIO.Write jmsIOWrite = JmsIO.write(); if (pipelineWrite.getOptions().as(JmsIOITOptions.class).getUseConnectionFactoryProviderFn()) { @@ -276,7 +349,7 @@ private PipelineResult publishingMessages() { .apply( "Publish to Jms Broker", jmsIOWrite - .withQueue(QUEUE) + .withQueue(queue) .withUsername(USERNAME) .withPassword(PASSWORD) .withValueMapper(new TextMessageMapper())); @@ -318,17 +391,19 @@ private Function getTimeMetric( } private int countRemain(String queue) throws JMSException { - Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD); - connection.start(); - Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE); - QueueBrowser browser = session.createBrowser(session.createQueue(queue)); - Enumeration messages = browser.getEnumeration(); - int count = 0; - while (messages.hasMoreElements()) { - messages.nextElement(); - count++; + try (Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD)) { + connection.start(); + try (Session session = connection.createSession(false, Session.CLIENT_ACKNOWLEDGE); + QueueBrowser browser = session.createBrowser(session.createQueue(queue))) { + Enumeration messages = browser.getEnumeration(); + int count = 0; + while (messages.hasMoreElements()) { + messages.nextElement(); + count++; + } + return count; + } } - return count; } static class ToString extends DoFn { diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java index b3233f866172..1f35cbb6f764 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java @@ -43,8 +43,6 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyLong; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -54,7 +52,6 @@ import java.io.IOException; import java.io.NotSerializableException; import java.io.Serializable; -import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Proxy; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; @@ -68,7 +65,6 @@ import java.util.List; import java.util.Set; import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import javax.jms.BytesMessage; @@ -90,6 +86,7 @@ import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.io.common.NetworkTestHelper; import org.apache.beam.sdk.io.jms.JmsIO.UnboundedJmsReader; import org.apache.beam.sdk.metrics.MetricNameFilter; import org.apache.beam.sdk.metrics.MetricQueryResults; @@ -109,6 +106,7 @@ import org.apache.qpid.jms.JmsAcknowledgeCallback; import org.apache.qpid.jms.JmsConnectionFactory; import org.apache.qpid.jms.message.JmsTextMessage; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.junit.After; import org.junit.Before; @@ -154,36 +152,39 @@ private static TextMessage createTextMessage(Session session, long messageIndex) @Parameterized.Parameters(name = "with client class {3}") public static Collection connectionFactories() { return Arrays.asList( - new Object[] { - "vm://localhost", 5672, "jms.sendAcksAsync=false", ActiveMQConnectionFactory.class - }, - new Object[] { - "amqp://localhost", 5672, "jms.forceAsyncAcks=false", JmsConnectionFactory.class - }); + new Object[] {"vm://localhost", "jms.sendAcksAsync=false", ActiveMQConnectionFactory.class}, + new Object[] {"amqp://localhost", "jms.forceAsyncAcks=false", JmsConnectionFactory.class}); } - private final CommonJms commonJms; - private final ConnectionFactory connectionFactory; + private CommonJms commonJms; + private ConnectionFactory connectionFactory; private final Class connectionFactoryClass; - private final ConnectionFactory connectionFactoryWithSyncAcksAndWithoutPrefetch; + private ConnectionFactory connectionFactoryWithSyncAcksAndWithoutPrefetch; + private final String brokerUrl; + private final Integer brokerPort; + private final String forceAsyncAcksParam; public JmsIOTest( String brokerUrl, - Integer brokerPort, String forceAsyncAcksParam, - Class connectionFactoryClass) - throws InvocationTargetException, NoSuchMethodException, InstantiationException, - IllegalAccessException { - this.commonJms = - new CommonJms(brokerUrl, brokerPort, forceAsyncAcksParam, connectionFactoryClass); + Class connectionFactoryClass) { + this.brokerUrl = brokerUrl; + this.forceAsyncAcksParam = forceAsyncAcksParam; this.connectionFactoryClass = connectionFactoryClass; - this.connectionFactory = commonJms.createConnectionFactory(); - this.connectionFactoryWithSyncAcksAndWithoutPrefetch = - commonJms.createConnectionFactoryWithSyncAcksAndWithoutPrefetch(); + try { + this.brokerPort = NetworkTestHelper.getAvailableLocalPort(); + } catch (IOException e) { + throw new RuntimeException("Failed to find available port", e); + } } @Before public void beforeEach() throws Exception { + this.commonJms = + new CommonJms(brokerUrl, brokerPort, forceAsyncAcksParam, connectionFactoryClass); + this.connectionFactory = commonJms.createConnectionFactory(); + this.connectionFactoryWithSyncAcksAndWithoutPrefetch = + commonJms.createConnectionFactoryWithSyncAcksAndWithoutPrefetch(); this.commonJms.startBroker(); } @@ -610,6 +611,12 @@ public void testCheckpointMarkAndFinalizeSeparately() throws Exception { // get checkpoint mark after consumed 4 messages CheckpointMark mark = reader.getCheckpointMark(); + JmsCheckpointMark jmsMark = (JmsCheckpointMark) mark; + // In CLIENT_ACKNOWLEDGE mode, session/consumer are recreated on checkpoint: + assertNotNull(jmsMark.getConsumer()); + assertNotNull(jmsMark.getSession()); + assertNotNull(jmsMark.getMessages()); + assertEquals(1, jmsMark.getMessages().size()); // consume two more messages after checkpoint made reader.advance(); @@ -625,7 +632,119 @@ public void testCheckpointMarkAndFinalizeSeparately() throws Exception { assertEquals(7, count(QUEUE)); } + @Test + public void testCheckpointMarkAndFinalizeSeparatelyIndividualAcknowledge() throws Exception { + UnboundedJmsReader reader = setupReaderForTest(JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE); + + assertTrue(reader.start()); + assertTrue(advanceWithRetry(reader)); + assertTrue(advanceWithRetry(reader)); + + CheckpointMark mark = reader.getCheckpointMark(); + JmsCheckpointMark jmsMark = (JmsCheckpointMark) mark; + assertNull(jmsMark.getConsumer()); + assertNull(jmsMark.getSession()); + assertNotNull(jmsMark.getMessages()); + assertEquals(3, jmsMark.getMessages().size()); + + reader.advance(); + reader.advance(); + + assertEquals(10, count(QUEUE)); + mark.finalizeCheckpoint(); + + // Verify only checkpointed messages are acknowledged + assertEquals(7, count(QUEUE)); + } + + @Test + public void testCheckpointMarkAndFinalizeSeparatelyClientAcknowledgeUnsafe() throws Exception { + UnboundedJmsReader reader = setupReaderForTest(JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE); + + assertTrue(reader.start()); + assertTrue(advanceWithRetry(reader)); + assertTrue(advanceWithRetry(reader)); + + CheckpointMark mark = reader.getCheckpointMark(); + JmsCheckpointMark jmsMark = (JmsCheckpointMark) mark; + assertNull(jmsMark.getConsumer()); + assertNull(jmsMark.getSession()); + assertNotNull(jmsMark.getMessages()); + assertEquals(1, jmsMark.getMessages().size()); + + reader.advance(); + reader.advance(); + + assertEquals(10, count(QUEUE)); + mark.finalizeCheckpoint(); + + // Verify all messages consumed on the session up to checkpoint are acknowledged + assertEquals(5, count(QUEUE)); + } + + @Test + public void testJmsCheckpointMarkIndividualAcknowledgeAllMessages() throws Exception { + Message msg1 = Mockito.mock(Message.class); + Message msg2 = Mockito.mock(Message.class); + Message msg3 = Mockito.mock(Message.class); + + JmsCheckpointMark.Preparer preparer = + JmsCheckpointMark.newPreparer(JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE); + preparer.add(msg1); + preparer.add(msg2); + preparer.add(msg3); + + AtomicInteger activeCheckpoints = new AtomicInteger(1); + JmsCheckpointMark mark = + preparer.newCheckpoint( + null, null, JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE, activeCheckpoints); + assertNotNull(mark.getMessages()); + assertEquals(3, mark.getMessages().size()); + assertNull(mark.getConsumer()); + assertNull(mark.getSession()); + assertEquals(1, activeCheckpoints.get()); + + mark.finalizeCheckpoint(); + + Mockito.verify(msg1, Mockito.times(1)).acknowledge(); + Mockito.verify(msg2, Mockito.times(1)).acknowledge(); + Mockito.verify(msg3, Mockito.times(1)).acknowledge(); + assertEquals(0, activeCheckpoints.get()); + } + + @Test + public void testJmsCheckpointMarkClientAcknowledgeUnsafeNoSessionRecreation() throws Exception { + Message msg1 = Mockito.mock(Message.class); + Message msg2 = Mockito.mock(Message.class); + + JmsCheckpointMark.Preparer preparer = + JmsCheckpointMark.newPreparer(JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE); + preparer.add(msg1); + preparer.add(msg2); + + AtomicInteger activeCheckpoints = new AtomicInteger(1); + JmsCheckpointMark mark = + preparer.newCheckpoint( + null, null, JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE, activeCheckpoints); + assertNotNull(mark.getMessages()); + assertEquals(1, mark.getMessages().size()); + assertNull(mark.getConsumer()); + assertNull(mark.getSession()); + assertEquals(1, activeCheckpoints.get()); + + mark.finalizeCheckpoint(); + + Mockito.verify(msg2, Mockito.times(1)).acknowledge(); + Mockito.verify(msg1, Mockito.never()).acknowledge(); + assertEquals(0, activeCheckpoints.get()); + } + private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException { + return setupReaderForTest(null); + } + + private JmsIO.UnboundedJmsReader setupReaderForTest( + JmsIO.@Nullable AcknowledgeMode acknowledgeMode) throws JMSException { // we are using no prefetch here // prefetch is an ActiveMQ feature: to make efficient use of network resources the broker // utilizes a 'push' model to dispatch messages to consumers. However, in the case of our @@ -652,6 +771,9 @@ private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException { .withUsername(USERNAME) .withPassword(PASSWORD) .withQueue(QUEUE); + if (acknowledgeMode != null) { + spec = spec.withAcknowledgeMode(acknowledgeMode); + } JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec); JmsIO.UnboundedJmsReader reader = source.createReader(PipelineOptionsFactory.create(), null); return reader; @@ -768,7 +890,9 @@ public void testCheckpointMarkSafety() throws Exception { /** Test the checkpoint mark default coder, which is actually AvroCoder. */ @Test public void testCheckpointMarkDefaultCoder() throws Exception { - JmsCheckpointMark jmsCheckpointMark = JmsCheckpointMark.newPreparer().newCheckpoint(null, null); + JmsCheckpointMark jmsCheckpointMark = + JmsCheckpointMark.newPreparer(JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE) + .newCheckpoint(null, null, JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE, null); Coder coder = new JmsIO.UnboundedJmsSource(null).getCheckpointMarkCoder(); CoderProperties.coderSerializable(coder); CoderProperties.coderDecodeEncodeEqual(coder, jmsCheckpointMark); @@ -819,7 +943,7 @@ public void testCustomAutoscaler() throws IOException { } @Test - public void testCloseWithTimeout() throws IOException { + public void testCloseWithTimeout() throws IOException, JMSException { Duration closeTimeout = Duration.millis(2000L); JmsIO.Read spec = JmsIO.read() @@ -836,17 +960,17 @@ public void testCloseWithTimeout() throws IOException { ExecutorOptions options = PipelineOptionsFactory.as(ExecutorOptions.class); options.setScheduledExecutorService(mockScheduledExecutorService); ArgumentCaptor runnableArgumentCaptor = ArgumentCaptor.forClass(Runnable.class); - when(mockScheduledExecutorService.schedule( - runnableArgumentCaptor.capture(), anyLong(), any(TimeUnit.class))) + when(mockScheduledExecutorService.submit(runnableArgumentCaptor.capture())) .thenReturn(null /* unused */); JmsIO.UnboundedJmsReader reader = source.createReader(options, null); reader.start(); assertFalse(getDiscardedValue(reader)); + reader.checkpointMarkPreparer.add(Mockito.mock(Message.class)); + reader.getCheckpointMark(); reader.close(); - assertFalse(getDiscardedValue(reader)); - verify(mockScheduledExecutorService) - .schedule(any(Runnable.class), eq(closeTimeout.getMillis()), eq(TimeUnit.MILLISECONDS)); + assertTrue(getDiscardedValue(reader)); + verify(mockScheduledExecutorService).submit(any(Runnable.class)); runnableArgumentCaptor.getValue().run(); assertTrue(getDiscardedValue(reader)); verifyNoMoreInteractions(mockScheduledExecutorService); @@ -982,7 +1106,8 @@ public void testWriteMessageWithRetryPolicyReachesLimit() throws Exception { int maxPublicationAttempts = 2; List data = Collections.singletonList(messageText); RetryConfiguration retryConfiguration = - RetryConfiguration.create(maxPublicationAttempts, null, null); + RetryConfiguration.create( + maxPublicationAttempts, Duration.standardSeconds(5), Duration.millis(10L)); WriteJmsResult output = pipeline @@ -1039,7 +1164,8 @@ public void testWriteMessagesWithErrors() throws Exception { List data = Arrays.asList("Message 1", "Message 2", "Message 3", "Message 4"); RetryConfiguration retryConfiguration = - RetryConfiguration.create(maxPublicationAttempts, null, null); + RetryConfiguration.create( + maxPublicationAttempts, Duration.standardSeconds(5), Duration.millis(10L)); WriteJmsResult output = pipeline From aa7f7e8e539db1d195733c3ae6e8587d69437684 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 14 Jul 2026 12:51:11 -0400 Subject: [PATCH 2/4] Address comments --- CHANGES.md | 1 + .../beam/sdk/io/jms/JmsCheckpointMark.java | 3 ++ .../org/apache/beam/sdk/io/jms/JmsIO.java | 46 +++++++++++-------- .../org/apache/beam/sdk/io/jms/JmsIOIT.java | 2 +- .../org/apache/beam/sdk/io/jms/JmsIOTest.java | 16 +++++-- 5 files changed, 42 insertions(+), 26 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 1e09ade6c389..24c735114d1e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -69,6 +69,7 @@ ## New Features / Improvements * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). +* (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java index c7bbb9730ece..3a58d28c1d93 100644 --- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java @@ -254,6 +254,9 @@ JmsCheckpointMark newCheckpoint( activeCheckpoints); messages.clear(); oldestMessageTimestamp = Instant.now(); + if (activeCheckpoints != null) { + activeCheckpoints.incrementAndGet(); + } } } finally { lock.writeLock().unlock(); diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java index 2a37285a5778..6a7d7b0ba8e3 100644 --- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java @@ -34,6 +34,7 @@ import java.util.Optional; import java.util.UUID; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import javax.jms.Connection; @@ -510,7 +511,12 @@ public Read withRequiresDeduping() { return builder().setRequiresDeduping(true).build(); } - /** Specify the {@link AcknowledgeMode} used for consuming and acknowledging JMS messages. */ + /** + * Specify the {@link AcknowledgeMode} used for consuming and acknowledging JMS messages. + * + *

To use {@link AcknowledgeMode#INDIVIDUAL_ACKNOWLEDGE}, providers other than ActiveMQ, Qpid + * JMS, ActiveMQ require configuring {@link #withIndividualAcknowledgeModeCode} explicitly. + */ public Read withAcknowledgeMode(AcknowledgeMode acknowledgeMode) { checkArgument(acknowledgeMode != null, "acknowledgeMode can not be null"); return builder().setAcknowledgeMode(acknowledgeMode).build(); @@ -521,7 +527,7 @@ public Read withAcknowledgeMode(AcknowledgeMode acknowledgeMode) { * AcknowledgeMode#INDIVIDUAL_ACKNOWLEDGE}. * *

Different JMS providers use different proprietary integer constants for individual - * acknowledgment (e.g., ActiveMQ uses 4, Qpid JMS / ActiveMQ Artemis / IBM MQ use 101). + * acknowledgment (e.g., ActiveMQ uses 4, Qpid JMS / ActiveMQ Artemis use 101). */ public Read withIndividualAcknowledgeModeCode(int individualAcknowledgeModeCode) { return builder().setIndividualAcknowledgeModeCode(individualAcknowledgeModeCode).build(); @@ -872,7 +878,6 @@ public CheckpointMark getCheckpointMark() { sessionTofinalize = null; } } - activeCheckpoints.incrementAndGet(); return checkpointMarkPreparer.newCheckpoint( consumerToClose, sessionTofinalize, mode, activeCheckpoints); } @@ -911,25 +916,26 @@ private void doClose() { } else { ScheduledExecutorService executorService = options.as(ExecutorOptions.class).getScheduledExecutorService(); - executorService.submit( - () -> { - long startTime = System.currentTimeMillis(); - long timeoutMillis = source.spec.getCloseTimeout().getMillis(); - while (activeCheckpoints.get() > 0 - && System.currentTimeMillis() - startTime < timeoutMillis) { - try { - Thread.sleep(1_000); // poll in 1 sec interval - } catch (InterruptedException ignored) { - break; + long deadline = System.currentTimeMillis() + source.spec.getCloseTimeout().getMillis(); + long pollInterval = 1L; + executorService.schedule( + new Runnable() { + @Override + public void run() { + if (activeCheckpoints.get() == 0 || System.currentTimeMillis() >= deadline) { + LOG.debug( + "Closing connection after checkpoints finalized or timeout: {}", + source.spec.getCloseTimeout()); + closeConsumer(); + closeSession(); + closeConnection(); + } else { + executorService.schedule(this, pollInterval, TimeUnit.SECONDS); } } - LOG.debug( - "Closing connection after checkpoints finalized or timeout: {}", - source.spec.getCloseTimeout()); - closeConsumer(); - closeSession(); - closeConnection(); - }); + }, + pollInterval, + TimeUnit.SECONDS); } } catch (Exception e) { LOG.warn("Error closing reader", e); diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java index 83c3906c49e9..c5650fb301f4 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java @@ -293,7 +293,7 @@ private void runPublishingThenReadingAll(JmsIO.AcknowledgeMode acknowledgeMode) private void cancelIfTimeouted(PipelineResult readResult, PipelineResult.State readState) throws IOException { - if (readState == null) { + if (readState == null || !readState.isTerminal()) { readResult.cancel(); } } diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java index 1f35cbb6f764..16737e9ef283 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java @@ -43,6 +43,8 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -65,6 +67,7 @@ import java.util.List; import java.util.Set; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import javax.jms.BytesMessage; @@ -694,7 +697,7 @@ public void testJmsCheckpointMarkIndividualAcknowledgeAllMessages() throws Excep preparer.add(msg2); preparer.add(msg3); - AtomicInteger activeCheckpoints = new AtomicInteger(1); + AtomicInteger activeCheckpoints = new AtomicInteger(0); JmsCheckpointMark mark = preparer.newCheckpoint( null, null, JmsIO.AcknowledgeMode.INDIVIDUAL_ACKNOWLEDGE, activeCheckpoints); @@ -722,7 +725,7 @@ public void testJmsCheckpointMarkClientAcknowledgeUnsafeNoSessionRecreation() th preparer.add(msg1); preparer.add(msg2); - AtomicInteger activeCheckpoints = new AtomicInteger(1); + AtomicInteger activeCheckpoints = new AtomicInteger(0); JmsCheckpointMark mark = preparer.newCheckpoint( null, null, JmsIO.AcknowledgeMode.CLIENT_ACKNOWLEDGE_UNSAFE, activeCheckpoints); @@ -960,17 +963,20 @@ public void testCloseWithTimeout() throws IOException, JMSException { ExecutorOptions options = PipelineOptionsFactory.as(ExecutorOptions.class); options.setScheduledExecutorService(mockScheduledExecutorService); ArgumentCaptor runnableArgumentCaptor = ArgumentCaptor.forClass(Runnable.class); - when(mockScheduledExecutorService.submit(runnableArgumentCaptor.capture())) + when(mockScheduledExecutorService.schedule( + runnableArgumentCaptor.capture(), anyLong(), any(TimeUnit.class))) .thenReturn(null /* unused */); JmsIO.UnboundedJmsReader reader = source.createReader(options, null); reader.start(); assertFalse(getDiscardedValue(reader)); reader.checkpointMarkPreparer.add(Mockito.mock(Message.class)); - reader.getCheckpointMark(); + CheckpointMark mark = reader.getCheckpointMark(); reader.close(); assertTrue(getDiscardedValue(reader)); - verify(mockScheduledExecutorService).submit(any(Runnable.class)); + verify(mockScheduledExecutorService) + .schedule(any(Runnable.class), eq(1L), eq(TimeUnit.SECONDS)); + mark.finalizeCheckpoint(); runnableArgumentCaptor.getValue().run(); assertTrue(getDiscardedValue(reader)); verifyNoMoreInteractions(mockScheduledExecutorService); From 60b595018121d5060c2277678c31c32eb17e2ef1 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 14 Jul 2026 13:13:29 -0400 Subject: [PATCH 3/4] Loose withCloseTimeout to fix flaky integration test; fix label --- .../src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java | 6 +++--- .../src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java index c5650fb301f4..3dbb20775f71 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java @@ -146,7 +146,7 @@ public interface JmsIOITOptions extends IOTestPipelineOptions, StreamingOptions @Rule public transient TestPipeline pipelineWrite = TestPipeline.create(); @Rule public transient TestPipeline pipelineRead = TestPipeline.create(); - @Parameterized.Parameters(name = "with client class {3}") + @Parameterized.Parameters(name = "with client class {2}") public static Collection connectionFactories() { return ImmutableList.of( new Object[] {"vm://localhost", "jms.sendAcksAsync=false", ActiveMQConnectionFactory.class}, @@ -304,9 +304,9 @@ private PipelineResult readMessages(JmsIO.AcknowledgeMode acknowledgeMode, Strin JmsIO.Read jmsIORead = JmsIO.readMessage() .withAcknowledgeMode(acknowledgeMode) - // Decrease withCloseTimeout to be smaller than pipeline runtime. Direct runner randomly + // Decrease withCloseTimeout to be smaller than pipeline timeout. Direct runner randomly // closes reader causing cached pending consumer hanging until closeTimeout - .withCloseTimeout(Duration.standardSeconds(5)); + .withCloseTimeout(Duration.standardSeconds(10)); if (pipelineRead.getOptions().as(JmsIOITOptions.class).getUseConnectionFactoryProviderFn()) { jmsIORead = jmsIORead.withConnectionFactoryProviderFn( diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java index 16737e9ef283..eb6fb4faec04 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java @@ -152,7 +152,7 @@ private static TextMessage createTextMessage(Session session, long messageIndex) RetryConfiguration.create(1, Duration.standardSeconds(1), null); @Rule public final transient TestPipeline pipeline = TestPipeline.create(); - @Parameterized.Parameters(name = "with client class {3}") + @Parameterized.Parameters(name = "with client class {2}") public static Collection connectionFactories() { return Arrays.asList( new Object[] {"vm://localhost", "jms.sendAcksAsync=false", ActiveMQConnectionFactory.class}, From 653d656d57a8d3c18fa7f98f4be15524feed4fea Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Wed, 15 Jul 2026 10:22:33 -0400 Subject: [PATCH 4/4] Fix javadoc --- .../jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java index 6a7d7b0ba8e3..8926d584a133 100644 --- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java @@ -514,8 +514,9 @@ public Read withRequiresDeduping() { /** * Specify the {@link AcknowledgeMode} used for consuming and acknowledging JMS messages. * - *

To use {@link AcknowledgeMode#INDIVIDUAL_ACKNOWLEDGE}, providers other than ActiveMQ, Qpid - * JMS, ActiveMQ require configuring {@link #withIndividualAcknowledgeModeCode} explicitly. + *

To use {@link AcknowledgeMode#INDIVIDUAL_ACKNOWLEDGE}, providers other than ActiveMQ, + * ActiveMQ Artemis, Qpid JMS, require configuring {@link #withIndividualAcknowledgeModeCode} + * explicitly. */ public Read withAcknowledgeMode(AcknowledgeMode acknowledgeMode) { checkArgument(acknowledgeMode != null, "acknowledgeMode can not be null");