SolaceIO data loss

- remove message ack from close and advance as it may lead to data loss during work rebalancing or retry.
- add async NACK with configurable deadline so any not finalized messages are rejected and retried.

Author: stankiewicz

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

@@ -412,6 +412,7 @@ public class SolaceIO {
         }
       };
   private static final boolean DEFAULT_DEDUPLICATE_RECORDS = false;
+  private static final int DEFAULT_ACK_DEADLINE_SECONDS = 60;
   private static final Duration DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD =
       Duration.standardSeconds(30);
   public static final int DEFAULT_WRITER_NUM_SHARDS = 20;
@@ -461,6 +462,7 @@ public static Read<Solace.Record> read() {
             .setParseFn(SolaceRecordMapper::map)
             .setTimestampFn(SENDER_TIMESTAMP_FUNCTION)
             .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS)
+            .setAckDeadlineSeconds(DEFAULT_ACK_DEADLINE_SECONDS)
             .setWatermarkIdleDurationThreshold(DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD));
   }
 
@@ -490,6 +492,7 @@ public static <T> Read<T> read(
             .setParseFn(parseFn)
             .setTimestampFn(timestampFn)
             .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS)
+            .setAckDeadlineSeconds(DEFAULT_ACK_DEADLINE_SECONDS)
             .setWatermarkIdleDurationThreshold(DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD));
   }
 
@@ -587,6 +590,16 @@ public Read<T> withDeduplicateRecords(boolean deduplicateRecords) {
       return this;
     }
 
