feat: Implement Consumer.endOffsets (#102)

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

@@ -29,6 +29,8 @@
 import com.google.cloud.pubsublite.internal.BufferingPullSubscriber;
 import com.google.cloud.pubsublite.internal.CursorClient;
 import com.google.cloud.pubsublite.internal.CursorClientSettings;
+import com.google.cloud.pubsublite.internal.TopicStatsClient;
+import com.google.cloud.pubsublite.internal.TopicStatsClientSettings;
 import com.google.cloud.pubsublite.internal.wire.AssignerFactory;
 import com.google.cloud.pubsublite.internal.wire.AssignerSettings;
 import com.google.cloud.pubsublite.internal.wire.CommitterSettings;
@@ -148,12 +150,21 @@ public Consumer<byte[], byte[]> instantiate() throws ApiException {
 
       CursorClient cursorClient =
           CursorClient.create(CursorClientSettings.newBuilder().setRegion(zone.region()).build());
+      TopicStatsClient topicStatsClient =
+          TopicStatsClient.create(
+              TopicStatsClientSettings.newBuilder().setRegion(zone.region()).build());
       SharedBehavior shared =
           new SharedBehavior(
               AdminClient.create(
                   AdminClientSettings.newBuilder().setRegion(topic.location().region()).build()));
       return new PubsubLiteConsumer(
-          subscriptionPath(), topic, shared, consumerFactory, assignerFactory, cursorClient);
+          subscriptionPath(),
+          topic,
+          shared,
+          consumerFactory,
+          assignerFactory,
+          cursorClient,
+          topicStatsClient);
     } catch (Exception e) {
       throw toCanonical(e).underlying;
     }

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

@@ -18,6 +18,7 @@
 
 import static com.google.cloud.pubsublite.kafka.KafkaExceptionUtils.toKafka;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutureCallback;
 import com.google.api.core.ApiFutures;
 import com.google.api.gax.rpc.ApiException;
@@ -26,6 +27,7 @@
 import com.google.cloud.pubsublite.SubscriptionPath;
 import com.google.cloud.pubsublite.TopicPath;
 import com.google.cloud.pubsublite.internal.CursorClient;
+import com.google.cloud.pubsublite.internal.TopicStatsClient;
 import com.google.cloud.pubsublite.internal.wire.Assigner;
 import com.google.cloud.pubsublite.internal.wire.AssignerFactory;
 import com.google.cloud.pubsublite.internal.wire.PartitionAssignmentReceiver;
