Io jms fix ack message checkpoint

CHANGES.md

@@ -43,6 +43,7 @@
 
 ## Bugfixes
 
+* Fixed JmsIO acknowledgment issue (https://github.com/apache/beam/issues/20814)
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 ## Known Issues
 

sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import javax.jms.Message;
 import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 import org.slf4j.Logger;
@@ -41,13 +42,20 @@ class JmsCheckpointMark implements UnboundedSource.CheckpointMark, Serializable
   private Instant oldestMessageTimestamp = Instant.now();
   private transient List<Message> messages = new ArrayList<>();
 
-  private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+  @VisibleForTesting transient boolean discarded = false;
+
+  @VisibleForTesting final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
 
   JmsCheckpointMark() {}
 
   void add(Message message) throws Exception {
     lock.writeLock().lock();
     try {
+      if (discarded) {
+        throw new IllegalStateException(
+            String.format(
+                "Attempting to add message %s to checkpoint that is discarded.", message));
+      }
       Instant currentMessageTimestamp = new Instant(message.getJMSTimestamp());
       if (currentMessageTimestamp.isBefore(oldestMessageTimestamp)) {
         oldestMessageTimestamp = currentMessageTimestamp;
@@ -67,6 +75,15 @@ Instant getOldestMessageTimestamp() {
     }
   }
 
+  void discard() {
+    lock.writeLock().lock();
+    try {
+      this.discarded = true;
+    } finally {
+      lock.writeLock().unlock();
+    }
+  }
+
   /**
    * Acknowledge all outstanding message. Since we believe that messages will be delivered in
    * timestamp order, and acknowledged messages will not be retried, the newest message in this
@@ -76,6 +93,10 @@ Instant getOldestMessageTimestamp() {
   public void finalizeCheckpoint() {
     lock.writeLock().lock();
     try {
+      if (discarded) {
+        messages.clear();
+        return;
+      }
       for (Message message : messages) {
         try {
           message.acknowledge();
@@ -98,6 +119,7 @@ private void readObject(java.io.ObjectInputStream stream)
       throws IOException, ClassNotFoundException {
     stream.defaultReadObject();
     messages = new ArrayList<>();
+    discarded = false;
   }
 
   @Override

sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java

@@ -28,6 +28,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 import javax.jms.Connection;
 import javax.jms.ConnectionFactory;
@@ -123,11 +126,13 @@
 public class JmsIO {
 
   private static final Logger LOG = LoggerFactory.getLogger(JmsIO.class);
+  private static final Duration DEFAULT_CLOSE_TIMEOUT = Duration.millis(60000L);
 
   public static Read<JmsRecord> read() {
     return new AutoValue_JmsIO_Read.Builder<JmsRecord>()
         .setMaxNumRecords(Long.MAX_VALUE)
         .setCoder(SerializableCoder.of(JmsRecord.class))
+        .setCloseTimeout(DEFAULT_CLOSE_TIMEOUT)
         .setMessageMapper(
             (MessageMapper<JmsRecord>)
                 new MessageMapper<JmsRecord>() {
@@ -162,7 +167,10 @@ public JmsRecord mapMessage(Message message) throws Exception {
   }
 
   public static <T> Read<T> readMessage() {
-    return new AutoValue_JmsIO_Read.Builder<T>().setMaxNumRecords(Long.MAX_VALUE).build();
+    return new AutoValue_JmsIO_Read.Builder<T>()
+        .setMaxNumRecords(Long.MAX_VALUE)
+        .setCloseTimeout(DEFAULT_CLOSE_TIMEOUT)
+        .build();
   }
 
   public static <EventT> Write<EventT> write() {
@@ -206,6 +214,8 @@ public abstract static class Read<T> extends PTransform<PBegin, PCollection<T>>
 
     abstract @Nullable AutoScaler getAutoScaler();
 
+    abstract Duration getCloseTimeout();
+
     abstract Builder<T> builder();
 
     @AutoValue.Builder
@@ -230,6 +240,8 @@ abstract static class Builder<T> {
 
       abstract Builder<T> setAutoScaler(AutoScaler autoScaler);
 
+      abstract Builder<T> setCloseTimeout(Duration closeTimeout);
+
       abstract Read<T> build();
     }
 
@@ -364,6 +376,18 @@ public Read<T> withAutoScaler(AutoScaler autoScaler) {
       return builder().setAutoScaler(autoScaler).build();
     }
 
+    /**
+     * Sets the amount of time to wait for callbacks from the runner stating that the output has
+     * been durably persisted before closing the connection to the JMS broker. Any callbacks that do
+     * not occur will cause any unacknowledged messages to be returned to the JMS broker and
+     * redelivered to other clients.
+     */
+    public Read<T> withCloseTimeout(Duration closeTimeout) {
+      checkArgument(closeTimeout != null, "closeTimeout can not be null");
+      checkArgument(closeTimeout.getMillis() >= 0, "Close timeout must be non-negative.");
+      return builder().setCloseTimeout(closeTimeout).build();
+    }
+
     @Override
     public PCollection<T> expand(PBegin input) {
       checkArgument(getConnectionFactory() != null, "withConnectionFactory() is required");
@@ -446,7 +470,7 @@ public List<UnboundedJmsSource<T>> split(int desiredNumSplits, PipelineOptions o
     @Override
     public UnboundedJmsReader<T> createReader(
         PipelineOptions options, JmsCheckpointMark checkpointMark) {
-      return new UnboundedJmsReader<T>(this, checkpointMark);
+      return new UnboundedJmsReader<T>(this);
     }
 
     @Override
@@ -472,13 +496,9 @@ static class UnboundedJmsReader<T> extends UnboundedReader<T> {
     private T currentMessage;
     private Instant currentTimestamp;
 
-    public UnboundedJmsReader(UnboundedJmsSource<T> source, JmsCheckpointMark checkpointMark) {
+    public UnboundedJmsReader(UnboundedJmsSource<T> source) {
       this.source = source;
-      if (checkpointMark != null) {
-        this.checkpointMark = checkpointMark;
-      } else {
-        this.checkpointMark = new JmsCheckpointMark();
-      }
+      this.checkpointMark = new JmsCheckpointMark();
       this.currentMessage = null;
     }
 
@@ -582,29 +602,83 @@ public long getTotalBacklogBytes() {
     }
 
     @Override
-    public void close() throws IOException {
+    public void close() {
+      doClose();
+    }
+
+    @SuppressWarnings("FutureReturnValueIgnored")
+    private void doClose() {
+
       try {
-        if (consumer != null) {
-          consumer.close();
-          consumer = null;
+        closeAutoscaler();
+        closeConsumer();
+        ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
+        executorService.schedule(
+            () -> {
+              LOG.debug(
+                  "Closing session and connection after delay {}", source.spec.getCloseTimeout());
+              // Discard the checkpoints and set the reader as inactive
+              checkpointMark.discard();
+              closeSession();
+              closeConnection();
+            },
+            source.spec.getCloseTimeout().getMillis(),
+            TimeUnit.MILLISECONDS);
+
+      } catch (Exception e) {
+        LOG.error("Error closing reader", e);
+      }
+    }
+
+    private void closeConnection() {
+      try {
+        if (connection != null) {
+          connection.stop();
+          connection.close();
+          connection = null;
         }
+      } catch (Exception e) {
+        LOG.error("Error closing connection", e);
+      }
+    }
+
+    private void closeSession() {
+      try {
         if (session != null) {
           session.close();
           session = null;
         }
-        if (connection != null) {
-          connection.stop();
-          connection.close();
-          connection = null;
+      } catch (Exception e) {
+        LOG.error("Error closing session" + e.getMessage(), e);
+      }
+    }
+
+    private void closeConsumer() {
+      try {
+        if (consumer != null) {
+          consumer.close();
+          consumer = null;
         }
+      } catch (Exception e) {
+        LOG.error("Error closing consumer", e);
+      }
+    }
+
+    private void closeAutoscaler() {
+      try {
         if (autoScaler != null) {
           autoScaler.stop();
           autoScaler = null;
         }
       } catch (Exception e) {
-        throw new IOException(e);
+        LOG.error("Error closing autoscaler", e);
       }
     }
+
+    @Override
+    protected void finalize() {
+      doClose();
+    }
   }
 
   /**

sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java

@@ -21,6 +21,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -67,6 +68,7 @@
 import org.apache.beam.sdk.transforms.SerializableBiFunction;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Throwables;
+import org.joda.time.Duration;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -505,7 +507,6 @@ public void testCheckpointMarkSafety() throws Exception {
   @Test
   public void testCheckpointMarkDefaultCoder() throws Exception {
     JmsCheckpointMark jmsCheckpointMark = new JmsCheckpointMark();
-    jmsCheckpointMark.add(new ActiveMQMessage());
     Coder coder = new JmsIO.UnboundedJmsSource(null).getCheckpointMarkCoder();
     CoderProperties.coderSerializable(coder);
     CoderProperties.coderDecodeEncodeEqual(coder, jmsCheckpointMark);
@@ -555,6 +556,101 @@ public void testCustomAutoscaler() throws IOException {
     verify(autoScaler, times(1)).stop();
   }
 
+  @Test
+  public void testCloseWithTimeout() throws IOException {
+
+    Duration closeTimeout = Duration.millis(2000L);
+    JmsIO.Read spec =
+        JmsIO.read()
+            .withConnectionFactory(connectionFactory)
+            .withUsername(USERNAME)
+            .withPassword(PASSWORD)
+            .withQueue(QUEUE)
+            .withCloseTimeout(closeTimeout);
+
+    JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec);
+    JmsIO.UnboundedJmsReader reader = source.createReader(null, null);
+
+    reader.start();
+    reader.close();
+
+    boolean discarded = getDiscardedValue(reader);
+    assertFalse(discarded);
+    try {
+      Thread.sleep(closeTimeout.getMillis() + 1000);
+    } catch (InterruptedException ignored) {
+    }
+    discarded = getDiscardedValue(reader);
+    assertTrue(discarded);
+  }
+
+  private boolean getDiscardedValue(JmsIO.UnboundedJmsReader reader) {
+    JmsCheckpointMark checkpoint = (JmsCheckpointMark) reader.getCheckpointMark();
+    checkpoint.lock.readLock().lock();
+    try {
+      return checkpoint.discarded;
+    } finally {
+      checkpoint.lock.readLock().unlock();
+    }
+  }
+
+  @Test
+  public void testDiscardCheckpointMark() throws Exception {
+
+    Connection connection =
+        connectionFactoryWithSyncAcksAndWithoutPrefetch.createConnection(USERNAME, PASSWORD);
+    connection.start();
+    Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
+    MessageProducer producer = session.createProducer(session.createQueue(QUEUE));
+    for (int i = 0; i < 10; i++) {
+      producer.send(session.createTextMessage("test " + i));
+    }
+    producer.close();
+    session.close();
+    connection.close();
+
+    JmsIO.Read spec =
+        JmsIO.read()
+            .withConnectionFactory(connectionFactoryWithSyncAcksAndWithoutPrefetch)
+            .withUsername(USERNAME)
+            .withPassword(PASSWORD)
+            .withQueue(QUEUE);
+    JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec);
+    JmsIO.UnboundedJmsReader reader = source.createReader(null, null);
+
+    // start the reader and move to the first record
+    assertTrue(reader.start());
+
+    // consume 3 messages (NB: start already consumed the first message)
+    for (int i = 0; i < 3; i++) {
+      assertTrue(reader.advance());
+    }
+
+    // the messages are still pending in the queue (no ACK yet)
+    assertEquals(10, count(QUEUE));
+
+    // we finalize the checkpoint
+    reader.getCheckpointMark().finalizeCheckpoint();
+
+    // the checkpoint finalize ack the messages, and so they are not pending in the queue anymore
+    assertEquals(6, count(QUEUE));
+
+    // we read the 6 pending messages
+    for (int i = 0; i < 6; i++) {
+      assertTrue(reader.advance());
+    }
+
+    // still 6 pending messages as we didn't finalize the checkpoint
+    assertEquals(6, count(QUEUE));
+
+    // But here we discard the checkpoint
+    ((JmsCheckpointMark) reader.getCheckpointMark()).discard();
+    // we finalize the checkpoint: no more message in the queue
+    reader.getCheckpointMark().finalizeCheckpoint();
+
+    assertEquals(6, count(QUEUE));
+  }
+
   private int count(String queue) throws Exception {
     Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD);
     connection.start();