add filter to watermark query

Author: changliiu

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/cache/AsyncWatermarkCache.java

@@ -18,9 +18,10 @@
 package org.apache.beam.sdk.io.gcp.spanner.changestreams.cache;
 
 import com.google.cloud.Timestamp;
-import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
@@ -40,23 +41,47 @@ public class AsyncWatermarkCache implements WatermarkCache {
 
   private static final String THREAD_NAME_FORMAT = "watermark_loading_thread_%d";
   private static final Object MIN_WATERMARK_KEY = new Object();
-  private final LoadingCache<Object, Optional<Timestamp>> cache;
+  private final LoadingCache<Object, Timestamp> cache;
+
+  // This is to cache the result of getUnfinishedMinWatermark query and filter the query in the next
+  // run. For the initial query, the value of this cache is min timestamp. If there is no partition
+  // in the metadata table, then this cache will not be updated. If the getUnfinishedMinWatermark
+  // query fails or times out, then this cache will not be updated.
+  // Note that, all the reload operations on this key are serialized due to use of the single
+  // threaded async reloading executor.
+  private Timestamp lastCachedMinWatermark = Timestamp.MIN_VALUE;
+
+  private final Lock lock = new ReentrantLock();
 
   public AsyncWatermarkCache(PartitionMetadataDao dao, Duration refreshRate) {
     this.cache =
         CacheBuilder.newBuilder()
             .refreshAfterWrite(java.time.Duration.ofMillis(refreshRate.getMillis()))
             .build(
                 CacheLoader.asyncReloading(
-                    CacheLoader.from(key -> Optional.ofNullable(dao.getUnfinishedMinWatermark())),
+                    CacheLoader.from(
+                        key -> {
+                          Timestamp unfinishedMinTimes =
+                              dao.getUnfinishedMinWatermarkFrom(lastCachedMinWatermark);
+                          if (unfinishedMinTimes != null
+                              && lastCachedMinWatermark.compareTo(unfinishedMinTimes) < 0) {
+                            lock.lock();
+                            try {
+                              lastCachedMinWatermark = unfinishedMinTimes;
+                            } finally {
+                              lock.unlock();
+                            }
+                          }
+                          return unfinishedMinTimes;
+                        }),
                     Executors.newSingleThreadExecutor(
                         new ThreadFactoryBuilder().setNameFormat(THREAD_NAME_FORMAT).build())));
   }
 
   @Override
   public @Nullable Timestamp getUnfinishedMinWatermark() {
     try {
-      return cache.get(MIN_WATERMARK_KEY).orElse(null);
+      return cache.get(MIN_WATERMARK_KEY);
     } catch (ExecutionException e) {
       throw new RuntimeException(e);
     }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/cache/NoOpWatermarkCache.java

@@ -35,6 +35,6 @@ public NoOpWatermarkCache(PartitionMetadataDao dao) {
 
   @Override
   public @Nullable Timestamp getUnfinishedMinWatermark() {
-    return dao.getUnfinishedMinWatermark();
+    return dao.getUnfinishedMinWatermarkFrom(Timestamp.MIN_VALUE);
   }
 }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java

@@ -178,47 +178,56 @@ public List<String> findAllTableIndexes() {
    *
    * @return the earliest partition watermark which is not in a {@link State#FINISHED} state.
    */
-  public @Nullable Timestamp getUnfinishedMinWatermark() {
+  public @Nullable Timestamp getUnfinishedMinWatermarkFrom(Timestamp sinceTimestamp) {
     Statement statement;
+    final String minWatermark = "min_watermark";
     if (this.isPostgres()) {
       statement =
           Statement.newBuilder(
-                  "SELECT \""
+                  "SELECT MIN(\""
                       + COLUMN_WATERMARK
-                      + "\" FROM \""
+                      + "\") as "
+                      + minWatermark
+                      + " FROM \""
                       + metadataTableName
                       + "\" WHERE \""
                       + COLUMN_STATE
                       + "\" != $1"
-                      + " ORDER BY \""
+                      + " AND \""
                       + COLUMN_WATERMARK
-                      + "\" ASC LIMIT 1")
+                      + "\" >= $2")
               .bind("p1")
               .to(State.FINISHED.name())
+              .bind("p2")
+              .to(sinceTimestamp)
               .build();
     } else {
       statement =
           Statement.newBuilder(
-                  "SELECT "
+                  "SELECT MIN("
                       + COLUMN_WATERMARK
+                      + ") as "
+                      + minWatermark
                       + " FROM "
                       + metadataTableName
                       + " WHERE "
                       + COLUMN_STATE
                       + " != @state"
-                      + " ORDER BY "
+                      + " AND "
                       + COLUMN_WATERMARK
-                      + " ASC LIMIT 1")
+                      + " >= @since;")
               .bind("state")
               .to(State.FINISHED.name())
+              .bind("since")
+              .to(sinceTimestamp)
               .build();
     }
     try (ResultSet resultSet =
         databaseClient
             .singleUse()
-            .executeQuery(statement, Options.tag("query=getUnfinishedMinWatermark"))) {
+            .executeQuery(statement, Options.tag("query=getUnfinishedMinWatermarkFrom"))) {
       if (resultSet.next()) {
-        return resultSet.getTimestamp(COLUMN_WATERMARK);
+        return resultSet.getTimestamp(minWatermark);
       }
       return null;
     }

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java

@@ -460,19 +460,27 @@ private void mockGetPartitionsAfter(Timestamp timestamp, ResultSet getPartitionR
   }
 
   private void mockGetWatermark(Timestamp watermark) {
+    final String minWatermark = "min_watermark";
+    // The query needs to sync with getUnfinishedMinWatermark() in PartitionMetadataDao file.
     Statement watermarkStatement =
         Statement.newBuilder(
-                "SELECT Watermark FROM my-metadata-table WHERE State != @state ORDER BY Watermark ASC LIMIT 1")
+                "SELECT MIN(Watermark) as "
+                    + minWatermark
+                    + " FROM my-metadata-table "
+                    + "WHERE State != @state "
+                    + "AND Watermark >= @since;")
             .bind("state")
             .to(State.FINISHED.name())
+            .bind("since")
+            .to(Timestamp.MIN_VALUE)
             .build();
     ResultSetMetadata watermarkResultSetMetadata =
         ResultSetMetadata.newBuilder()
             .setRowType(
                 StructType.newBuilder()
                     .addFields(
                         Field.newBuilder()
-                            .setName("Watermark")
+                            .setName(minWatermark)
                             .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP).build())
                             .build())
                     .build())