@@ -74,6 +76,7 @@ class PubsubLiteConsumer implements Consumer<byte[], byte[]> {
   private final ConsumerFactory consumerFactory;
   private final AssignerFactory assignerFactory;
   private final CursorClient cursorClient;
+  private final TopicStatsClient topicStatsClient;
   private Optional<Assigner> assigner = Optional.empty();
   private Optional<SingleSubscriptionConsumer> consumer = Optional.empty();
 
@@ -83,13 +86,15 @@ class PubsubLiteConsumer implements Consumer<byte[], byte[]> {
       SharedBehavior shared,
       ConsumerFactory consumerFactory,
       AssignerFactory assignerFactory,
-      CursorClient cursorClient) {
+      CursorClient cursorClient,
+      TopicStatsClient topicStatsClient) {
     this.subscriptionPath = subscriptionPath;
     this.topicPath = topicPath;
     this.shared = shared;
     this.consumerFactory = consumerFactory;
     this.assignerFactory = assignerFactory;
     this.cursorClient = cursorClient;
+    this.topicStatsClient = topicStatsClient;
   }
 
   private TopicPartition toTopicPartition(Partition partition) {
@@ -490,8 +495,25 @@ public Map<TopicPartition, Long> endOffsets(Collection<TopicPartition> collectio
   @Override
   public Map<TopicPartition, Long> endOffsets(
       Collection<TopicPartition> collection, Duration duration) {
-    throw new UnsupportedVersionException(
-        "Pub/Sub Lite does not support Consumer backlog introspection.");
+    try {
+      Map<TopicPartition, ApiFuture<Cursor>> cursors =
+          collection.stream()
+              .collect(
+                  Collectors.toMap(
+                      topicPartition -> topicPartition,
+                      topicPartition ->
+                          topicStatsClient.computeHeadCursor(
+                              topicPath, checkTopicGetPartition(topicPartition))));
+      ApiFutures.allAsList(cursors.values()).get(duration.toMillis(), TimeUnit.MILLISECONDS);
+
+      ImmutableMap.Builder<TopicPartition, Long> output = ImmutableMap.builder();
+      for (Map.Entry<TopicPartition, ApiFuture<Cursor>> entry : cursors.entrySet()) {
+        output.put(entry.getKey(), entry.getValue().get().getOffset());
+      }
+      return output.build();
+    } catch (Throwable t) {
+      throw toKafka(t);
+    }
   }
 
   @Override
@@ -511,6 +533,12 @@ public void close(Duration timeout) {
     } catch (Exception e) {
       logger.atSevere().withCause(e).log("Error closing cursor client during Consumer shutdown.");
     }
+    try {
+      topicStatsClient.close();
+    } catch (Exception e) {
+      logger.atSevere().withCause(e).log(
+          "Error closing topic stats client during Consumer shutdown.");
+    }
     try {
       shared.close();
     } catch (Exception e) {

src/test/java/com/google/cloud/pubsublite/kafka/PubsubLiteConsumerTest.java

@@ -39,6 +39,7 @@
 import com.google.cloud.pubsublite.TopicName;
 import com.google.cloud.pubsublite.TopicPath;
 import com.google.cloud.pubsublite.internal.CursorClient;
+import com.google.cloud.pubsublite.internal.TopicStatsClient;
 import com.google.cloud.pubsublite.internal.testing.UnitTestExamples;
 import com.google.cloud.pubsublite.internal.wire.Assigner;
 import com.google.cloud.pubsublite.internal.wire.AssignerFactory;
@@ -54,6 +55,7 @@
 import com.google.common.reflect.ImmutableTypeToInstanceMap;
 import java.time.Duration;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Pattern;
 import org.apache.kafka.clients.consumer.Consumer;
@@ -102,6 +104,7 @@ private static <T> T example(Class<T> klass) {
   @Mock AssignerFactory assignerFactory;
   @Mock CursorClient cursorClient;
   @Mock AdminClient adminClient;
+  @Mock TopicStatsClient topicStatsClient;
 
   @Mock Assigner assigner;
   @Mock SingleSubscriptionConsumer underlying;
@@ -118,7 +121,8 @@ public void setUp() {
             new SharedBehavior(adminClient),
             consumerFactory,
             assignerFactory,
-            cursorClient);
+            cursorClient,
+            topicStatsClient);
     when(consumerFactory.newConsumer()).thenReturn(underlying);
   }
 
@@ -140,10 +144,6 @@ public void unsupportedOperations() {
     assertThrows(
         UnsupportedVersionException.class,
         () -> consumer.offsetsForTimes(ImmutableMap.of(), Duration.ZERO));
-    assertThrows(UnsupportedVersionException.class, () -> consumer.endOffsets(ImmutableList.of()));
-    assertThrows(
-        UnsupportedVersionException.class,
-        () -> consumer.endOffsets(ImmutableList.of(), Duration.ZERO));
   }
 
   @Test
@@ -468,10 +468,24 @@ public void partitionsFor() {
     assertThat(info.size()).isEqualTo(2L);
   }
 
+  @Test
+  public void endOffsets() {
+    TopicPartition partition2 = new TopicPartition(example(TopicPath.class).toString(), 2);
+    TopicPartition partition4 = new TopicPartition(example(TopicPath.class).toString(), 4);
+    when(topicStatsClient.computeHeadCursor(example(TopicPath.class), Partition.of(2)))
+        .thenReturn(ApiFutures.immediateFuture(Cursor.newBuilder().setOffset(22).build()));
+    when(topicStatsClient.computeHeadCursor(example(TopicPath.class), Partition.of(4)))
+        .thenReturn(ApiFutures.immediateFuture(Cursor.newBuilder().setOffset(44).build()));
+    Map<TopicPartition, Long> output =
+        consumer.endOffsets(ImmutableList.of(partition2, partition4));
+    assertThat(output).isEqualTo(ImmutableMap.of(partition2, 22L, partition4, 44L));
+  }
+
   @Test
   public void close() {
     consumer.close();
     verify(adminClient).close();
     verify(cursorClient).close();
+    verify(topicStatsClient).close();
   }
 }