feat: Add the PartitionCountWatchingPublisher (#387)

* feat: Add the PartitionCountWatchingPublisher

* changes to address comments

* more changes

google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/DefaultRoutingPolicy.java

@@ -26,25 +26,27 @@
 import com.google.protobuf.ByteString;
 import java.math.BigInteger;
 import java.util.Random;
+import java.util.concurrent.ThreadLocalRandom;
 
 public class DefaultRoutingPolicy implements RoutingPolicy {
-  private final int numPartitions;
+  private final long numPartitions;
   private final CloseableMonitor monitor = new CloseableMonitor();
 
   @GuardedBy("monitor.monitor")
-  private int nextWithoutKeyPartition;
+  private long nextWithoutKeyPartition;
 
-  public DefaultRoutingPolicy(int numPartitions) throws ApiException {
+  public DefaultRoutingPolicy(long numPartitions) throws ApiException {
     checkArgument(numPartitions > 0, "Must have a positive number of partitions.");
     this.numPartitions = numPartitions;
-    this.nextWithoutKeyPartition = new Random().nextInt(this.numPartitions);
+    this.nextWithoutKeyPartition = ThreadLocalRandom.current().nextLong(numPartitions);
+    this.nextWithoutKeyPartition = new Random().longs(1, 0, numPartitions).findFirst().getAsLong();
   }
 
   @Override
   public Partition routeWithoutKey() throws ApiException {
     try (CloseableMonitor.Hold h = monitor.enter()) {
       Partition toReturn = Partition.of(nextWithoutKeyPartition);
-      int next = nextWithoutKeyPartition + 1;
+      long next = nextWithoutKeyPartition + 1;
       next = next % numPartitions;
       nextWithoutKeyPartition = next;
       return toReturn;

google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/RoutingPolicy.java

@@ -21,6 +21,9 @@
 
 // Route the user message key to a given partition.
 public interface RoutingPolicy {
+  interface Factory {
+    RoutingPolicy newPolicy(long numPartitions);
+  }
   // Route a message without a key to a partition.
   Partition routeWithoutKey() throws CheckedApiException;
   // Route a message with a key to a partition.

google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/PartitionCountWatchingPublisher.java

@@ -0,0 +1,203 @@
+/*
+ * Copyright 2020 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.internal.wire;
+
+import static com.google.cloud.pubsublite.internal.CheckedApiPreconditions.checkState;
+import static com.google.cloud.pubsublite.internal.ExtractStatus.toCanonical;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.ApiService;
+import com.google.cloud.pubsublite.Message;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.PublishMetadata;
+import com.google.cloud.pubsublite.internal.*;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.flogger.GoogleLogger;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import java.io.IOException;
+import java.util.Optional;
+import java.util.stream.LongStream;
+
+public class PartitionCountWatchingPublisher extends ProxyService
+    implements Publisher<PublishMetadata> {
+  private static final GoogleLogger log = GoogleLogger.forEnclosingClass();
+  private final PartitionPublisherFactory publisherFactory;
+  private final RoutingPolicy.Factory policyFactory;
+
+  private static class PartitionsWithRouting {
+    public final ImmutableMap<Partition, Publisher<PublishMetadata>> publishers;
+    private final RoutingPolicy routingPolicy;
+
+    private PartitionsWithRouting(
+        ImmutableMap<Partition, Publisher<PublishMetadata>> publishers,
+        RoutingPolicy routingPolicy) {
+      this.publishers = publishers;
+      this.routingPolicy = routingPolicy;
+    }
+
+    public ApiFuture<PublishMetadata> publish(Message message) throws CheckedApiException {
+      try {
+        Partition routedPartition =
+            message.key().isEmpty()
+                ? routingPolicy.routeWithoutKey()
+                : routingPolicy.route(message.key());
+        checkState(
+            publishers.containsKey(routedPartition),
+            String.format(
+                "Routed to partition %s for which there is no publisher available.",
+                routedPartition));
+        return publishers.get(routedPartition).publish(message);
+      } catch (Throwable t) {
+        throw toCanonical(t);
+      }
+    }
+
+    public void flush() throws IOException {
+      for (Publisher<PublishMetadata> publisher : publishers.values()) {
+        publisher.flush();
+      }
+    }
+
+    public void stop() {
+      publishers.values().forEach(ApiService::stopAsync);
+      publishers.values().forEach(ApiService::awaitTerminated);
+    }
+  }
+
+  private final CloseableMonitor monitor = new CloseableMonitor();
+
+  @GuardedBy("monitor.monitor")
+  private boolean shutdown = false;
+
+  @GuardedBy("monitor.monitor")
+  private Optional<PartitionsWithRouting> partitionsWithRouting = Optional.empty();
+
+  public PartitionCountWatchingPublisher(PartitionCountWatchingPublisherSettings settings) {
+    this.publisherFactory = settings.publisherFactory();
+    this.policyFactory = settings.routingPolicyFactory();
+    PartitionCountWatcher configWatcher =
+        settings.configWatcherFactory().newWatcher(this::handleConfig);
+    addServices(configWatcher);
+  }
+
+  @Override
+  public ApiFuture<PublishMetadata> publish(Message message) {
+    Optional<PartitionsWithRouting> partitions;
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      partitions = partitionsWithRouting;
+    }
+    if (!partitions.isPresent()) {
+      throw new IllegalStateException("Publish called before start or after shutdown");
+    }
+    try {
+      return partitions.get().publish(message);
+    } catch (CheckedApiException e) {
+      onPermanentError(e);
+      return ApiFutures.immediateFailedFuture(e);
+    }
+  }
+
+  @Override
+  public void flush() throws IOException {
+    Optional<PartitionsWithRouting> partitions;
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      partitions = partitionsWithRouting;
+    }
+    if (!partitions.isPresent()) {
+      throw new IllegalStateException("Publish called before start or after shutdown");
+    }
+    partitions.get().flush();
+  }
+
+  private ImmutableMap<Partition, Publisher<PublishMetadata>> getNewPartitionPublishers(
+      LongStream newPartitions) {
+    ImmutableMap.Builder<Partition, Publisher<PublishMetadata>> mapBuilder = ImmutableMap.builder();
+    newPartitions.forEach(
+        i -> {
+          Publisher<PublishMetadata> p = publisherFactory.newPublisher(Partition.of(i));
+          p.addListener(
+              new Listener() {
+                @Override
+                public void failed(State from, Throwable failure) {
+                  onPermanentError(toCanonical(failure));
+                }
+              },
+              MoreExecutors.directExecutor());
+          mapBuilder.put(Partition.of(i), p);
+          p.startAsync();
+        });
+    ImmutableMap<Partition, Publisher<PublishMetadata>> partitions = mapBuilder.build();
+    partitions.values().forEach(ApiService::awaitRunning);
+    return partitions;
+  }
+
+  private void handleConfig(long partitionCount) {
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      if (shutdown) {
+        return;
+      }
+      Optional<PartitionsWithRouting> current = partitionsWithRouting;
+      long currentSize = current.map(withRouting -> withRouting.publishers.size()).orElse(0);
+      if (partitionCount == currentSize) {
+        return;
+      }
+      if (partitionCount < currentSize) {
+        log.atWarning().log(
+            "Received an unexpected decrease in partition count. Previous partition count {}, new count {}",
+            currentSize,
+            partitionCount);
+        return;
+      }
+      ImmutableMap.Builder<Partition, Publisher<PublishMetadata>> mapBuilder =
+          ImmutableMap.builder();
+      current.ifPresent(p -> p.publishers.forEach(mapBuilder::put));
+      getNewPartitionPublishers(LongStream.range(currentSize, partitionCount))
+          .forEach(mapBuilder::put);
+      ImmutableMap<Partition, Publisher<PublishMetadata>> newMap = mapBuilder.build();
+
+      partitionsWithRouting =
+          Optional.of(
+              new PartitionsWithRouting(
+                  mapBuilder.build(), policyFactory.newPolicy(partitionCount)));
+    }
+  }
+
+  @Override
+  protected void start() {}
+
+  @Override
+  protected void stop() {
+    Optional<PartitionsWithRouting> current;
+    try (CloseableMonitor.Hold h = monitor.enter()) {
+      shutdown = true;
+      current = partitionsWithRouting;
+      partitionsWithRouting = Optional.empty();
+    }
+    current.ifPresent(PartitionsWithRouting::stop);
+  }
+
+  @Override
+  protected void handlePermanentError(CheckedApiException error) {
+    try {
+      stop();
+    } catch (Exception e) {
+      log.atWarning().withCause(e).log("Encountered exception while trying to handle failure");
+    }
+  }
+}

google-cloud-pubsublite/src/main/java/com/google/cloud/pubsublite/internal/wire/PartitionCountWatchingPublisherSettings.java

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2020 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.internal.wire;
+
+import com.google.auto.value.AutoValue;
+import com.google.cloud.pubsublite.*;
+import com.google.cloud.pubsublite.internal.DefaultRoutingPolicy;
+import com.google.cloud.pubsublite.internal.RoutingPolicy;
+import java.time.Duration;
+import java.util.Optional;
+
+@AutoValue
+public abstract class PartitionCountWatchingPublisherSettings {
+  // Required parameters.
+  abstract TopicPath topic();
+
+  abstract PartitionPublisherFactory publisherFactory();
+
+  // Optional parameters
+  abstract PartitionCountWatcher.Factory configWatcherFactory();
+
+  abstract RoutingPolicy.Factory routingPolicyFactory();
+
+  abstract Duration configPollPeriod();
+
+  public static Builder newBuilder() {
+    return new AutoValue_PartitionCountWatchingPublisherSettings.Builder()
+        .setConfigPollPeriod(Duration.ofMinutes(10));
+  }
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+    // Required parameters.
+    public abstract Builder setTopic(TopicPath path);
+
+    public abstract Builder setPublisherFactory(PartitionPublisherFactory factory);
+
+    // Optional parameters.
+    public abstract Builder setConfigWatcherFactory(PartitionCountWatcher.Factory factory);
+
+    public abstract Builder setRoutingPolicyFactory(RoutingPolicy.Factory factory);
+
+    public abstract Builder setConfigPollPeriod(Duration period);
+
+    abstract Optional<PartitionCountWatcher.Factory> configWatcherFactory();
+
+    abstract Optional<RoutingPolicy.Factory> routingPolicyFactory();
+
+    abstract Duration configPollPeriod();
+
+    abstract TopicPath topic();
+
+    abstract PartitionCountWatchingPublisherSettings autoBuild();
+
+    public PartitionCountWatchingPublisherSettings build() {
+      if (!configWatcherFactory().isPresent()) {
+        setConfigWatcherFactory(
+            new PartitionCountWatcherImpl.Factory(
+                topic(),
+                AdminClient.create(
+                    AdminClientSettings.newBuilder()
+                        .setRegion(topic().location().region())
+                        .build()),
+                configPollPeriod()));
+      }
+      if (!routingPolicyFactory().isPresent()) {
+        setRoutingPolicyFactory(DefaultRoutingPolicy::new);
+      }
+      return autoBuild();
+    }
+  }
+}