fix: Change commit logic to use streams instead of single RPCs (#177)

* fix: Change commit logic to use streams instead of single RPCs

* fix: Change commit logic to use streams instead of single RPCs

Author: dpcollins-google

src/main/java/com/google/cloud/pubsublite/flink/PubsubLiteSource.java

@@ -22,6 +22,7 @@
 import com.google.cloud.pubsublite.flink.internal.source.enumerator.SplitDiscovery;
 import com.google.cloud.pubsublite.flink.internal.source.enumerator.SplitEnumeratorCheckpointSerializer;
 import com.google.cloud.pubsublite.flink.internal.source.enumerator.UniformPartitionAssigner;
+import com.google.cloud.pubsublite.flink.internal.source.reader.CommitterCache;
 import com.google.cloud.pubsublite.flink.internal.source.reader.PubsubLiteRecordEmitter;
 import com.google.cloud.pubsublite.flink.internal.source.reader.PubsubLiteSourceReader;
 import com.google.cloud.pubsublite.flink.internal.source.split.SubscriptionPartitionSplit;
@@ -74,10 +75,10 @@ public UserCodeClassLoader getUserCodeClassLoader() {
     SourceAssembler<OutputT> assembler = new SourceAssembler<>(settings);
     return new PubsubLiteSourceReader<>(
         new PubsubLiteRecordEmitter<>(),
-        assembler.getCursorClientRemoveThis(),
         assembler.getSplitReaderSupplier(),
         new Configuration(),
-        readerContext);
+        readerContext,
+        new CommitterCache(assembler::getCommitter));
   }
 
   @Override
@@ -89,7 +90,7 @@ public SplitEnumerator<SubscriptionPartitionSplit, SplitEnumeratorCheckpoint> cr
         UniformPartitionAssigner.create(),
         SingleSubscriptionSplitDiscovery.create(
             assembler.newAdminClient(),
-            assembler.getCursorClientRemoveThis(),
+            assembler.getCursorClient(),
             assembler.getTopicPath(),
             settings.subscriptionPath()));
   }
@@ -106,7 +107,7 @@ public SplitEnumerator<SubscriptionPartitionSplit, SplitEnumeratorCheckpoint> re
             checkpoint.getDiscovery(),
             assigner.listSplits(),
             assembler.newAdminClient(),
-            assembler.getCursorClientRemoveThis());
+            assembler.getCursorClient());
     return new PubsubLiteSplitEnumerator(enumContext, assigner, discovery);
   }
 

src/main/java/com/google/cloud/pubsublite/flink/internal/source/SourceAssembler.java

@@ -24,6 +24,7 @@
 import com.google.api.gax.rpc.ApiException;
 import com.google.cloud.pubsublite.AdminClient;
 import com.google.cloud.pubsublite.AdminClientSettings;
+import com.google.cloud.pubsublite.Partition;
 import com.google.cloud.pubsublite.SubscriptionPath;
 import com.google.cloud.pubsublite.TopicPath;
 import com.google.cloud.pubsublite.flink.PubsubLiteSourceSettings;
@@ -35,6 +36,8 @@
 import com.google.cloud.pubsublite.internal.BlockingPullSubscriberImpl;
 import com.google.cloud.pubsublite.internal.CursorClient;
 import com.google.cloud.pubsublite.internal.CursorClientSettings;
+import com.google.cloud.pubsublite.internal.wire.Committer;
+import com.google.cloud.pubsublite.internal.wire.CommitterSettings;
 import com.google.cloud.pubsublite.internal.wire.PubsubContext;
 import com.google.cloud.pubsublite.internal.wire.PubsubContext.Framework;
 import com.google.cloud.pubsublite.internal.wire.RoutingMetadata;
@@ -135,16 +138,26 @@ private CursorServiceClient newCursorClient() throws ApiException {
     }
   }
 
