[GOBBLIN-2040] Abstract comparable watermark (apache#3919)

ZihanLi58 · Zihan Li · web-flow · commit 36cdf2bb0550 · 2024-04-19T09:06:09.000-07:00
* address comments

* use connectionmanager when httpclient is not cloesable

* [GOBBLIN-2040]Abstract comparable watermark

* address comments

---------

Co-authored-by: Zihan Li &lt;zihli@zihli-mn2.linkedin.biz&gt;
diff --git a/gobblin-api/src/main/java/org/apache/gobblin/source/extractor/ComparableWatermark.java b/gobblin-api/src/main/java/org/apache/gobblin/source/extractor/ComparableWatermark.java
@@ -20,6 +20,8 @@
 /**
  * {@link Watermark} that is also {@link Comparable}.
  */
-public interface ComparableWatermark extends Watermark, Comparable<ComparableWatermark>{
+public interface ComparableWatermark<V extends Comparable<V>> extends Watermark, Comparable<ComparableWatermark>{
+
+  V getValue();
 
 }
diff --git a/gobblin-core-base/src/main/java/org/apache/gobblin/source/extractor/extract/LongWatermark.java b/gobblin-core-base/src/main/java/org/apache/gobblin/source/extractor/extract/LongWatermark.java
@@ -29,20 +29,29 @@
 import org.apache.gobblin.source.extractor.Watermark;
 
 import lombok.EqualsAndHashCode;
-import lombok.Getter;
 import lombok.Setter;
 import lombok.ToString;
 
+
+/**
+ * Long based {@link ComparableWatermark} implementation.
+ */
 @ToString
 @EqualsAndHashCode
-public class LongWatermark implements ComparableWatermark {
+public class LongWatermark implements ComparableWatermark<Long> {
 
   private static final Gson GSON = new Gson();
 
-  @Getter
   @Setter
   private long value;
 
+  @Override
+  // Returns a Long object due to Java generics' requirement for object types.
+  // The underlying variable is maintained as a primitive long to optimize performance for mathematical operations.
+  public Long getValue(){
+    return value;
+  }
+
   public LongWatermark(long value) {
     this.value = value;
   }
diff --git a/gobblin-core-base/src/test/java/org/apache/gobblin/writer/FineGrainedWatermarkTrackerTest.java b/gobblin-core-base/src/test/java/org/apache/gobblin/writer/FineGrainedWatermarkTrackerTest.java
@@ -98,7 +98,7 @@ private static void verifyCommitables(FineGrainedWatermarkTracker tracker, Sorte
     } else {
       Assert.assertEquals(uncommitted.size(), 1);
       CheckpointableWatermark uncommitable = uncommitted.get("default");
-      Assert.assertEquals(((LongWatermark) uncommitable.getWatermark()).getValue(), (long) holes.first());
+      Assert.assertEquals((long)((LongWatermark) uncommitable.getWatermark()).getValue(),  (long)holes.first());
     }
 
     Map<String, CheckpointableWatermark> commitables = tracker.getCommittableWatermarks();
@@ -109,9 +109,9 @@ private static void verifyCommitables(FineGrainedWatermarkTracker tracker, Sorte
       Assert.assertEquals(commitables.size(), 1);
       CheckpointableWatermark commitable = commitables.get("default");
       if (holes.isEmpty()) {
-        Assert.assertEquals(((LongWatermark) commitable.getWatermark()).getValue(), maxWatermark);
+        Assert.assertEquals((long)((LongWatermark) commitable.getWatermark()).getValue(), maxWatermark);
       } else {
-        Assert.assertEquals(((LongWatermark) commitable.getWatermark()).getValue(), holes.first() - 1);
+        Assert.assertEquals((long)((LongWatermark) commitable.getWatermark()).getValue(), holes.first() - 1);
       }
     }
 
diff --git a/gobblin-core-base/src/test/java/org/apache/gobblin/writer/WatermarkTrackerTest.java b/gobblin-core-base/src/test/java/org/apache/gobblin/writer/WatermarkTrackerTest.java
@@ -42,7 +42,7 @@ public void testSingleSource() {
     commits(watermarkTracker, "default", 0, 4, 5, 6);
 
     Assert.assertEquals(watermarkTracker.getCommittableWatermark("default").get().getSource(), "default");
-    Assert.assertEquals(((LongWatermark) watermarkTracker.getCommittableWatermark("default")
+    Assert.assertEquals((long)((LongWatermark) watermarkTracker.getCommittableWatermark("default")
         .get().getWatermark()).getValue(), 6L);
   }
 
@@ -54,10 +54,10 @@ public void testMultiSource() {
     commits(watermarkTracker, "other", 1, 3, 5, 7);
 
     Assert.assertEquals(watermarkTracker.getCommittableWatermark("default").get().getSource(), "default");
-    Assert.assertEquals(((LongWatermark) watermarkTracker.getCommittableWatermark("default")
+    Assert.assertEquals((long)((LongWatermark) watermarkTracker.getCommittableWatermark("default")
         .get().getWatermark()).getValue(), 6L);
     Assert.assertEquals(watermarkTracker.getCommittableWatermark("other").get().getSource(), "other");
-    Assert.assertEquals(((LongWatermark) watermarkTracker.getCommittableWatermark("other")
+    Assert.assertEquals((long)((LongWatermark) watermarkTracker.getCommittableWatermark("other")
         .get().getWatermark()).getValue(), 7L);
 
   }
diff --git a/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/watermark/StringWatermark.java b/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/watermark/StringWatermark.java
@@ -20,24 +20,24 @@
 import com.google.common.base.Preconditions;
 import com.google.gson.JsonElement;
 
+import lombok.Getter;
 import org.apache.gobblin.source.extractor.ComparableWatermark;
 import org.apache.gobblin.source.extractor.Watermark;
 import org.apache.gobblin.source.extractor.WatermarkSerializerHelper;
 
 import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
-import lombok.Getter;
 
 
 /**
  * String based {@link ComparableWatermark} implementation.
  */
 @AllArgsConstructor
 @EqualsAndHashCode
-public class StringWatermark implements ComparableWatermark {
+public class StringWatermark implements ComparableWatermark<String> {
 
   @Getter
-  String value;
+  private String value;
 
   @Override
   public int compareTo(ComparableWatermark other) {
diff --git a/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/conversion/hive/watermarker/PartitionLevelWatermarkerTest.java b/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/conversion/hive/watermarker/PartitionLevelWatermarkerTest.java
@@ -96,11 +96,11 @@ public void testExpectedHighWatermarkNoPreviousState() throws Exception {
     List<WorkUnit> workunits = Lists.newArrayList();
     watermarker.onGetWorkunitsEnd(workunits);
 
-    Assert.assertEquals(watermarker.getPreviousHighWatermark(part1).getValue(), 0l);
-    Assert.assertEquals(watermarker.getPreviousHighWatermark(table).getValue(), 0l);
+    Assert.assertEquals((long)watermarker.getPreviousHighWatermark(part1).getValue(), 0l);
+    Assert.assertEquals((long)watermarker.getPreviousHighWatermark(table).getValue(), 0l);
 
-    Assert.assertEquals(watermarker.getPreviousHighWatermark(part2).getValue(), 0l);
-    Assert.assertEquals(watermarker.getPreviousHighWatermark(table2).getValue(), 0l);
+    Assert.assertEquals((long)watermarker.getPreviousHighWatermark(part2).getValue(), 0l);
+    Assert.assertEquals((long)watermarker.getPreviousHighWatermark(table2).getValue(), 0l);
 
     Assert.assertEquals(workunits.size(), 2);
 
diff --git a/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/source/LoopingDatasetFinderSourceTest.java b/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/source/LoopingDatasetFinderSourceTest.java
@@ -205,17 +205,17 @@ public void testNonDrilldownDatasetState()
     List<LongWatermark> watermarks1 = new ArrayList<>();
     List<Dataset> datasets1 = new ArrayList<>();
     Assert.assertEquals(workUnits.get(0).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset1");
-    Assert.assertEquals(workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks1.add(workUnits.get(0).getExpectedHighWatermark(LongWatermark.class));
     datasets1.add(dataset1);
 
     Assert.assertEquals(workUnits.get(1).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset2");
-    Assert.assertEquals(workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks1.add(workUnits.get(1).getExpectedHighWatermark(LongWatermark.class));
     datasets1.add(dataset2);
 
     Assert.assertEquals(workUnits.get(2).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset3");
-    Assert.assertEquals(workUnits.get(2).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(2).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks1.add(workUnits.get(2).getExpectedHighWatermark(LongWatermark.class));
     datasets1.add(dataset3);
 
@@ -244,12 +244,12 @@ public void testNonDrilldownDatasetState()
 
     Assert.assertEquals(workUnits.size(), 3);
     Assert.assertEquals(workUnits.get(0).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset4");
-    Assert.assertEquals(workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks2.add(workUnits.get(0).getExpectedHighWatermark(LongWatermark.class));
     datasets2.add(dataset4);
 
     Assert.assertEquals(workUnits.get(1).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset5");
-    Assert.assertEquals(workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks2.add(workUnits.get(1).getExpectedHighWatermark(LongWatermark.class));
     datasets2.add(dataset5);
 
@@ -327,17 +327,17 @@ public void testDrilldownDatasetState()
 
     Assert.assertEquals(workUnits.size(), 4);
     Assert.assertEquals(workUnits.get(0).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset1");
-    Assert.assertEquals(workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks1.add(workUnits.get(0).getExpectedHighWatermark(LongWatermark.class));
     datasets1.add(dataset1);
 
     Assert.assertEquals(workUnits.get(1).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset2@p1");
-    Assert.assertEquals(workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks1.add(workUnits.get(1).getExpectedHighWatermark(LongWatermark.class));
     datasets1.add(new SimpleDatasetForTesting("dataset2@p1"));
 
     Assert.assertEquals(workUnits.get(2).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset2@p2");
-    Assert.assertEquals(workUnits.get(2).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(2).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks1.add(workUnits.get(2).getExpectedHighWatermark(LongWatermark.class));
     datasets1.add(new SimpleDatasetForTesting("dataset2@p2"));
 
@@ -367,17 +367,17 @@ public void testDrilldownDatasetState()
 
     Assert.assertEquals(workUnits.size(), 4);
     Assert.assertEquals(workUnits.get(0).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset2@p3");
-    Assert.assertEquals(workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks2.add(workUnits.get(0).getExpectedHighWatermark(LongWatermark.class));
     datasets2.add(new SimpleDatasetForTesting("dataset2@p3"));
 
     Assert.assertEquals(workUnits.get(1).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset3@p1");
-    Assert.assertEquals(workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(1).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks2.add(workUnits.get(1).getExpectedHighWatermark(LongWatermark.class));
     datasets2.add(new SimpleDatasetForTesting("dataset3@p1"));
 
     Assert.assertEquals(workUnits.get(2).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset3@p2");
-    Assert.assertEquals(workUnits.get(2).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(2).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks2.add(workUnits.get(2).getExpectedHighWatermark(LongWatermark.class));
     datasets2.add(new SimpleDatasetForTesting("dataset3@p2"));
 
@@ -404,7 +404,7 @@ public void testDrilldownDatasetState()
 
     Assert.assertEquals(workUnits.size(), 2);
     Assert.assertEquals(workUnits.get(0).getProp(ConfigurationKeys.DATASET_URN_KEY), "dataset3@p3");
-    Assert.assertEquals(workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
+    Assert.assertEquals((long)workUnits.get(0).getLowWatermark(LongWatermark.class).getValue(), 0);
     watermarks3.add(workUnits.get(0).getExpectedHighWatermark(LongWatermark.class));
     datasets3.add(new SimpleDatasetForTesting("dataset3@p3"));
 
diff --git a/gobblin-modules/gobblin-helix/src/test/java/org/apache/gobblin/runtime/StateStoreWatermarkStorageTest.java b/gobblin-modules/gobblin-helix/src/test/java/org/apache/gobblin/runtime/StateStoreWatermarkStorageTest.java
@@ -68,7 +68,7 @@ public void testPersistWatermarkStateToZk() throws IOException {
         ImmutableList.of("source"));
 
     Assert.assertEquals(watermarkMap.size(), 1);
-    Assert.assertEquals(((LongWatermark) watermarkMap.get("source").getWatermark()).getValue(), startTime);
+    Assert.assertEquals((long)((LongWatermark) watermarkMap.get("source").getWatermark()).getValue(), startTime);
   }
 
   @AfterClass

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,8 @@`
`20`	`20`	`/**`
`21`	`21`	`* {@link Watermark} that is also {@link Comparable}.`
`22`	`22`	`*/`
`23`		`-public interface ComparableWatermark extends Watermark, Comparable<ComparableWatermark>{`
	`23`	`+public interface ComparableWatermark<V extends Comparable<V>> extends Watermark, Comparable<ComparableWatermark>{`
	`24`	`+`
	`25`	`+ V getValue();`
`24`	`26`
`25`	`27`	`}`
Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ private static void verifyCommitables(FineGrainedWatermarkTracker tracker, Sorte`
`98`	`98`	`} else {`
`99`	`99`	`Assert.assertEquals(uncommitted.size(), 1);`
`100`	`100`	`CheckpointableWatermark uncommitable = uncommitted.get("default");`
`101`		`- Assert.assertEquals(((LongWatermark) uncommitable.getWatermark()).getValue(), (long) holes.first());`
	`101`	`+ Assert.assertEquals((long)((LongWatermark) uncommitable.getWatermark()).getValue(), (long)holes.first());`
`102`	`102`	`}`
`103`	`103`
`104`	`104`	`Map<String, CheckpointableWatermark> commitables = tracker.getCommittableWatermarks();`
`@@ -109,9 +109,9 @@ private static void verifyCommitables(FineGrainedWatermarkTracker tracker, Sorte`
`109`	`109`	`Assert.assertEquals(commitables.size(), 1);`
`110`	`110`	`CheckpointableWatermark commitable = commitables.get("default");`
`111`	`111`	`if (holes.isEmpty()) {`
`112`		`- Assert.assertEquals(((LongWatermark) commitable.getWatermark()).getValue(), maxWatermark);`
	`112`	`+ Assert.assertEquals((long)((LongWatermark) commitable.getWatermark()).getValue(), maxWatermark);`
`113`	`113`	`} else {`
`114`		`- Assert.assertEquals(((LongWatermark) commitable.getWatermark()).getValue(), holes.first() - 1);`
	`114`	`+ Assert.assertEquals((long)((LongWatermark) commitable.getWatermark()).getValue(), holes.first() - 1);`
`115`	`115`	`}`
`116`	`116`	`}`
`117`	`117`
Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ public void testPersistWatermarkStateToZk() throws IOException {`
`68`	`68`	`ImmutableList.of("source"));`
`69`	`69`
`70`	`70`	`Assert.assertEquals(watermarkMap.size(), 1);`
`71`		`- Assert.assertEquals(((LongWatermark) watermarkMap.get("source").getWatermark()).getValue(), startTime);`
	`71`	`+ Assert.assertEquals((long)((LongWatermark) watermarkMap.get("source").getWatermark()).getValue(), startTime);`
`72`	`72`	`}`
`73`	`73`
`74`	`74`	`@AfterClass`