SolaceIO data loss - remove message ack from close and advance as it may lead to data loss during work rebalancing or retry.

Author: stankiewicz

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/read/UnboundedSolaceReader.java

@@ -64,7 +64,7 @@ class UnboundedSolaceReader<T> extends UnboundedReader<T> {
    * Queue to place advanced messages before {@link #getCheckpointMark()} is called. CAUTION:
    * Accessed by both reader and checkpointing threads.
    */
-  private final Queue<BytesXMLMessage> safeToAckMessages = new ConcurrentLinkedQueue<>();
+  private Queue<BytesXMLMessage> safeToAckMessages = new ConcurrentLinkedQueue<>();
 
   /**
    * Queue for messages that were ingested in the {@link #advance()} method, but not sent yet to a
@@ -136,8 +136,6 @@ public boolean start() {
 
   @Override
   public boolean advance() {
-    finalizeReadyMessages();
-
     BytesXMLMessage receivedXmlMessage;
     try {
       receivedXmlMessage = getSessionService().getReceiver().receive();
@@ -158,25 +156,10 @@ public boolean advance() {
 
   @Override
   public void close() {
-    finalizeReadyMessages();
-    sessionServiceCache.invalidate(readerUuid);
-  }
-
-  public void finalizeReadyMessages() {
-    BytesXMLMessage msg;
-    while ((msg = safeToAckMessages.poll()) != null) {
-      try {
-        msg.ackMessage();
-      } catch (IllegalStateException e) {
-        LOG.error(
-            "SolaceIO.Read: failed to acknowledge the message with applicationMessageId={}, ackMessageId={}. Returning the message to queue to retry.",
-            msg.getApplicationMessageId(),
-            msg.getAckMessageId(),
-            e);
-        safeToAckMessages.add(msg); // In case the error was transient, might succeed later
-        break; // Commit is only best effort
-      }
+    if (!safeToAckMessages.isEmpty()) {
+      LOG.error("SolaceIO.Read: reader is closing without  ack-ing messages.");
     }
+    sessionServiceCache.invalidate(readerUuid);
   }
 
   @Override
@@ -190,6 +173,7 @@ public Instant getWatermark() {
 
   @Override
   public UnboundedSource.CheckpointMark getCheckpointMark() {
+    safeToAckMessages = new ConcurrentLinkedQueue<>();
     safeToAckMessages.addAll(receivedMessages);
     receivedMessages.clear();
     return new SolaceCheckpointMark(safeToAckMessages);

sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOReadTest.java

@@ -458,13 +458,13 @@ public void testCheckpointMarkAndFinalizeSeparately() throws Exception {
     // mark all consumed messages as ready to be acknowledged
     CheckpointMark checkpointMark = reader.getCheckpointMark();
 
-    // consume 1 more message. This will call #ackMsg() on messages that were ready to be acked.
+    // consume 1 more message.
     reader.advance();
-    assertEquals(4, countAckMessages.get());
+    assertEquals(0, countAckMessages.get());
 
     // consume 1 more message. No change in the acknowledged messages.
     reader.advance();
-    assertEquals(4, countAckMessages.get());
+    assertEquals(0, countAckMessages.get());
 
     // acknowledge from the first checkpoint
     checkpointMark.finalizeCheckpoint();
@@ -473,6 +473,73 @@ public void testCheckpointMarkAndFinalizeSeparately() throws Exception {
     assertEquals(4, countAckMessages.get());
   }
 
+  @Test
+  public void testLateCheckpointOverlappingFlushingOfNextBundle() throws Exception {
+    AtomicInteger countConsumedMessages = new AtomicInteger(0);
+    AtomicInteger countAckMessages = new AtomicInteger(0);
+
+    // Broker that creates input data
+    SerializableFunction<Integer, BytesXMLMessage> recordFn =
+        index -> {
+          List<BytesXMLMessage> messages = new ArrayList<>();
+          for (int i = 0; i < 10; i++) {
+            messages.add(
+                SolaceDataUtils.getBytesXmlMessage(
+                    "payload_test" + i, "45" + i, (num) -> countAckMessages.incrementAndGet()));
+          }
+          countConsumedMessages.incrementAndGet();
+          return getOrNull(index, messages);
+        };
+
+    SessionServiceFactory fakeSessionServiceFactory =
+        MockSessionServiceFactory.builder().recordFn(recordFn).minMessagesReceived(10).build();
+
+    Read<Record> spec =
+        getDefaultRead()
+            .withSessionServiceFactory(fakeSessionServiceFactory)
+            .withMaxNumConnections(4);
+
+    UnboundedSolaceSource<Record> initialSource = getSource(spec, pipeline);
+
+    UnboundedReader<Record> reader =
+        initialSource.createReader(PipelineOptionsFactory.create(), 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(String.format("Failed at %d-th message", i), reader.advance());
+    }
+
+    // #advance() was called, but the messages were not ready to be acknowledged.
+    assertEquals(0, countAckMessages.get());
+
+    // mark all consumed messages as ready to be acknowledged
+    CheckpointMark checkpointMark = reader.getCheckpointMark();
+
+    // data is flushed
+
+    // consume 1 more message.
+    reader.advance();
+    assertEquals(0, countAckMessages.get());
+
+    // consume 1 more message. No change in the acknowledged messages.
+    reader.advance();
+    assertEquals(0, countAckMessages.get());
+
+    CheckpointMark checkpointMark2 = reader.getCheckpointMark();
+    // data is prepared for flushing that will be rejected
+
+    // acknowledge from the first checkpoint may arrive late
+    checkpointMark.finalizeCheckpoint();
+
+    assertEquals(4, countAckMessages.get());
+
+    checkpointMark2.finalizeCheckpoint();
+    assertEquals(6, countAckMessages.get());
+  }
+
   @Test
   public void testCheckpointMarkSafety() throws Exception {