+    /**
+     * Optional, default: 60. Set to ack deadline after which {@link
+     * org.apache.beam.sdk.io.solace.read.UnboundedSolaceReader} will start to reject outstanding
+     * messages that were not successfully checkpointed.
+     */
+    public Read<T> withAckDeadlineSeconds(int ackDeadlineSeconds) {
+      configurationBuilder.setAckDeadlineSeconds(ackDeadlineSeconds);
+      return this;
+    }
+
     /**
      * Set a factory that creates a {@link org.apache.beam.sdk.io.solace.broker.SempClientFactory}.
      *
@@ -689,6 +702,8 @@ abstract static class Configuration<T> {
 
       abstract Duration getWatermarkIdleDurationThreshold();
 
+      abstract int getAckDeadlineSeconds();
+
       public static <T> Builder<T> builder() {
         Builder<T> builder =
             new org.apache.beam.sdk.io.solace.AutoValue_SolaceIO_Read_Configuration.Builder<T>();
@@ -719,6 +734,8 @@ abstract Builder<T> setParseFn(
 
         abstract Builder<T> setWatermarkIdleDurationThreshold(Duration idleDurationThreshold);
 
+        abstract Builder<T> setAckDeadlineSeconds(int seconds);
+
         abstract Configuration<T> build();
       }
     }
@@ -756,7 +773,8 @@ public PCollection<T> expand(PBegin input) {
                   coder,
                   configuration.getTimestampFn(),
                   configuration.getWatermarkIdleDurationThreshold(),
-                  configuration.getParseFn())));
+                  configuration.getParseFn(),
+                  configuration.getAckDeadlineSeconds())));
     }
 
     @VisibleForTesting

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/JcsmpSessionService.java

@@ -29,6 +29,7 @@
 import com.solacesystems.jcsmp.JCSMPProperties;
 import com.solacesystems.jcsmp.JCSMPSession;
 import com.solacesystems.jcsmp.Queue;
+import com.solacesystems.jcsmp.XMLMessage;
 import com.solacesystems.jcsmp.XMLMessageProducer;
 import java.io.IOException;
 import java.util.Objects;
@@ -143,6 +144,8 @@ private MessageReceiver createFlowReceiver() throws JCSMPException, IOException
 
     ConsumerFlowProperties flowProperties = new ConsumerFlowProperties();
     flowProperties.setEndpoint(queue);
+    flowProperties.addRequiredSettlementOutcomes(
+        XMLMessage.Outcome.FAILED, XMLMessage.Outcome.REJECTED);
     flowProperties.setAckMode(JCSMPProperties.SUPPORTED_MESSAGE_ACK_CLIENT);
 
     EndpointProperties endpointProperties = new EndpointProperties();

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

@@ -20,6 +20,8 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 
 import com.solacesystems.jcsmp.BytesXMLMessage;
+import com.solacesystems.jcsmp.JCSMPException;
+import com.solacesystems.jcsmp.XMLMessage;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
@@ -31,7 +33,9 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
 import org.apache.beam.sdk.io.solace.broker.SempClient;
@@ -59,12 +63,8 @@ class UnboundedSolaceReader<T> extends UnboundedReader<T> {
   private final SessionServiceFactory sessionServiceFactory;
   private @Nullable BytesXMLMessage solaceOriginalRecord;
   private @Nullable T solaceMappedRecord;
-
-  /**
-   * 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 @Nullable ScheduledFuture<?> nackCallback = null;
+  private final int ackDeadlineSeconds;
 
   /**
    * Queue for messages that were ingested in the {@link #advance()} method, but not sent yet to a
@@ -109,6 +109,7 @@ public UnboundedSolaceReader(UnboundedSolaceSource<T> currentSource) {
     this.sessionServiceFactory = currentSource.getSessionServiceFactory();
     this.sempClient = currentSource.getSempClientFactory().create();
     this.readerUuid = UUID.randomUUID();
+    this.ackDeadlineSeconds = currentSource.getAckDeadlineSeconds();
   }
 
   private SessionService getSessionService() {
@@ -136,8 +137,6 @@ public boolean start() {
 
   @Override
   public boolean advance() {
-    finalizeReadyMessages();
-
     BytesXMLMessage receivedXmlMessage;
     try {
       receivedXmlMessage = getSessionService().getReceiver().receive();
@@ -158,23 +157,28 @@ public boolean advance() {
 
   @Override
   public void close() {
-    finalizeReadyMessages();
+    try {
+      if (nackCallback != null) {
+        // wait only for last one to finish, it will mean all the previous one are also done.
+        nackCallback.get(ackDeadlineSeconds * 2, TimeUnit.SECONDS);
+      }
+    } catch (InterruptedException | ExecutionException | TimeoutException e) {
+      LOG.error("SolaceIO.Read: Failed to wait till nack background thread is finished");
+    }
     sessionServiceCache.invalidate(readerUuid);
   }
 
-  public void finalizeReadyMessages() {
+  public void nackMessages(Queue<BytesXMLMessage> checkpoint) {
     BytesXMLMessage msg;
-    while ((msg = safeToAckMessages.poll()) != null) {
+    while ((msg = checkpoint.poll()) != null) {
       try {
-        msg.ackMessage();
-      } catch (IllegalStateException e) {
+        msg.settle(XMLMessage.Outcome.FAILED);
+      } catch (IllegalStateException | JCSMPException e) {
         LOG.error(
-            "SolaceIO.Read: failed to acknowledge the message with applicationMessageId={}, ackMessageId={}. Returning the message to queue to retry.",
+            "SolaceIO.Read: failed to nack the message with applicationMessageId={}, ackMessageId={}.",
             msg.getApplicationMessageId(),
             msg.getAckMessageId(),
             e);
-        safeToAckMessages.add(msg); // In case the error was transient, might succeed later
-        break; // Commit is only best effort
       }
     }
   }
@@ -190,8 +194,12 @@ public Instant getWatermark() {
 
   @Override
   public UnboundedSource.CheckpointMark getCheckpointMark() {
+    Queue<BytesXMLMessage> safeToAckMessages = new ConcurrentLinkedQueue<>();
     safeToAckMessages.addAll(receivedMessages);
     receivedMessages.clear();
+    nackCallback =
+        cleanUpThread.schedule(
+            () -> nackMessages(safeToAckMessages), ackDeadlineSeconds, TimeUnit.SECONDS);
     return new SolaceCheckpointMark(safeToAckMessages);
   }
 

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

@@ -49,6 +49,7 @@ public class UnboundedSolaceSource<T> extends UnboundedSource<T, SolaceCheckpoin
   private final SerializableFunction<T, Instant> timestampFn;
   private final Duration watermarkIdleDurationThreshold;
   private final SerializableFunction<@Nullable BytesXMLMessage, @Nullable T> parseFn;
+  private final int ackDeadlineSeconds;
 
   public Queue getQueue() {
     return queue;
@@ -70,6 +71,10 @@ public Duration getWatermarkIdleDurationThreshold() {
     return watermarkIdleDurationThreshold;
   }
 
+  public int getAckDeadlineSeconds() {
+    return ackDeadlineSeconds;
+  }
+
   public SerializableFunction<@Nullable BytesXMLMessage, @Nullable T> getParseFn() {
     return parseFn;
   }
@@ -83,7 +88,8 @@ public UnboundedSolaceSource(
       Coder<T> coder,
       SerializableFunction<T, Instant> timestampFn,
       Duration watermarkIdleDurationThreshold,
-      SerializableFunction<@Nullable BytesXMLMessage, @Nullable T> parseFn) {
+      SerializableFunction<@Nullable BytesXMLMessage, @Nullable T> parseFn,
+      int ackDeadlineSeconds) {
     this.queue = queue;
     this.sempClientFactory = sempClientFactory;
     this.sessionServiceFactory = sessionServiceFactory;
@@ -93,6 +99,7 @@ public UnboundedSolaceSource(
     this.timestampFn = timestampFn;
     this.watermarkIdleDurationThreshold = watermarkIdleDurationThreshold;
     this.parseFn = parseFn;
+    this.ackDeadlineSeconds = ackDeadlineSeconds;
   }
 
   @Override
@@ -134,7 +141,8 @@ private List<UnboundedSolaceSource<T>> getSolaceSources(
               coder,
               timestampFn,
               watermarkIdleDurationThreshold,
-              parseFn);
+              parseFn,
+              ackDeadlineSeconds);
       sourceList.add(source);
     }
     return sourceList;

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

@@ -72,14 +72,16 @@ private Read<Record> getDefaultRead() {
     return SolaceIO.read()
         .from(Solace.Queue.fromName("queue"))
         .withSempClientFactory(MockSempClientFactory.getDefaultMock())
-        .withMaxNumConnections(1);
+        .withMaxNumConnections(1)
+        .withAckDeadlineSeconds(1);
   }
 
   private Read<Record> getDefaultReadForTopic() {
     return SolaceIO.read()
         .from(Solace.Topic.fromName("topic"))
         .withSempClientFactory(MockSempClientFactory.getDefaultMock())
-        .withMaxNumConnections(1);
+        .withMaxNumConnections(1)
+        .withAckDeadlineSeconds(1);
   }
 
   private static BytesXMLMessage getOrNull(Integer index, List<BytesXMLMessage> messages) {
@@ -97,7 +99,8 @@ private static UnboundedSolaceSource<Record> getSource(Read<Record> spec, TestPi
         spec.inferCoder(pipeline, configuration.getTypeDescriptor()),
         configuration.getTimestampFn(),
         configuration.getWatermarkIdleDurationThreshold(),
-        configuration.getParseFn());
+        configuration.getParseFn(),
+        configuration.getAckDeadlineSeconds());
   }
 
   @Test
@@ -437,7 +440,8 @@ public void testCheckpointMarkAndFinalizeSeparately() throws Exception {
     Read<Record> spec =
         getDefaultRead()
             .withSessionServiceFactory(fakeSessionServiceFactory)
-            .withMaxNumConnections(4);
+            .withMaxNumConnections(4)
+            .withAckDeadlineSeconds(1);
 
     UnboundedSolaceSource<Record> initialSource = getSource(spec, pipeline);
 
@@ -458,19 +462,26 @@ 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. This will still not call ack.
     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();
     // No change in the acknowledged messages, because they were acknowledged in the #advance()
     // method.
     assertEquals(4, countAckMessages.get());
+
+    checkpointMark = reader.getCheckpointMark();
+
+    Thread.sleep(2000);
+    checkpointMark.finalizeCheckpoint();
+    // messages were nacked, no chane in expected values
+    assertEquals(4, countAckMessages.get());
   }
 
   @Test
@@ -542,7 +553,7 @@ public void testCheckpointMarkSafety() throws Exception {
   @Test
   public void testDefaultCoder() {
     Coder<SolaceCheckpointMark> coder =
-        new UnboundedSolaceSource<>(null, null, null, 0, false, null, null, null, null)
+        new UnboundedSolaceSource<>(null, null, null, 0, false, null, null, null, null, 1)
             .getCheckpointMarkCoder();
     CoderProperties.coderSerializable(coder);
   }

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

@@ -102,6 +102,7 @@ public void test01writeAndReadWithCustomSessionServiceFactory() {
             SolaceIO.read()
                 .from(Queue.fromName(QUEUE_NAME))
                 .withMaxNumConnections(1)
+                .withAckDeadlineSeconds(10)
                 .withDeduplicateRecords(true)
                 .withSempClientFactory(
                     BasicAuthSempClientFactory.builder()
@@ -118,7 +119,7 @@ public void test01writeAndReadWithCustomSessionServiceFactory() {
     PipelineResult pipelineResult = writerPipeline.run();
     // We need enough time for Beam to pull all messages from the queue, but we need a timeout too,
     // as the Read connector will keep attempting to read forever.
-    pipelineResult.waitUntilFinish(Duration.standardSeconds(15));
+    pipelineResult.waitUntilFinish(Duration.standardMinutes(2));
 
     MetricsReader metricsReader = new MetricsReader(pipelineResult, NAMESPACE);
     long actualRecordsCount = metricsReader.getCounterMetric(READ_COUNT);

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

@@ -107,6 +107,7 @@ public void test01writeAndReadWithMultipleSempClientFactory() {
             SolaceIO.read()
                 .from(Queue.fromName(QUEUE_NAME))
                 .withMaxNumConnections(1)
+                .withAckDeadlineSeconds(10)
                 .withDeduplicateRecords(true)
                 .withSempClientFactory(
                     BasicAuthMultipleSempClientFactory.builder()
@@ -131,7 +132,7 @@ public void test01writeAndReadWithMultipleSempClientFactory() {
     PipelineResult pipelineResult = writerPipeline.run();
     // We need enough time for Beam to pull all messages from the queue, but we need a timeout too,
     // as the Read connector will keep attempting to read forever.
-    pipelineResult.waitUntilFinish(Duration.standardSeconds(15));
+    pipelineResult.waitUntilFinish(Duration.standardMinutes(2));
 
     MetricsReader metricsReader = new MetricsReader(pipelineResult, NAMESPACE);
     long actualRecordsCount = metricsReader.getCounterMetric(READ_COUNT);