Address comments

Author: Abacn

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

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();

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

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();
     }
   }

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<Runnable> 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);