-  private CursorServiceClient getCursorClient() {
+  private CursorServiceClient getCursorServiceClient() {
     return CURSOR_CLIENTS.computeIfAbsent(settings.subscriptionPath(), path -> newCursorClient());
   }
 
-  /** TODO(dpcollins): Remove this */
-  public CursorClient getCursorClientRemoveThis() {
+  public Committer getCommitter(Partition partition) {
+    CursorServiceClient client = getCursorServiceClient();
+    return CommitterSettings.newBuilder()
+        .setSubscriptionPath(settings.subscriptionPath())
+        .setPartition(partition)
+        .setStreamFactory(
+            responseObserver -> client.streamingCommitCursorCallable().splitCall(responseObserver))
+        .build()
+        .instantiate();
+  }
+
+  public CursorClient getCursorClient() {
     return CursorClient.create(
         CursorClientSettings.newBuilder()
             .setRegion(settings.subscriptionPath().location().extractRegion())
-            .setServiceClient(getCursorClient())
+            .setServiceClient(getCursorServiceClient())
             .build());
   }
 

src/main/java/com/google/cloud/pubsublite/flink/internal/source/reader/CheckpointCursorCommitter.java

@@ -18,7 +18,6 @@
 import com.google.api.core.ApiFutureCallback;
 import com.google.api.core.ApiFutures;
 import com.google.cloud.pubsublite.flink.internal.source.split.SubscriptionPartitionSplit;
-import com.google.cloud.pubsublite.internal.CursorClient;
 import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
@@ -42,10 +41,10 @@ public class CheckpointCursorCommitter implements AutoCloseable {
   private final LinkedHashMap<Long, List<SubscriptionPartitionSplit>> checkpoints =
       new LinkedHashMap<>();
 
-  private final CursorClient cursorCommitter;
+  private final CommitterFactory factory;
 
-  public CheckpointCursorCommitter(CursorClient cursorCommitter) {
-    this.cursorCommitter = cursorCommitter;
+  public CheckpointCursorCommitter(CommitterFactory factory) {
+    this.factory = factory;
   }
 
   public synchronized void addCheckpoint(
@@ -60,7 +59,7 @@ public synchronized void addCheckpoint(
 
   private void commitCursor(SubscriptionPartitionSplit split) {
     ApiFutures.addCallback(
-        cursorCommitter.commitCursor(split.subscriptionPath(), split.partition(), split.start()),
+        factory.getCommitter(split.partition()).commitOffset(split.start()),
         new ApiFutureCallback<Void>() {
           @Override
           public void onFailure(Throwable throwable) {
@@ -98,6 +97,6 @@ public synchronized void notifySplitFinished(Collection<SubscriptionPartitionSpl
 
   @Override
   public void close() {
-    cursorCommitter.close();
+    factory.close();
   }
 }

src/main/java/com/google/cloud/pubsublite/flink/internal/source/reader/CommitterCache.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.pubsublite.flink.internal.source.reader;
+
+import static com.google.cloud.pubsublite.internal.wire.ApiServiceUtils.blockingShutdown;
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+
+import com.google.api.core.ApiService.Listener;
+import com.google.api.core.ApiService.State;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.internal.wire.Committer;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import java.util.HashMap;
+import java.util.function.Function;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class CommitterCache implements CommitterFactory {
+  private static final Logger LOG = LoggerFactory.getLogger(CommitterCache.class);
+  private final Function<Partition, Committer> underlying;
+
+  @GuardedBy("this")
+  private final HashMap<Partition, Committer> partitionCommitters = new HashMap<>();
+
+  public CommitterCache(Function<Partition, Committer> underlying) {
+    this.underlying = underlying;
+  }
+
+  @Override
+  public synchronized Committer getCommitter(Partition partition) {
+    return partitionCommitters.computeIfAbsent(
+        partition,
+        p -> {
+          Committer newCommitter = underlying.apply(p);
+          newCommitter.addListener(
+              new Listener() {
+                @Override
+                public void failed(State from, Throwable failure) {
+                  LOG.info("Committer failed.", failure);
+                  remove(p);
+                }
+
+                @Override
+                public void terminated(State from) {
+                  remove(p);
+                }
+              },
+              directExecutor());
+          newCommitter.startAsync().awaitRunning();
+          return newCommitter;
+        });
+  }
+
+  @Override
+  public synchronized void close() {
+    blockingShutdown(partitionCommitters.values());
+  }
+
+  private synchronized void remove(Partition partition) {
+    partitionCommitters.remove(partition);
+  }
+}

src/main/java/com/google/cloud/pubsublite/flink/internal/source/reader/CommitterFactory.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.pubsublite.flink.internal.source.reader;
+
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.internal.wire.Committer;
+import java.io.Closeable;
+
+public interface CommitterFactory extends Closeable {
+  Committer getCommitter(Partition partition);
+
+  @Override
+  void close();
+}

src/main/java/com/google/cloud/pubsublite/flink/internal/source/reader/PubsubLiteSourceReader.java

@@ -17,7 +17,6 @@
 
 import com.google.cloud.pubsublite.flink.internal.source.split.SubscriptionPartitionSplit;
 import com.google.cloud.pubsublite.flink.internal.source.split.SubscriptionPartitionSplitState;
-import com.google.cloud.pubsublite.internal.CursorClient;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
@@ -34,10 +33,10 @@ public class PubsubLiteSourceReader<T>
 
   public PubsubLiteSourceReader(
       RecordEmitter<Record<T>, T, SubscriptionPartitionSplitState> recordEmitter,
-      CursorClient cursorCommitter,
       Supplier<SplitReader<Record<T>, SubscriptionPartitionSplit>> splitReaderSupplier,
       Configuration config,
-      SourceReaderContext context) {
+      SourceReaderContext context,
+      CommitterFactory cursorCommitter) {
     super(splitReaderSupplier, recordEmitter, config, context);
     this.checkpointCursorCommitter = new CheckpointCursorCommitter(cursorCommitter);
   }

src/test/java/com/google/cloud/pubsublite/flink/internal/source/reader/CheckpointCursorCommitterTest.java

@@ -27,7 +27,7 @@
 import com.google.cloud.pubsublite.Offset;
 import com.google.cloud.pubsublite.Partition;
 import com.google.cloud.pubsublite.flink.internal.source.split.SubscriptionPartitionSplit;
-import com.google.cloud.pubsublite.internal.CursorClient;
+import com.google.cloud.pubsublite.internal.wire.Committer;
 import com.google.common.collect.ImmutableList;
 import org.junit.Before;
 import org.junit.Test;
@@ -37,14 +37,15 @@
 
 @RunWith(MockitoJUnitRunner.class)
 public class CheckpointCursorCommitterTest {
-  @Mock CursorClient mockCursorClient;
+  @Mock CommitterFactory mockCommitterFactory;
+  @Mock Committer mockCommitter;
   CheckpointCursorCommitter cursorCommitter;
 
   @Before
   public void setUp() {
-    cursorCommitter = new CheckpointCursorCommitter(mockCursorClient);
-    when(mockCursorClient.commitCursor(any(), any(), any()))
-        .thenReturn(ApiFutures.immediateFuture(null));
+    cursorCommitter = new CheckpointCursorCommitter(mockCommitterFactory);
+    when(mockCommitterFactory.getCommitter(any())).thenReturn(mockCommitter);
+    when(mockCommitter.commitOffset(any())).thenReturn(ApiFutures.immediateFuture(null));
   }
 
   public static SubscriptionPartitionSplit splitFromPartition(Partition partition) {
@@ -57,25 +58,25 @@ public void testFinishedSplits() {
     cursorCommitter.notifySplitFinished(ImmutableList.of(split));
     cursorCommitter.addCheckpoint(1, ImmutableList.of());
     cursorCommitter.notifyCheckpointComplete(1);
-    verify(mockCursorClient)
-        .commitCursor(split.subscriptionPath(), split.partition(), split.start());
+    verify(mockCommitterFactory).getCommitter(split.partition());
+    verify(mockCommitter).commitOffset(split.start());
   }
 
   @Test
   public void testCheckpointCommitted() {
     SubscriptionPartitionSplit split = splitFromPartition(examplePartition());
     cursorCommitter.addCheckpoint(1, ImmutableList.of(split));
     cursorCommitter.notifyCheckpointComplete(1);
-    verify(mockCursorClient)
-        .commitCursor(split.subscriptionPath(), split.partition(), split.start());
+    verify(mockCommitterFactory).getCommitter(split.partition());
+    verify(mockCommitter).commitOffset(split.start());
   }
 
   @Test
   public void testUnknownCheckpoint() {
     SubscriptionPartitionSplit split = splitFromPartition(examplePartition());
     cursorCommitter.addCheckpoint(1, ImmutableList.of(split));
     cursorCommitter.notifyCheckpointComplete(4);
-    verifyNoInteractions(mockCursorClient);
+    verifyNoInteractions(mockCommitter);
   }
 
   @Test
@@ -88,15 +89,15 @@ public void testIntermediateCheckpointSkipped() {
 
     // Checkpoint 1 is committed, removing checkpoint 2
     cursorCommitter.notifyCheckpointComplete(1);
-    verify(mockCursorClient)
-        .commitCursor(split1.subscriptionPath(), split1.partition(), split1.start());
+    verify(mockCommitterFactory).getCommitter(split1.partition());
+    verify(mockCommitter).commitOffset(split1.start());
     cursorCommitter.notifyCheckpointComplete(2);
-    verifyNoMoreInteractions(mockCursorClient);
+    verifyNoMoreInteractions(mockCommitter);
   }
 
   @Test
   public void testClose() {
     cursorCommitter.close();
-    verify(mockCursorClient).close();
+    verify(mockCommitterFactory).close();
   }
 }

src/test/java/com/google/cloud/pubsublite/flink/internal/source/reader/PubsubLiteSourceReaderTest.java

@@ -30,7 +30,7 @@
 import com.google.cloud.pubsublite.flink.PubsubLiteDeserializationSchema;
 import com.google.cloud.pubsublite.flink.internal.source.split.SubscriptionPartitionSplit;
 import com.google.cloud.pubsublite.internal.BlockingPullSubscriber;
-import com.google.cloud.pubsublite.internal.CursorClient;
+import com.google.cloud.pubsublite.internal.wire.Committer;
 import com.google.cloud.pubsublite.proto.Cursor;
 import com.google.cloud.pubsublite.proto.PubSubMessage;
 import com.google.common.collect.ImmutableList;
@@ -51,7 +51,9 @@
 @RunWith(MockitoJUnitRunner.class)
 public class PubsubLiteSourceReaderTest {
   @Mock PullSubscriberFactory mockFactory;
-  @Mock CursorClient mockCursorClient;
+  @Mock CommitterFactory mockCommitterFactory;
+  @Mock Committer mockCommitter0;
+  @Mock Committer mockCommitter1;
 
   @Mock(answer = RETURNS_DEEP_STUBS)
   SourceReaderContext mockContext;
@@ -83,16 +85,18 @@ public void setUp() {
     reader =
         new PubsubLiteSourceReader<>(
             new PubsubLiteRecordEmitter<>(),
-            mockCursorClient,
             () ->
                 new DeserializingSplitReader<>(
                     new MessageSplitReader(mockFactory),
                     PubsubLiteDeserializationSchema.dataOnly(new SimpleStringSchema()),
                     MessageTimestampExtractor.publishTimeExtractor()),
             new Configuration(),
-            mockContext);
-    when(mockCursorClient.commitCursor(any(), any(), any()))
-        .thenReturn(ApiFutures.immediateFuture(null));
+            mockContext,
+            mockCommitterFactory);
+    when(mockCommitterFactory.getCommitter(Partition.of(0))).thenReturn(mockCommitter0);
+    when(mockCommitterFactory.getCommitter(Partition.of(1))).thenReturn(mockCommitter1);
+    when(mockCommitter0.commitOffset(any())).thenReturn(ApiFutures.immediateFuture(null));
+    when(mockCommitter1.commitOffset(any())).thenReturn(ApiFutures.immediateFuture(null));
   }
 
   @Test(timeout = 1000)
@@ -110,8 +114,8 @@ public void testReader() throws Exception {
 
     reader.snapshotState(1);
     reader.notifyCheckpointComplete(1);
-    verify(mockCursorClient).commitCursor(exampleSubscriptionPath(), Partition.of(0), Offset.of(2));
-    verify(mockCursorClient).commitCursor(exampleSubscriptionPath(), Partition.of(1), Offset.of(3));
+    verify(mockCommitter0).commitOffset(Offset.of(2));
+    verify(mockCommitter1).commitOffset(Offset.of(3));
 
     while (output.getEmittedRecords().size() < 6) {
       reader.pollNext(output);