feat: Reset Consumer upon out-of-band seek (#172)

Resets the Consumer state to handle out-of-band seeks pushed from the server.

.readme-partials.yaml

@@ -80,3 +80,9 @@ about: |
   - Producers operate on a single topic, and Consumers on a single subscription.
   - ProducerRecord may not specify partition explicitly.
   - Consumers may not dynamically create consumer groups (subscriptions).
+
+  Note:
+  - In order to use Pub/Sub Lite [seek operations](https://cloud.google.com/pubsub/lite/docs/seek),
+    Consumers must have auto-commit enabled. Consumer seek methods are client-initiated, whereas
+    Pub/Sub Lite seek operations are initiated out-of-band and pushed to Consumers. Both types of
+    seeks should not be used concurrently, as they would interfere with one another.

src/main/java/com/google/cloud/pubsublite/kafka/ConsumerSettings.java

@@ -128,7 +128,7 @@ public Consumer<byte[], byte[]> instantiate() throws ApiException {
             }
           };
       PullSubscriberFactory pullSubscriberFactory =
-          (partition, initialSeek) -> {
+          (partition, initialSeek, resetHandler) -> {
             SubscriberFactory subscriberFactory =
                 consumer -> {
                   try {
@@ -145,6 +145,7 @@ public Consumer<byte[], byte[]> instantiate() throws ApiException {
                                         RoutingMetadata.of(subscriptionPath(), partition),
                                         SubscriberServiceSettings.newBuilder()))))
                         .setInitialLocation(initialSeek)
+                        .setResetHandler(resetHandler)
                         .build();
                   } catch (Throwable t) {
                     throw toCanonical(t).underlying;

src/main/java/com/google/cloud/pubsublite/kafka/PullSubscriberFactory.java

@@ -19,10 +19,12 @@
 import com.google.cloud.pubsublite.Partition;
 import com.google.cloud.pubsublite.internal.BlockingPullSubscriber;
 import com.google.cloud.pubsublite.internal.CheckedApiException;
+import com.google.cloud.pubsublite.internal.wire.SubscriberResetHandler;
 import com.google.cloud.pubsublite.proto.SeekRequest;
 
 /** A factory for making new PullSubscribers for a given partition of a subscription. */
 interface PullSubscriberFactory {
-  BlockingPullSubscriber newPullSubscriber(Partition partition, SeekRequest initial)
+  BlockingPullSubscriber newPullSubscriber(
+      Partition partition, SeekRequest initial, SubscriberResetHandler resetHandler)
       throws CheckedApiException;
 }

src/main/java/com/google/cloud/pubsublite/kafka/SinglePartitionSubscriber.java

@@ -0,0 +1,168 @@
+/*
+ * Copyright 2021 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.kafka;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.SequencedMessage;
+import com.google.cloud.pubsublite.internal.BlockingPullSubscriber;
+import com.google.cloud.pubsublite.internal.CheckedApiException;
+import com.google.cloud.pubsublite.internal.CloseableMonitor;
+import com.google.cloud.pubsublite.internal.ProxyService;
+import com.google.cloud.pubsublite.internal.wire.Committer;
+import com.google.cloud.pubsublite.proto.SeekRequest;
+import com.google.common.collect.Iterables;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import java.util.ArrayDeque;
+import java.util.Optional;
+
+/** Pulls messages and manages commits for a single partition of a subscription. */
+class SinglePartitionSubscriber extends ProxyService {
+  private final PullSubscriberFactory subscriberFactory;
+  private final Partition partition;
+  private final Committer committer;
+  private final boolean enableReset;
+
+  private final CloseableMonitor monitor = new CloseableMonitor();
+
+  @GuardedBy("monitor.monitor")
+  private BlockingPullSubscriber subscriber;
+
+  @GuardedBy("monitor.monitor")
+  private boolean needsCommitting = false;
+
+  @GuardedBy("monitor.monitor")
+  private Optional<Offset> lastReceived = Optional.empty();
+
+  SinglePartitionSubscriber(
+      PullSubscriberFactory subscriberFactory,
+      Partition partition,
+      SeekRequest initialSeek,
+      Committer committer,
+      boolean enableReset)
+      throws CheckedApiException {
+    this.subscriberFactory = subscriberFactory;
+    this.partition = partition;
+    this.committer = committer;
+    this.enableReset = enableReset;
+    this.subscriber =
+        subscriberFactory.newPullSubscriber(partition, initialSeek, this::onSubscriberReset);
+    addServices(committer);
+  }
+
+  // ProxyService implementation.
+  @Override
+  protected void start() {}
+
+  @Override
+  protected void stop() {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      subscriber.close();
+    }
+  }
+
+  @Override
+  protected void handlePermanentError(CheckedApiException error) {
+    stop();
+  }
+
+  /** Executes a client-initiated seek. */
+  void clientSeek(SeekRequest request) throws CheckedApiException {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      subscriber.close();
+      subscriber = subscriberFactory.newPullSubscriber(partition, request, this::onSubscriberReset);
+    }
+  }
+
+  ApiFuture<Void> onData() {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      return subscriber.onData();
+    }
+  }
+
+  @GuardedBy("monitor.monitor")
+  private ArrayDeque<SequencedMessage> pullMessages() throws CheckedApiException {
+    ArrayDeque<SequencedMessage> messages = new ArrayDeque<>();
+    for (Optional<SequencedMessage> message = subscriber.messageIfAvailable();
+        message.isPresent();
+        message = subscriber.messageIfAvailable()) {
+      messages.add(message.get());
+    }
+    return messages;
+  }
+
+  /** Pulls all available messages. */
+  ArrayDeque<SequencedMessage> getMessages() throws CheckedApiException {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      ArrayDeque<SequencedMessage> messages = pullMessages();
+      if (!messages.isEmpty()) {
+        lastReceived = Optional.of(Iterables.getLast(messages).offset());
+        needsCommitting = true;
+      }
+      return messages;
+    }
+  }
+
+  Optional<Long> position() {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      return lastReceived.map(lastReceived -> lastReceived.value() + 1);
+    }
+  }
+
+  /** Executes a client-initiated commit. */
+  ApiFuture<Void> commitOffset(Offset offset) {
+    return committer.commitOffset(offset);
+  }
+
+  /** Auto-commits the offset of the last received message. */
+  Optional<ApiFuture<Offset>> autoCommit() {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      if (!needsCommitting) return Optional.empty();
+      checkState(lastReceived.isPresent());
+      needsCommitting = false;
+      // The Pub/Sub Lite commit offset is one more than the last received.
+      Offset toCommit = Offset.of(lastReceived.get().value() + 1);
+      return Optional.of(
+          ApiFutures.transform(
+              committer.commitOffset(toCommit),
+              ignored -> toCommit,
+              MoreExecutors.directExecutor()));
+    }
+  }
+
+  private boolean onSubscriberReset() throws CheckedApiException {
+    if (!enableReset) {
+      return false;
+    }
+
+    // Handle an out-of-band seek notification from the server. There must be no pending commits
+    // after this function returns.
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      // Discard undelivered messages.
+      pullMessages();
+      // Prevent further auto-commits until post-seek messages are received.
+      needsCommitting = false;
+    }
+    committer.waitUntilEmpty();
+    return true;
+  }
+}