fix: wait for pending ACK callbacks in SolaceIO write finishBundle to prevent batch pipeline data loss [fj4WqyCCw3C5ShR1RfB7MoBPTpkRrBFYP1uT35g3MvT]

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

@@ -34,6 +34,7 @@
 import java.util.Objects;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.io.solace.RetryCallableManager;
 import org.apache.beam.sdk.io.solace.SolaceIO.SubmissionMode;
@@ -59,6 +60,7 @@ public abstract class JcsmpSessionService extends SessionService {
   @Nullable private transient MessageProducer messageProducer;
   private final java.util.Queue<PublishResult> publishedResultsQueue =
       new ConcurrentLinkedQueue<>();
+  private final AtomicInteger pendingPublishCount = new AtomicInteger(0);
   private final RetryCallableManager retryCallableManager = RetryCallableManager.create();
 
   public static JcsmpSessionService create(JCSMPProperties jcsmpProperties, @Nullable Queue queue) {
@@ -117,6 +119,11 @@ public java.util.Queue<PublishResult> getPublishedResultsQueue() {
     return publishedResultsQueue;
   }
 
+  @Override
+  public AtomicInteger getPendingPublishCount() {
+    return pendingPublishCount;
+  }
+
   private MessageProducer createXMLMessageProducer(SubmissionMode submissionMode)
       throws JCSMPException, IOException {
 
@@ -128,7 +135,7 @@ private MessageProducer createXMLMessageProducer(SubmissionMode submissionMode)
     Callable<XMLMessageProducer> initProducer =
         () ->
             Objects.requireNonNull(jcsmpSession)
-                .getMessageProducer(new PublishResultHandler(publishedResultsQueue));
+                .getMessageProducer(new PublishResultHandler(publishedResultsQueue, pendingPublishCount));
 
     XMLMessageProducer producer =
         retryCallableManager.retryCallable(initProducer, ImmutableSet.of(JCSMPException.class));

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

@@ -20,6 +20,7 @@
 import com.solacesystems.jcsmp.JCSMPException;
 import com.solacesystems.jcsmp.JCSMPStreamingPublishCorrelatingEventHandler;
 import java.util.Queue;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.beam.sdk.io.solace.data.Solace;
 import org.apache.beam.sdk.io.solace.data.Solace.PublishResult;
 import org.apache.beam.sdk.io.solace.write.UnboundedSolaceWriter;
@@ -42,11 +43,14 @@ public final class PublishResultHandler implements JCSMPStreamingPublishCorrelat
 
   private static final Logger LOG = LoggerFactory.getLogger(PublishResultHandler.class);
   private final Queue<PublishResult> publishResultsQueue;
+  private final AtomicInteger pendingPublishCount;
   private final Counter batchesRejectedByBroker =
       Metrics.counter(UnboundedSolaceWriter.class, "batches_rejected");
 
-  public PublishResultHandler(Queue<PublishResult> publishResultsQueue) {
+  public PublishResultHandler(
+      Queue<PublishResult> publishResultsQueue, AtomicInteger pendingPublishCount) {
     this.publishResultsQueue = publishResultsQueue;
+    this.pendingPublishCount = pendingPublishCount;
   }
 
   @Override
@@ -90,6 +94,7 @@ private void processKey(Object key, boolean isPublished, @Nullable JCSMPExceptio
     // Static reference, it receives all callbacks from all publications
     // from all threads
     publishResultsQueue.add(publishResult);
+    pendingPublishCount.decrementAndGet();
   }
 
   private static long calculateLatency(Solace.CorrelationKey key) {

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

@@ -20,6 +20,7 @@
 import com.solacesystems.jcsmp.JCSMPProperties;
 import java.io.Serializable;
 import java.util.Queue;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.beam.sdk.io.solace.SolaceIO;
 import org.apache.beam.sdk.io.solace.SolaceIO.SubmissionMode;
 import org.apache.beam.sdk.io.solace.data.Solace.PublishResult;
@@ -140,6 +141,16 @@ public abstract class SessionService implements Serializable {
    */
   public abstract Queue<PublishResult> getPublishedResultsQueue();
 
+  /**
+   * Returns the {@link AtomicInteger} tracking the number of in-flight publish operations.
+   *
+   * <p>The counter is incremented when a publish is sent and decremented when the asynchronous
+   * Solace ACK callback completes. This allows the writer to wait for all pending publishes to
+   * complete before emitting results in {@code @FinishBundle}, preventing data loss in batch
+   * pipelines.
+   */
+  public abstract AtomicInteger getPendingPublishCount();
+
   /**
    * Override this method and provide your specific properties, including all those related to
    * authentication, and possibly others too.

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

@@ -126,26 +126,34 @@ public void finishBundle(FinishBundleContext context) throws IOException {
       if (batch.isEmpty()) {
         continue;
       }
-      publishBatch(batch);
+      int entriesPublished = publishBatch(batch);
+      sentToBroker.inc(entriesPublished);
+      incrementPendingPublishes(entriesPublished);
     }
     getCurrentBundle().clear();
 
+    // Wait for pending asynchronous Solace ACK callbacks to complete before emitting publish
+    // results. In batch pipelines, the pipeline may end before the timer fires and before
+    // ACK callbacks arrive, causing data loss. This wait ensures all in-flight publishes
+    // are acknowledged before we emit results.
+    waitForPendingPublishes();
     publishResults(BeamContextWrapper.of(context));
   }
 
   @OnTimer("bundle_flusher")
   public void flushBundle(OnTimerContext context) throws IOException {
+    waitForPendingPublishes();
     publishResults(BeamContextWrapper.of(context));
   }
 
-  private void publishBatch(List<Solace.Record> records) {
+  private int publishBatch(List<Solace.Record> records) {
     try {
       int entriesPublished =
           solaceSessionServiceWithProducer()
               .getInitializedProducer(getSubmissionMode())
               .publishBatch(
                   records, shouldPublishLatencyMetrics(), getDestinationFn(), getDeliveryMode());
-      sentToBroker.inc(entriesPublished);
+      return entriesPublished;
     } catch (Exception e) {
       batchesRejectedByBroker.inc();
       Solace.PublishResult errorPublish =
@@ -159,6 +167,10 @@ private void publishBatch(List<Solace.Record> records) {
               .setLatencyNanos(System.nanoTime())
               .build();
       solaceSessionServiceWithProducer().getPublishedResultsQueue().add(errorPublish);
+      // Even though the batch failed, we need to track the pending count so that
+      // waitForPendingPublishes() doesn't wait indefinitely. The error result is already
+      // in the queue, so the count doesn't need to be incremented.
+      return 0;
     }
   }
 }

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

@@ -68,6 +68,7 @@ public abstract class UnboundedSolaceWriter
 
   // This is the batch limit supported by the send multiple JCSMP API method.
   static final int SOLACE_BATCH_LIMIT = 50;
+  private static final long PUBLISH_ACKS_TIMEOUT_SECS = 30;
   private final Distribution latencyPublish =
       Metrics.distribution(SolaceIO.Write.class, "latency_publish_ms");
 
@@ -132,6 +133,50 @@ public SessionService solaceSessionServiceWithProducer() {
         currentBundleProducerIndex, sessionServiceFactory, writerTransformUuid);
   }
 
+  /**
+   * Increments the pending publish count for the current session's producer. This count is
+   * decremented by the {@link org.apache.beam.sdk.io.solace.broker.PublishResultHandler} when each
+   * asynchronous Solace ACK callback arrives.
+   *
+   * <p>Use this to track in-flight publish operations so that {@link #waitForPendingPublishes()}
+   * can block until all ACKs have been received.
+   */
+  public void incrementPendingPublishes(int count) {
+    solaceSessionServiceWithProducer().getPendingPublishCount().addAndGet(count);
+  }
+
+  /**
+   * Waits for all in-flight publish operations to complete, with a timeout of {@value
+   * #PUBLISH_ACKS_TIMEOUT_SECS} seconds.
+   *
+   * <p>This is necessary in batch pipelines where the asynchronous Solace ACK callbacks may not
+   * have arrived by the time {@code @FinishBundle} is called. Without this wait, the pipeline may
+   * end before publish results are emitted, causing data loss.
+   *
+   * <p>This method is a no-op if there are no pending publishes, or if the thread is interrupted.
+   */
+  public void waitForPendingPublishes() throws IOException {
+    AtomicInteger pending = solaceSessionServiceWithProducer().getPendingPublishCount();
+    long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(PUBLISH_ACKS_TIMEOUT_SECS);
+    while (pending.get() > 0 && System.nanoTime() < deadline) {
+      try {
+        Thread.sleep(50);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        LOG.warn(
+            "SolaceIO.Write: Interrupted while waiting for {} pending publish ACKs.",
+            pending.get());
+        return;
+      }
+    }
+    if (pending.get() > 0) {
+      LOG.warn(
+          "SolaceIO.Write: Timed out waiting for {} pending publish ACKs after {} seconds.",
+          pending.get(),
+          PUBLISH_ACKS_TIMEOUT_SECS);
+    }
+  }
+
   public void publishResults(BeamContextWrapper context) {
     long sumPublish = 0;
     long countPublish = 0;

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

@@ -19,6 +19,7 @@
 
 import com.solacesystems.jcsmp.DeliveryMode;
 import com.solacesystems.jcsmp.Destination;
+import java.io.IOException;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.io.solace.SolaceIO;
 import org.apache.beam.sdk.io.solace.broker.SessionServiceFactory;
@@ -115,6 +116,7 @@ public void processElement(
               shouldPublishLatencyMetrics(),
               getDeliveryMode());
       sentToBroker.inc();
+      incrementPendingPublishes(1);
     } catch (Exception e) {
       rejectedByBroker.inc();
       Solace.PublishResult errorPublish =
@@ -132,7 +134,13 @@ public void processElement(
   }
 
   @FinishBundle
-  public void finishBundle(FinishBundleContext context) {
+  public void finishBundle(FinishBundleContext context) throws IOException {
+    // Wait for pending asynchronous Solace ACK callbacks to complete before emitting publish
+    // results. In batch pipelines, the asynchronous ACK callbacks from publishSingleMessage
+    // may not have arrived by the time finishBundle is called, causing the pipeline to end
+    // before publish results are emitted. This wait ensures all in-flight publishes are
+    // acknowledged before we emit results.
+    waitForPendingPublishes();
     publishResults(BeamContextWrapper.of(context));
   }
 }