Make the PublishResult class for handling callbacks non-static to handle pipelines with multiple write transforms.

Author: bzablocki

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

@@ -36,6 +36,7 @@
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.io.solace.RetryCallableManager;
 import org.apache.beam.sdk.io.solace.SolaceIO.SubmissionMode;
+import org.apache.beam.sdk.io.solace.write.PublishResultsReceiver;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 
 /**
@@ -46,6 +47,7 @@
  */
 @AutoValue
 public abstract class BasicAuthJcsmpSessionService extends SessionService {
+
   /** The name of the queue to receive messages from. */
   public abstract @Nullable String queueName();
 
@@ -83,6 +85,7 @@ public abstract static class Builder {
   @Nullable private transient JCSMPSession jcsmpSession;
   @Nullable private transient MessageReceiver messageReceiver;
   @Nullable private transient MessageProducer messageProducer;
+  private final PublishResultsReceiver publishResultsReceiver = new PublishResultsReceiver();
   private final RetryCallableManager retryCallableManager = RetryCallableManager.create();
 
   @Override
@@ -119,7 +122,7 @@ public MessageReceiver getReceiver() {
   }
 
   @Override
-  public MessageProducer getProducer(SubmissionMode submissionMode) {
+  public MessageProducer getInitializeProducer(SubmissionMode submissionMode) {
     if (this.messageProducer == null || this.messageProducer.isClosed()) {
       Callable<MessageProducer> create = () -> createXMLMessageProducer(submissionMode);
       this.messageProducer =
@@ -128,20 +131,28 @@ public MessageProducer getProducer(SubmissionMode submissionMode) {
     return checkStateNotNull(this.messageProducer);
   }
 
+  @Override
+  public PublishResultsReceiver getPublishResultsReceiver() {
+    return publishResultsReceiver;
+  }
+
   @Override
   public boolean isClosed() {
     return jcsmpSession == null || jcsmpSession.isClosed();
   }
 
   private MessageProducer createXMLMessageProducer(SubmissionMode submissionMode)
       throws JCSMPException, IOException {
+
     if (isClosed()) {
       connectWriteSession(submissionMode);
     }
 
     @SuppressWarnings("nullness")
     Callable<XMLMessageProducer> initProducer =
-        () -> Objects.requireNonNull(jcsmpSession).getMessageProducer(new PublishResultHandler());
+        () ->
+            Objects.requireNonNull(jcsmpSession)
+                .getMessageProducer(new PublishResultHandler(publishResultsReceiver));
 
     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

@@ -22,6 +22,9 @@
 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.PublishResultsReceiver;
+import org.apache.beam.sdk.io.solace.write.UnboundedSolaceWriter;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -38,6 +41,13 @@
 public final class PublishResultHandler implements JCSMPStreamingPublishCorrelatingEventHandler {
 
   private static final Logger LOG = LoggerFactory.getLogger(PublishResultHandler.class);
+  private final PublishResultsReceiver publishResultsReceiver;
+  private final Counter batchesRejectedByBroker =
+      Metrics.counter(UnboundedSolaceWriter.class, "batches_rejected");
+
+  public PublishResultHandler(PublishResultsReceiver publishResultsReceiver) {
+    this.publishResultsReceiver = publishResultsReceiver;
+  }
 
   @Override
   public void handleErrorEx(Object key, JCSMPException cause, long timestamp) {
@@ -64,6 +74,7 @@ private void processKey(Object key, boolean isPublished, @Nullable JCSMPExceptio
 
     resultBuilder = resultBuilder.setMessageId(messageId).setPublished(isPublished);
     if (!isPublished) {
+      batchesRejectedByBroker.inc();
       if (cause != null) {
         resultBuilder = resultBuilder.setError(cause.getMessage());
       } else {
@@ -78,7 +89,7 @@ private void processKey(Object key, boolean isPublished, @Nullable JCSMPExceptio
     PublishResult publishResult = resultBuilder.build();
     // Static reference, it receives all callbacks from all publications
     // from all threads
-    PublishResultsReceiver.addResult(publishResult);
+    publishResultsReceiver.addResult(publishResult);
   }
 
   private static long calculateLatency(Solace.CorrelationKey key) {

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

@@ -21,6 +21,7 @@
 import java.io.Serializable;
 import org.apache.beam.sdk.io.solace.SolaceIO;
 import org.apache.beam.sdk.io.solace.SolaceIO.SubmissionMode;
+import org.apache.beam.sdk.io.solace.write.PublishResultsReceiver;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -137,7 +138,15 @@ public abstract class SessionService implements Serializable {
    * this method is used, the producer is created from the session instance, otherwise it returns
    * the producer created initially.
    */
-  public abstract MessageProducer getProducer(SubmissionMode mode);
+  public abstract MessageProducer getInitializeProducer(SubmissionMode mode);
+
+  /**
+   * Returns the {@link PublishResultsReceiver} instance associated with this session.
+   *
+   * <p>The {@link PublishResultsReceiver} handles asynchronous callbacks from Solace, providing
+   * results for message publications.
+   */
+  public abstract PublishResultsReceiver getPublishResultsReceiver();
 
   /**
    * Override this method and provide your specific properties, including all those related to

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

@@ -30,14 +30,13 @@
  */
 @Internal
 public final class PublishResultsReceiver {
-  private static final ConcurrentLinkedQueue<PublishResult> resultsQueue =
-      new ConcurrentLinkedQueue<>();
+  private final ConcurrentLinkedQueue<PublishResult> resultsQueue = new ConcurrentLinkedQueue<>();
 
-  public static @Nullable PublishResult pollResults() {
+  public @Nullable PublishResult pollResults() {
     return resultsQueue.poll();
   }
 
-  public static boolean addResult(PublishResult result) {
+  public boolean addResult(PublishResult result) {
     return resultsQueue.add(result);
   }
 }

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

@@ -21,6 +21,7 @@
 import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 
 import com.google.auto.value.AutoValue;
+import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import org.apache.beam.sdk.io.solace.SolaceIO.SubmissionMode;
 import org.apache.beam.sdk.io.solace.broker.SessionService;
@@ -39,15 +40,17 @@
  * than one, each {@link SessionServiceFactory} instance keeps their own pool of producers.
  */
 final class SolaceWriteSessionsHandler {
+
   private static final ConcurrentHashMap<SessionConfigurationIndex, SessionService> sessionsMap =
       new ConcurrentHashMap<>(DEFAULT_WRITER_CLIENTS_PER_WORKER);
 
-  public static SessionService getSessionService(
-      int producerIndex, SessionServiceFactory sessionServiceFactory) {
+  public static SessionService getSessionServiceWithProducer(
+      int producerIndex, SessionServiceFactory sessionServiceFactory, UUID writerTransformUuid) {
     SessionConfigurationIndex key =
         SessionConfigurationIndex.builder()
             .producerIndex(producerIndex)
             .sessionServiceFactory(sessionServiceFactory)
+            .writerTransformUuid(writerTransformUuid)
             .build();
     return sessionsMap.computeIfAbsent(
         key, SolaceWriteSessionsHandler::createSessionAndStartProducer);
@@ -61,17 +64,19 @@ private static SessionService createSessionAndStartProducer(SessionConfiguration
     checkStateNotNull(
         mode,
         "SolaceIO.Write: Submission mode is not set. You need to set it to create write sessions.");
-    sessionService.getProducer(mode);
+    sessionService.getInitializeProducer(mode);
     return sessionService;
   }
 
   /** Disconnect all the sessions from Solace, and clear the corresponding state. */
-  public static void disconnectFromSolace(SessionServiceFactory factory, int producersCardinality) {
+  public static void disconnectFromSolace(
+      SessionServiceFactory factory, int producersCardinality, UUID writerTransformUuid) {
     for (int i = 0; i < producersCardinality; i++) {
       SessionConfigurationIndex key =
           SessionConfigurationIndex.builder()
               .producerIndex(i)
               .sessionServiceFactory(factory)
+              .writerTransformUuid(writerTransformUuid)
               .build();
 
       SessionService sessionService = sessionsMap.remove(key);
@@ -87,6 +92,8 @@ abstract static class SessionConfigurationIndex {
 
     abstract SessionServiceFactory sessionServiceFactory();
 
+    abstract UUID writerTransformUuid();
+
     static Builder builder() {
       return new AutoValue_SolaceWriteSessionsHandler_SessionConfigurationIndex.Builder();
     }
@@ -97,6 +104,8 @@ abstract static class Builder {
 
       abstract Builder sessionServiceFactory(SessionServiceFactory sessionServiceFactory);
 
+      abstract Builder writerTransformUuid(UUID writerTransformUuid);
+
       abstract SessionConfigurationIndex build();
     }
   }

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

@@ -69,7 +69,7 @@ public final class UnboundedBatchedSolaceWriter extends UnboundedSolaceWriter {
       Metrics.counter(UnboundedBatchedSolaceWriter.class, "msgs_sent_to_broker");
 
   private final Counter batchesRejectedByBroker =
-      Metrics.counter(UnboundedBatchedSolaceWriter.class, "batches_rejected");
+      Metrics.counter(UnboundedSolaceWriter.class, "batches_rejected");
 
   // State variables are never explicitly "used"
   @SuppressWarnings("UnusedVariable")
@@ -142,8 +142,8 @@ public void flushBundle(OnTimerContext context) {
   private void publishBatch(List<Solace.Record> records) {
     try {
       int entriesPublished =
-          solaceSessionService()
-              .getProducer(getSubmissionMode())
+          solaceSessionServiceWithProducer()
+              .getInitializeProducer(getSubmissionMode())
               .publishBatch(
                   records, shouldPublishLatencyMetrics(), getDestinationFn(), getDeliveryMode());
       sentToBroker.inc(entriesPublished);
@@ -159,7 +159,7 @@ private void publishBatch(List<Solace.Record> records) {
                       e.getMessage()))
               .setLatencyNanos(System.nanoTime())
               .build();
-      PublishResultsReceiver.addResult(errorPublish);
+      solaceSessionServiceWithProducer().getPublishResultsReceiver().addResult(errorPublish);
     }
   }
 }

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

@@ -27,6 +27,7 @@
 import com.solacesystems.jcsmp.JCSMPSendMultipleEntry;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.io.solace.SolaceIO;
@@ -80,6 +81,8 @@ public abstract class UnboundedSolaceWriter
   private @Nullable Instant bundleTimestamp;
   private @Nullable BoundedWindow bundleWindow;
 
+  final UUID writerTransformUuid = UUID.randomUUID();
+
   public UnboundedSolaceWriter(
       SerializableFunction<Solace.Record, Destination> destinationFn,
       SessionServiceFactory sessionServiceFactory,
@@ -101,7 +104,8 @@ public UnboundedSolaceWriter(
 
   @Teardown
   public void teardown() {
-    SolaceWriteSessionsHandler.disconnectFromSolace(sessionServiceFactory, producersMapCardinality);
+    SolaceWriteSessionsHandler.disconnectFromSolace(
+        sessionServiceFactory, producersMapCardinality, writerTransformUuid);
   }
 
   public void updateProducerIndex() {
@@ -115,9 +119,9 @@ public void startBundle() {
     batchToEmit.clear();
   }
 
-  public SessionService solaceSessionService() {
-    return SolaceWriteSessionsHandler.getSessionService(
-        currentBundleProducerIndex, sessionServiceFactory);
+  public SessionService solaceSessionServiceWithProducer() {
+    return SolaceWriteSessionsHandler.getSessionServiceWithProducer(
+        currentBundleProducerIndex, sessionServiceFactory, writerTransformUuid);
   }
 
   public void publishResults(BeamContextWrapper context) {
@@ -131,7 +135,9 @@ public void publishResults(BeamContextWrapper context) {
     long minFailed = Long.MAX_VALUE;
     long maxFailed = 0;
 
-    Solace.PublishResult result = PublishResultsReceiver.pollResults();
+    PublishResultsReceiver publishResultsReceiver =
+        solaceSessionServiceWithProducer().getPublishResultsReceiver();
+    Solace.PublishResult result = publishResultsReceiver.pollResults();
 
     if (result != null) {
       if (getCurrentBundleTimestamp() == null) {
@@ -174,7 +180,7 @@ public void publishResults(BeamContextWrapper context) {
             FAILED_PUBLISH_TAG, result, getCurrentBundleTimestamp(), getCurrentBundleWindow());
       }
 
-      result = PublishResultsReceiver.pollResults();
+      result = publishResultsReceiver.pollResults();
     }
 
     if (shouldPublishLatencyMetrics()) {

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

@@ -110,8 +110,8 @@ public void processElement(
 
     // The publish method will retry, let's send a failure message if all the retries fail
     try {
-      solaceSessionService()
-          .getProducer(getSubmissionMode())
+      solaceSessionServiceWithProducer()
+          .getInitializeProducer(getSubmissionMode())
           .publishSingleMessage(
               record,
               getDestinationFn().apply(record),
@@ -130,7 +130,7 @@ public void processElement(
                       e.getMessage()))
               .setLatencyNanos(System.nanoTime())
               .build();
-      PublishResultsReceiver.addResult(errorPublish);
+      solaceSessionServiceWithProducer().getPublishResultsReceiver().addResult(errorPublish);
     }
   }
 

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

@@ -23,6 +23,7 @@
 import org.apache.beam.sdk.io.solace.broker.MessageProducer;
 import org.apache.beam.sdk.io.solace.broker.MessageReceiver;
 import org.apache.beam.sdk.io.solace.broker.SessionService;
+import org.apache.beam.sdk.io.solace.write.PublishResultsReceiver;
 
 @AutoValue
 public abstract class MockEmptySessionService extends SessionService {
@@ -49,7 +50,12 @@ public MessageReceiver getReceiver() {
   }
 
   @Override
-  public MessageProducer getProducer(SubmissionMode mode) {
+  public MessageProducer getInitializeProducer(SubmissionMode mode) {
+    throw new UnsupportedOperationException(exceptionMessage);
+  }
+
+  @Override
+  public PublishResultsReceiver getPublishResultsReceiver() {
     throw new UnsupportedOperationException(exceptionMessage);
   }
 

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

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.io.solace.broker.PublishResultHandler;
 import org.apache.beam.sdk.io.solace.broker.SessionService;
 import org.apache.beam.sdk.io.solace.data.Solace;
+import org.apache.beam.sdk.io.solace.write.PublishResultsReceiver;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -45,6 +46,8 @@ public abstract class MockSessionService extends SessionService {
 
   public abstract @Nullable SubmissionMode mode();
 
+  private final PublishResultsReceiver publishResultsReceiver = new PublishResultsReceiver();
+
   public static Builder builder() {
     return new AutoValue_MockSessionService.Builder().minMessagesReceived(0);
   }
@@ -81,13 +84,18 @@ public MessageReceiver getReceiver() {
   }
 
   @Override
-  public MessageProducer getProducer(SubmissionMode mode) {
+  public MessageProducer getInitializeProducer(SubmissionMode mode) {
     if (messageProducer == null) {
-      messageProducer = new MockProducer(new PublishResultHandler());
+      messageProducer = new MockProducer(new PublishResultHandler(publishResultsReceiver));
     }
     return messageProducer;
   }
 
+  @Override
+  public PublishResultsReceiver getPublishResultsReceiver() {
+    return publishResultsReceiver;
+  }
+
   @Override
   public void connect() {}