SolaceIO.Read: handle occasional cases when finalizeCheckpoint is not executed

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

@@ -102,10 +102,7 @@ public void close() {
           if (messageReceiver != null) {
             messageReceiver.close();
           }
-          if (messageProducer != null) {
-            messageProducer.close();
-          }
-          if (!isClosed()) {
+          if (jcsmpSession != null) {
             checkStateNotNull(jcsmpSession).closeSession();
           }
           return 0;
@@ -119,8 +116,9 @@ public MessageReceiver getReceiver() {
       this.messageReceiver =
           retryCallableManager.retryCallable(
               this::createFlowReceiver, ImmutableSet.of(JCSMPException.class));
+      this.messageReceiver.start();
     }
-    return this.messageReceiver;
+    return checkStateNotNull(this.messageReceiver);
   }
 
   @Override
@@ -138,15 +136,10 @@ public java.util.Queue<PublishResult> getPublishedResultsQueue() {
     return publishedResultsQueue;
   }
 
-  @Override
-  public boolean isClosed() {
-    return jcsmpSession == null || jcsmpSession.isClosed();
-  }
-
   private MessageProducer createXMLMessageProducer(SubmissionMode submissionMode)
       throws JCSMPException, IOException {
 
-    if (isClosed()) {
+    if (jcsmpSession == null) {
       connectWriteSession(submissionMode);
     }
 
@@ -165,9 +158,6 @@ private MessageProducer createXMLMessageProducer(SubmissionMode submissionMode)
   }
 
   private MessageReceiver createFlowReceiver() throws JCSMPException, IOException {
-    if (isClosed()) {
-      connectSession();
-    }
 
     Queue queue =
         JCSMPFactory.onlyInstance()

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

@@ -35,13 +35,6 @@ public interface MessageReceiver {
    */
   void start();
 
-  /**
-   * Returns {@literal true} if the message receiver is closed, {@literal false} otherwise.
-   *
-   * <p>A message receiver is closed when it is no longer able to receive messages.
-   */
-  boolean isClosed();
-
   /**
    * Receives a message from the broker.
    *

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

@@ -120,13 +120,6 @@ public abstract class SessionService implements Serializable {
   /** Gracefully closes the connection to the service. */
   public abstract void close();
 
-  /**
-   * Checks whether the connection to the service is currently closed. This method is called when an
-   * `UnboundedSolaceReader` is starting to read messages - a session will be created if this
-   * returns true.
-   */
-  public abstract boolean isClosed();
-
   /**
    * Returns a MessageReceiver object for receiving messages from Solace. If it is the first time
    * this method is used, the receiver is created from the session instance, otherwise it returns

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

@@ -24,12 +24,8 @@
 import java.io.IOException;
 import org.apache.beam.sdk.io.solace.RetryCallableManager;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 public class SolaceMessageReceiver implements MessageReceiver {
-  private static final Logger LOG = LoggerFactory.getLogger(SolaceMessageReceiver.class);
-
   public static final int DEFAULT_ADVANCE_TIMEOUT_IN_MILLIS = 100;
   private final FlowReceiver flowReceiver;
   private final RetryCallableManager retryCallableManager = RetryCallableManager.create();
@@ -52,28 +48,21 @@ private void startFlowReceiver() {
         ImmutableSet.of(JCSMPException.class));
   }
 
-  @Override
-  public boolean isClosed() {
-    return flowReceiver == null || flowReceiver.isClosed();
-  }
-
   @Override
   public BytesXMLMessage receive() throws IOException {
     try {
       return flowReceiver.receive(DEFAULT_ADVANCE_TIMEOUT_IN_MILLIS);
     } catch (StaleSessionException e) {
-      LOG.warn("SolaceIO: Caught StaleSessionException, restarting the FlowReceiver.");
       startFlowReceiver();
-      throw new IOException(e);
+      throw new IOException(
+          "SolaceIO: Caught StaleSessionException, restarting the FlowReceiver.", e);
     } catch (JCSMPException e) {
       throw new IOException(e);
     }
   }
 
   @Override
   public void close() {
-    if (!isClosed()) {
-      this.flowReceiver.close();
-    }
+    flowReceiver.close();
   }
 }

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

@@ -18,17 +18,16 @@
 package org.apache.beam.sdk.io.solace.read;
 
 import com.solacesystems.jcsmp.BytesXMLMessage;
-import java.util.ArrayDeque;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.Objects;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.Queue;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.coders.DefaultCoder;
 import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Checkpoint for an unbounded Solace source. Consists of the Solace messages waiting to be
@@ -38,36 +37,33 @@
 @Internal
 @VisibleForTesting
 public class SolaceCheckpointMark implements UnboundedSource.CheckpointMark {
-  private transient AtomicBoolean activeReader;
-  // BytesXMLMessage is not serializable so if a job restarts from the checkpoint, we cannot retry
-  // these messages here. We relay on Solace's retry mechanism.
-  private transient ArrayDeque<BytesXMLMessage> ackQueue;
+  private static final Logger LOG = LoggerFactory.getLogger(SolaceCheckpointMark.class);
+  private transient Queue<BytesXMLMessage> safeToAck;
 
   @SuppressWarnings("initialization") // Avro will set the fields by breaking abstraction
   private SolaceCheckpointMark() {}
 
   /**
    * Creates a new {@link SolaceCheckpointMark}.
    *
-   * @param activeReader {@link AtomicBoolean} indicating if the related reader is active. The
-   *     reader creating the messages has to be active to acknowledge the messages.
-   * @param ackQueue {@link List} of {@link BytesXMLMessage} to be acknowledged.
+   * @param safeToAck - a queue of {@link BytesXMLMessage} to be acknowledged.
    */
-  SolaceCheckpointMark(AtomicBoolean activeReader, List<BytesXMLMessage> ackQueue) {
-    this.activeReader = activeReader;
-    this.ackQueue = new ArrayDeque<>(ackQueue);
+  SolaceCheckpointMark(Queue<BytesXMLMessage> safeToAck) {
+    this.safeToAck = safeToAck;
   }
 
   @Override
   public void finalizeCheckpoint() {
-    if (activeReader == null || !activeReader.get() || ackQueue == null) {
-      return;
-    }
-
-    while (!ackQueue.isEmpty()) {
-      BytesXMLMessage msg = ackQueue.poll();
-      if (msg != null) {
+    BytesXMLMessage msg;
+    while ((msg = safeToAck.poll()) != null) {
+      try {
         msg.ackMessage();
+      } catch (IllegalStateException e) {
+        LOG.error(
+            "SolaceIO.Read: cannot acknowledge the message with applicationMessageId={}, ackMessageId={}. It will not be retried.",
+            msg.getApplicationMessageId(),
+            msg.getAckMessageId(),
+            e);
       }
     }
   }
@@ -84,15 +80,11 @@ public boolean equals(@Nullable Object o) {
       return false;
     }
     SolaceCheckpointMark that = (SolaceCheckpointMark) o;
-    // Needed to convert to ArrayList because ArrayDeque.equals checks only for reference, not
-    // content.
-    ArrayList<BytesXMLMessage> ackList = new ArrayList<>(ackQueue);
-    ArrayList<BytesXMLMessage> thatAckList = new ArrayList<>(that.ackQueue);
-    return Objects.equals(activeReader, that.activeReader) && Objects.equals(ackList, thatAckList);
+    return Objects.equals(safeToAck, that.safeToAck);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(activeReader, ackQueue);
+    return Objects.hash(safeToAck);
   }
 }