Expose watermarkIdleDurationThreshold parameter to the user in SolaceIO read connector

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

@@ -51,6 +51,7 @@
 import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -393,7 +394,8 @@ public class SolaceIO {
         }
       };
   private static final boolean DEFAULT_DEDUPLICATE_RECORDS = false;
-
+  private static final Duration DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD =
+      Duration.standardSeconds(30);
   public static final int DEFAULT_WRITER_MAX_NUMBER_OF_WORKERS = 20;
   public static final int DEFAULT_WRITER_CLIENTS_PER_WORKER = 4;
   public static final Boolean DEFAULT_WRITER_PUBLISH_LATENCY_METRICS = false;
@@ -440,7 +442,8 @@ public static Read<Solace.Record> read() {
             .setTypeDescriptor(TypeDescriptor.of(Solace.Record.class))
             .setParseFn(SolaceRecordMapper::map)
             .setTimestampFn(SENDER_TIMESTAMP_FUNCTION)
-            .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS));
+            .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS)
+            .setWatermarkIdleDurationThreshold(DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD));
   }
   /**
    * Create a {@link Read} transform, to read from Solace. Specify a {@link SerializableFunction} to
@@ -467,7 +470,8 @@ public static <T> Read<T> read(
             .setTypeDescriptor(typeDescriptor)
             .setParseFn(parseFn)
             .setTimestampFn(timestampFn)
-            .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS));
+            .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS)
+            .setWatermarkIdleDurationThreshold(DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD));
   }
 
   /**
@@ -540,6 +544,19 @@ public Read<T> withMaxNumConnections(Integer maxNumConnections) {
       return this;
     }
 
+    /**
+     * Optional. Denotes the duration for which the watermark can be idle. If there are no incoming
+     * messages for this ‘idle’ period of time, the watermark is set to a timestamp representing a
+     * time earlier than now by the ‘idle’ period of time (e.g. if the ‘idle’ period of time is set
+     * to 30 seconds, and there is no new data incoming for 30 seconds, the watermark will be set to
+     * max(currentWatermark, now() - 30 seconds). The default watermark idle duration threshold is
+     * {@link #DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD}.
+     */
+    public Read<T> withWatermarkIdleDurationThreshold(Duration idleDurationThreshold) {
+      configurationBuilder.setWatermarkIdleDurationThreshold(idleDurationThreshold);
+      return this;
+    }
+
     /**
      * Optional, default: false. Set to deduplicate messages based on the {@link
      * BytesXMLMessage#getApplicationMessageId()} of the incoming {@link BytesXMLMessage}. If the
@@ -652,6 +669,8 @@ abstract static class Configuration<T> {
 
       abstract TypeDescriptor<T> getTypeDescriptor();
 
+      abstract Duration getWatermarkIdleDurationThreshold();
+
       public static <T> Builder<T> builder() {
         Builder<T> builder =
             new org.apache.beam.sdk.io.solace.AutoValue_SolaceIO_Read_Configuration.Builder<T>();
@@ -680,6 +699,8 @@ abstract Builder<T> setParseFn(
 
         abstract Builder<T> setTypeDescriptor(TypeDescriptor<T> typeDescriptor);
 
+        abstract Builder<T> setWatermarkIdleDurationThreshold(Duration idleDurationThreshold);
+
         abstract Configuration<T> build();
       }
     }
@@ -716,6 +737,7 @@ public PCollection<T> expand(PBegin input) {
                   configuration.getDeduplicateRecords(),
                   coder,
                   configuration.getTimestampFn(),
+                  configuration.getWatermarkIdleDurationThreshold(),
                   configuration.getParseFn())));
     }
 

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

@@ -62,7 +62,9 @@ class UnboundedSolaceReader<T> extends UnboundedReader<T> {
 
   public UnboundedSolaceReader(UnboundedSolaceSource<T> currentSource) {
     this.currentSource = currentSource;
-    this.watermarkPolicy = WatermarkPolicy.create(currentSource.getTimestampFn());
+    this.watermarkPolicy =
+        WatermarkPolicy.create(
+            currentSource.getTimestampFn(), currentSource.getWatermarkIdleDurationThreshold());
     this.sessionService = currentSource.getSessionServiceFactory().create();
     this.sempClient = currentSource.getSempClientFactory().create();
   }

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

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -46,6 +47,7 @@ public class UnboundedSolaceSource<T> extends UnboundedSource<T, SolaceCheckpoin
   private final SempClientFactory sempClientFactory;
   private final SessionServiceFactory sessionServiceFactory;
   private final SerializableFunction<T, Instant> timestampFn;
+  private final Duration watermarkIdleDurationThreshold;
   private final SerializableFunction<@Nullable BytesXMLMessage, @Nullable T> parseFn;
 
   public Queue getQueue() {
@@ -64,6 +66,10 @@ public SerializableFunction<T, Instant> getTimestampFn() {
     return timestampFn;
   }
 
+  public Duration getWatermarkIdleDurationThreshold() {
+    return watermarkIdleDurationThreshold;
+  }
+
   public SerializableFunction<@Nullable BytesXMLMessage, @Nullable T> getParseFn() {
     return parseFn;
   }
@@ -76,6 +82,7 @@ public UnboundedSolaceSource(
       boolean enableDeduplication,
       Coder<T> coder,
       SerializableFunction<T, Instant> timestampFn,
+      Duration watermarkIdleDurationThreshold,
       SerializableFunction<@Nullable BytesXMLMessage, @Nullable T> parseFn) {
     this.queue = queue;
     this.sempClientFactory = sempClientFactory;
@@ -84,6 +91,7 @@ public UnboundedSolaceSource(
     this.enableDeduplication = enableDeduplication;
     this.coder = coder;
     this.timestampFn = timestampFn;
+    this.watermarkIdleDurationThreshold = watermarkIdleDurationThreshold;
     this.parseFn = parseFn;
   }
 
@@ -125,6 +133,7 @@ private List<UnboundedSolaceSource<T>> getSolaceSources(
               enableDeduplication,
               coder,
               timestampFn,
+              watermarkIdleDurationThreshold,
               parseFn);
       sourceList.add(source);
     }

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

@@ -21,17 +21,13 @@
 import java.io.Serializable;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 
 /** {@code WatermarkParameters} contains the parameters used for watermark computation. */
 @AutoValue
 abstract class WatermarkParameters<T> implements Serializable {
 
-  private static final Duration STANDARD_WATERMARK_IDLE_DURATION_THRESHOLD =
-      Duration.standardSeconds(30);
-
   abstract Instant getCurrentWatermark();
 
   abstract Instant getLastSavedWatermark();
@@ -48,8 +44,7 @@ static <T> Builder<T> builder() {
     return new AutoValue_WatermarkParameters.Builder<T>()
         .setCurrentWatermark(BoundedWindow.TIMESTAMP_MIN_VALUE)
         .setLastSavedWatermark(BoundedWindow.TIMESTAMP_MIN_VALUE)
-        .setLastUpdateTime(Instant.now())
-        .setWatermarkIdleDurationThreshold(STANDARD_WATERMARK_IDLE_DURATION_THRESHOLD);
+        .setLastUpdateTime(Instant.now());
   }
 
   @AutoValue.Builder
@@ -66,23 +61,4 @@ abstract static class Builder<T> {
 
     abstract WatermarkParameters<T> build();
   }
-
-  /**
-   * Create an instance of {@link WatermarkParameters} with a {@code SerializableFunction} to
-   * extract the event time.
-   */
-  static <T> WatermarkParameters<T> create(SerializableFunction<T, Instant> timestampFn) {
-    Preconditions.checkArgument(timestampFn != null, "timestampFn function is null");
-    return WatermarkParameters.<T>builder().setTimestampFn(timestampFn).build();
-  }
-
-  /**
-   * Specify the watermark idle duration to consider before advancing the watermark. The default
-   * watermark idle duration threshold is {@link #STANDARD_WATERMARK_IDLE_DURATION_THRESHOLD}.
-   */
-  WatermarkParameters<T> withWatermarkIdleDurationThreshold(Duration idleDurationThreshold) {
-    Preconditions.checkArgument(
-        idleDurationThreshold != null, "watermark idle duration threshold is null");
-    return toBuilder().setWatermarkIdleDurationThreshold(idleDurationThreshold).build();
-  }
 }

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

