[feat] Remove custom deduplication logic in favor of beam Deduplicate. (#120)

Author: dpcollins-google

google-cloud-pubsublite/src/test/java/com/google/cloud/pubsublite/internal/wire/CommitterImplTest.java

@@ -185,10 +185,12 @@ public void multipleSentCompletedInOrder() {
   public void stopInCommitCallback() throws Exception {
     ApiFuture<Void> future = committer.commitOffset(Offset.of(10));
     CountDownLatch latch = new CountDownLatch(1);
-    ExtractStatus.addFailureHandler(future, (error) -> {
-      committer.stopAsync();
-      latch.countDown();
-    });
+    ExtractStatus.addFailureHandler(
+        future,
+        (error) -> {
+          committer.stopAsync();
+          latch.countDown();
+        });
     leakedResponseObserver.onError(Status.FAILED_PRECONDITION.asException());
     latch.await();
     assertFutureThrowsCode(future, Code.FAILED_PRECONDITION);

pubsublite-beam-io/pom.xml

@@ -28,7 +28,7 @@
     <dependency>
       <groupId>org.apache.beam</groupId>
       <artifactId>beam-sdks-java-core</artifactId>
-      <version>2.20.0</version>
+      <version>2.21.0</version>
     </dependency>
     <dependency>
       <groupId>com.google.code.findbugs</groupId>
@@ -104,7 +104,7 @@
     <dependency>
       <groupId>org.apache.beam</groupId>
       <artifactId>beam-runners-direct-java</artifactId>
-      <version>2.20.0</version>
+      <version>2.21.0</version>
       <scope>test</scope>
     </dependency>
     <dependency>

pubsublite-beam-io/src/main/java/com/google/cloud/pubsublite/beam/DeduplicationFn.java

@@ -23,6 +23,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import java.io.Serializable;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.transforms.Deduplicate;
 
 @AutoValue
 public abstract class UuidDeduplicationOptions implements Serializable {
@@ -44,7 +46,7 @@ public abstract class UuidDeduplicationOptions implements Serializable {
   // All parameters are optional.
   public abstract SerializableStatusFunction<SequencedMessage, Uuid> uuidExtractor();
 
-  public abstract DeduplicationFnOptions<Uuid> deduplicationFnOptions();
+  public abstract Deduplicate.KeyedValues<Uuid, SequencedMessage> deduplicate();
 
   // The number of partitions to hash values into.
   public abstract int hashPartitions();
@@ -53,8 +55,8 @@ public abstract class UuidDeduplicationOptions implements Serializable {
   public static Builder newBuilder() {
     Builder builder = new AutoValue_UuidDeduplicationOptions.Builder();
     builder.setUuidExtractor(DEFAULT_UUID_EXTRACTOR);
-    builder.setDeduplicationFnOptions(
-        DeduplicationFnOptions.<Uuid>newBuilder().setKeyCoder(Uuid.getCoder()).build());
+    builder.setDeduplicate(
+        Deduplicate.<Uuid, SequencedMessage>keyedValues().withTimeDomain(TimeDomain.EVENT_TIME));
     builder.setHashPartitions(DEFAULT_HASH_PARTITIONS);
     return builder;
   }
@@ -64,7 +66,8 @@ public abstract static class Builder {
     public abstract Builder setUuidExtractor(
         SerializableStatusFunction<SequencedMessage, Uuid> uuidExtractor);
 
-    public abstract Builder setDeduplicationFnOptions(DeduplicationFnOptions<Uuid> options);
+    public abstract Builder setDeduplicate(
+        Deduplicate.KeyedValues<Uuid, SequencedMessage> deduplicate);
 
     public abstract Builder setHashPartitions(int partitions);
 

pubsublite-beam-io/src/main/java/com/google/cloud/pubsublite/beam/DeduplicationFnOptions.java

@@ -17,10 +17,8 @@
 package com.google.cloud.pubsublite.beam;
 
 import com.google.cloud.pubsublite.SequencedMessage;
-import java.math.BigInteger;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.ProcessFunction;
 import org.apache.beam.sdk.transforms.Values;
 import org.apache.beam.sdk.values.KV;
@@ -48,22 +46,8 @@ public PCollection<SequencedMessage> expand(PCollection<SequencedMessage> input)
         input.apply(
             "MapUuids",
             MapElements.into(new TypeDescriptor<KV<Uuid, SequencedMessage>>() {}).via(mapWithKeys));
-    ProcessFunction<KV<Uuid, SequencedMessage>, KV<Integer, KV<Uuid, SequencedMessage>>>
-        mapWithHash =
-            kv ->
-                KV.of(
-                    new BigInteger(/*signum=*/ 1, kv.getKey().value().toByteArray())
-                        .mod(BigInteger.valueOf(options.hashPartitions()))
-                        .intValue(),
-                    kv);
-    PCollection<KV<Integer, KV<Uuid, SequencedMessage>>> hashedUuids =
-        uuidMapped.apply(
-            "HashUuids",
-            MapElements.into(new TypeDescriptor<KV<Integer, KV<Uuid, SequencedMessage>>>() {})
-                .via(mapWithHash));
     PCollection<KV<Uuid, SequencedMessage>> unique =
-        hashedUuids.apply(
-            "Deduplicate", ParDo.of(new DeduplicationFn<>(options.deduplicationFnOptions())));
+        uuidMapped.apply("Deduplicate", options.deduplicate());
     return unique.apply("StripUuids", Values.create());
   }
 }

pubsublite-beam-io/src/main/java/com/google/cloud/pubsublite/beam/UuidDeduplicationOptions.java

@@ -27,6 +27,7 @@
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.testing.TestStream;
+import org.apache.beam.sdk.transforms.Deduplicate;
 import org.apache.beam.sdk.values.PCollection;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
@@ -60,7 +61,7 @@ public void unrelatedUuidsProxied() {
         TestStream.create(new SequencedMessageCoder())
             .advanceWatermarkTo(START)
             .addElements(message1)
-            .advanceWatermarkTo(START.plus(DeduplicationFnOptions.DEFAULT_GC_DELAY.dividedBy(2)))
+            .advanceWatermarkTo(START.plus(Deduplicate.DEFAULT_DURATION.dividedBy(2)))
             .addElements(message2)
             .advanceWatermarkToInfinity();
     PCollection<SequencedMessage> results =
@@ -81,7 +82,7 @@ public void sameUuidsWithinWindowOnlyOne() {
         TestStream.create(new SequencedMessageCoder())
             .advanceWatermarkTo(START)
             .addElements(message)
-            .advanceWatermarkTo(START.plus(DeduplicationFnOptions.DEFAULT_GC_DELAY.dividedBy(2)))
+            .advanceWatermarkTo(START.plus(Deduplicate.DEFAULT_DURATION.dividedBy(2)))
             .advanceWatermarkToInfinity();
     PCollection<SequencedMessage> results =
         pipeline
@@ -101,8 +102,7 @@ public void sameUuidsAfterGcOutsideWindowHasBoth() {
         TestStream.create(new SequencedMessageCoder())
             .advanceWatermarkTo(START)
             .addElements(message1)
-            .advanceWatermarkTo(
-                START.plus(DeduplicationFnOptions.DEFAULT_GC_DELAY.plus(Duration.millis(1))))
+            .advanceWatermarkTo(START.plus(Deduplicate.DEFAULT_DURATION.plus(Duration.millis(1))))
             .addElements(message1)
             .advanceWatermarkToInfinity();
     PCollection<SequencedMessage> results =