Fix MqttIO read checkpoint logic (apache#36056)

* Fix MqttIO read checkpoint logic

* add tests

Author: Abacn

sdks/java/io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java

@@ -422,20 +422,16 @@ public void populateDisplayData(DisplayData.Builder builder) {
   static class MqttCheckpointMark implements UnboundedSource.CheckpointMark, Serializable {
 
     @VisibleForTesting String clientId;
-    @VisibleForTesting Instant oldestMessageTimestamp = Instant.now();
     @VisibleForTesting transient List<Message> messages = new ArrayList<>();
 
-    public MqttCheckpointMark() {}
-
-    public MqttCheckpointMark(String id) {
-      clientId = id;
+    public MqttCheckpointMark(String id, List<Message> messages) {
+      this.clientId = id;
+      this.messages = messages;
     }
 
-    public void add(Message message, Instant timestamp) {
-      if (timestamp.isBefore(oldestMessageTimestamp)) {
-        oldestMessageTimestamp = timestamp;
-      }
-      messages.add(message);
+    @VisibleForTesting
+    MqttCheckpointMark(String id) {
+      this.clientId = id;
     }
 
     @Override
@@ -448,7 +444,6 @@ public void finalizeCheckpoint() {
           LOG.warn("Can't ack message for client ID {}", clientId, e);
         }
       }
-      oldestMessageTimestamp = Instant.now();
       messages.clear();
     }
 
@@ -464,7 +459,6 @@ public boolean equals(@Nullable Object other) {
       if (other instanceof MqttCheckpointMark) {
         MqttCheckpointMark that = (MqttCheckpointMark) other;
         return Objects.equals(this.clientId, that.clientId)
-            && Objects.equals(this.oldestMessageTimestamp, that.oldestMessageTimestamp)
             && Objects.deepEquals(this.messages, that.messages);
       } else {
         return false;
@@ -473,7 +467,38 @@ public boolean equals(@Nullable Object other) {
 
     @Override
     public int hashCode() {
-      return Objects.hash(clientId, oldestMessageTimestamp, messages);
+      return Objects.hash(clientId, messages);
+    }
+
+    static class Preparer {
+      @VisibleForTesting String clientId;
+      @VisibleForTesting Instant oldestMessageTimestamp = Instant.now();
+      @VisibleForTesting transient List<Message> messages = new ArrayList<>();
+
+      public Preparer(MqttCheckpointMark checkpointMark) {
+        clientId = checkpointMark.clientId;
+        messages = checkpointMark.messages;
+      }
+
+      public Preparer(String id) {
+        clientId = id;
+      }
+
+      public Preparer() {}
+
+      public void add(Message message, Instant timestamp) {
+        if (timestamp.isBefore(oldestMessageTimestamp)) {
+          oldestMessageTimestamp = timestamp;
+        }
+        messages.add(message);
+      }
+
+      MqttCheckpointMark newCheckpoint() {
+        List<Message> currentMessages = messages;
+        messages = new ArrayList<>();
+        oldestMessageTimestamp = Instant.now();
+        return new MqttCheckpointMark(clientId, currentMessages);
+      }
     }
   }
 
@@ -489,16 +514,20 @@ public UnboundedMqttSource(Read<T> spec) {
     @Override
     @SuppressWarnings("unchecked")
     public UnboundedReader<T> createReader(
-        PipelineOptions options, MqttCheckpointMark checkpointMark) {
+        PipelineOptions options, @Nullable MqttCheckpointMark checkpointMark) {
       final UnboundedMqttReader<T> unboundedMqttReader;
+      MqttCheckpointMark.Preparer preparer =
+          checkpointMark == null
+              ? new MqttCheckpointMark.Preparer()
+              : new MqttCheckpointMark.Preparer(checkpointMark);
       if (spec.withMetadata()) {
         unboundedMqttReader =
             new UnboundedMqttReader<>(
                 this,
-                checkpointMark,
+                preparer,
                 message -> (T) MqttRecord.of(message.getTopic(), message.getPayload()));
       } else {
-        unboundedMqttReader = new UnboundedMqttReader<>(this, checkpointMark);
+        unboundedMqttReader = new UnboundedMqttReader<>(this, preparer);
       }
 
       return unboundedMqttReader;
@@ -538,25 +567,26 @@ static class UnboundedMqttReader<T> extends UnboundedSource.UnboundedReader<T> {
     private BlockingConnection connection;
     private T current;
     private Instant currentTimestamp;
-    private MqttCheckpointMark checkpointMark;
+    private final MqttCheckpointMark.Preparer checkpointPreparer;
     private SerializableFunction<Message, T> extractFn;
 
-    public UnboundedMqttReader(UnboundedMqttSource<T> source, MqttCheckpointMark checkpointMark) {
+    public UnboundedMqttReader(
+        UnboundedMqttSource<T> source, MqttCheckpointMark.Preparer checkpointPreparer) {
       this.source = source;
       this.current = null;
-      if (checkpointMark != null) {
-        this.checkpointMark = checkpointMark;
+      if (checkpointPreparer != null) {
+        this.checkpointPreparer = checkpointPreparer;
       } else {
-        this.checkpointMark = new MqttCheckpointMark();
+        this.checkpointPreparer = new MqttCheckpointMark.Preparer();
       }
       this.extractFn = message -> (T) message.getPayload();
     }
 
     public UnboundedMqttReader(
         UnboundedMqttSource<T> source,
-        MqttCheckpointMark checkpointMark,
+        MqttCheckpointMark.Preparer checkpointPreparer,
         SerializableFunction<Message, T> extractFn) {
-      this(source, checkpointMark);
+      this(source, checkpointPreparer);
       this.extractFn = extractFn;
     }
 
@@ -567,7 +597,7 @@ public boolean start() throws IOException {
       try {
         client = spec.connectionConfiguration().createClient();
         LOG.debug("Reader client ID is {}", client.getClientId());
-        checkpointMark.clientId = client.getClientId().toString();
+        checkpointPreparer.clientId = client.getClientId().toString();
         connection = createConnection(client);
         connection.subscribe(
             new Topic[] {new Topic(spec.connectionConfiguration().getTopic(), QoS.AT_LEAST_ONCE)});
@@ -587,7 +617,7 @@ public boolean advance() throws IOException {
         }
         current = this.extractFn.apply(message);
         currentTimestamp = Instant.now();
-        checkpointMark.add(message, currentTimestamp);
+        checkpointPreparer.add(message, currentTimestamp);
       } catch (Exception e) {
         throw new IOException(e);
       }
@@ -608,12 +638,12 @@ public void close() throws IOException {
 
     @Override
     public Instant getWatermark() {
-      return checkpointMark.oldestMessageTimestamp;
+      return checkpointPreparer.oldestMessageTimestamp;
     }
 
     @Override
     public UnboundedSource.CheckpointMark getCheckpointMark() {
-      return checkpointMark;
+      return checkpointPreparer.newCheckpoint();
     }
 
     @Override

sdks/java/io/mqtt/src/test/java/org/apache/beam/sdk/io/mqtt/MqttIOTest.java

@@ -27,6 +27,7 @@
 import java.io.ObjectOutputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -50,11 +51,13 @@
 import org.apache.beam.sdk.values.PCollection;
 import org.fusesource.hawtbuf.Buffer;
 import org.fusesource.mqtt.client.BlockingConnection;
+import org.fusesource.mqtt.client.Callback;
 import org.fusesource.mqtt.client.MQTT;
 import org.fusesource.mqtt.client.Message;
 import org.fusesource.mqtt.client.QoS;
 import org.fusesource.mqtt.client.Topic;
 import org.joda.time.Duration;
+import org.joda.time.Instant;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
@@ -286,6 +289,61 @@ public void testReceiveWithTimeoutAndNoData() throws Exception {
     pipeline.run();
   }
 
+  private static class FakeMessage extends Message {
+
+    private int ackCount;
+
+    public FakeMessage() {
+      super(null, null, null, null);
+      this.ackCount = 0;
+    }
+
+    @Override
+    public void ack() {
+      ++ackCount;
+    }
+
+    @Override
+    public void ack(final Callback<Void> unused) {
+      ++ackCount;
+    }
+
+    public int getAckCount() {
+      return ackCount;
+    }
+  }
+
+  @Test
+  public void testReadCheckpoint() {
+    MqttIO.MqttCheckpointMark.Preparer preparer = new MqttIO.MqttCheckpointMark.Preparer("id");
+    ArrayList<Message> messages = new ArrayList<>();
+    for (int i = 0; i < 5; ++i) {
+      messages.add(new FakeMessage());
+    }
+    preparer.add(messages.get(0), Instant.ofEpochMilli(20));
+    preparer.add(messages.get(1), Instant.ofEpochMilli(10));
+    preparer.add(messages.get(2), Instant.ofEpochMilli(30));
+    assertEquals(Instant.ofEpochMilli(10), preparer.oldestMessageTimestamp);
+    MqttIO.MqttCheckpointMark checkpointA = preparer.newCheckpoint();
+    preparer.add(messages.get(3), Instant.ofEpochMilli(40));
+    preparer.add(messages.get(4), Instant.ofEpochMilli(50));
+    MqttIO.MqttCheckpointMark checkpointB = preparer.newCheckpoint();
+    assertTrue(
+        Arrays.stream(messages.toArray()).allMatch((m -> ((FakeMessage) m).getAckCount() == 0)));
+    checkpointA.finalizeCheckpoint();
+    // only messages in finalized checkpoint acked
+    assertTrue(
+        Arrays.stream(messages.subList(0, 3).toArray())
+            .allMatch((m -> ((FakeMessage) m).getAckCount() == 1)));
+    assertTrue(
+        Arrays.stream(messages.subList(3, 5).toArray())
+            .allMatch((m -> ((FakeMessage) m).getAckCount() == 0)));
+    checkpointB.finalizeCheckpoint();
+    // all messaged acked once
+    assertTrue(
+        Arrays.stream(messages.toArray()).allMatch((m -> ((FakeMessage) m).getAckCount() == 1)));
+  }
+
   @Test
   public void testWrite() throws Exception {
     final int numberOfTestMessages = 200;
@@ -560,7 +618,6 @@ public void testReadObject() throws Exception {
     // the number of messages of the decoded checkpoint should be zero
     assertEquals(0, cp2.messages.size());
     assertEquals(cp1.clientId, cp2.clientId);
-    assertEquals(cp1.oldestMessageTimestamp, cp2.oldestMessageTimestamp);
   }
 
   /**