@@ -21,6 +21,7 @@
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Ordering;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
 import org.joda.time.Instant;
 
 /**
@@ -39,8 +40,13 @@
 class WatermarkPolicy<T> implements Serializable {
   private WatermarkParameters<T> watermarkParameters;
 
-  static <T> WatermarkPolicy<T> create(SerializableFunction<T, Instant> timestampFunction) {
-    return new WatermarkPolicy<T>(WatermarkParameters.<T>create(timestampFunction));
+  static <T> WatermarkPolicy<T> create(
+      SerializableFunction<T, Instant> timestampFunction, Duration watermarkIdleDurationThreshold) {
+    return new WatermarkPolicy<T>(
+        WatermarkParameters.<T>builder()
+            .setTimestampFn(timestampFunction)
+            .setWatermarkIdleDurationThreshold(watermarkIdleDurationThreshold)
+            .build());
   }
 
   private WatermarkPolicy(WatermarkParameters<T> watermarkParameters) {

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

@@ -95,6 +95,7 @@ private static UnboundedSolaceSource<Record> getSource(Read<Record> spec, TestPi
         configuration.getDeduplicateRecords(),
         spec.inferCoder(pipeline, configuration.getTypeDescriptor()),
         configuration.getTimestampFn(),
+        configuration.getWatermarkIdleDurationThreshold(),
         configuration.getParseFn());
   }
 
@@ -527,7 +528,7 @@ public void testCheckpointMarkSafety() throws Exception {
   @Test
   public void testDefaultCoder() {
     Coder<SolaceCheckpointMark> coder =
-        new UnboundedSolaceSource<>(null, null, null, 0, false, null, null, null)
+        new UnboundedSolaceSource<>(null, null, null, 0, false, null, null, null, null)
             .getCheckpointMarkCoder();
     CoderProperties.coderSerializable(coder);
   }