// Copyright 2019 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;

import static com.google.cloud.pubsublite.internal.Preconditions.checkState;

import com.google.api.core.ApiFuture;
import com.google.cloud.pubsublite.internal.ExtractStatus;
import com.google.cloud.pubsublite.proto.Subscription;
import io.grpc.StatusException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

public final class PartitionLookupUtils {
  private PartitionLookupUtils() {}

  public static int numPartitions(TopicPath topic) throws StatusException {
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    AdminClient client =
        AdminClientBuilder.builder()
            .setRegion(TopicPaths.getZone(topic).region())
            .setExecutor(executor)
            .build();
    return numPartitions(topic, client);
  }

  public static int numPartitions(TopicPath topic, AdminClient client) throws StatusException {
    ApiFuture<Long> partitionCountFuture = client.getTopicPartitionCount(topic);
    try {
      long numPartitions = partitionCountFuture.get();
      checkState(
          numPartitions > 0, "Config has 0 or less partitions configured. This config is invalid.");
      checkState(
          numPartitions < Integer.MAX_VALUE,
          "Config has more than Integer.MAX_VALUE partitions configured. This"
              + " config cannot be used with this client library.");
      return (int) numPartitions;
    } catch (ExecutionException e) {
      throw ExtractStatus.toCanonical(e.getCause());
    } catch (InterruptedException t) {
      throw ExtractStatus.toCanonical(t);
    } catch (Throwable t) {
      throw ExtractStatus.toCanonical(t);
    }
  }

  public static int numPartitions(SubscriptionPath subscription) throws StatusException {
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    AdminClient client =
        AdminClientBuilder.builder()
            .setRegion(SubscriptionPaths.getZone(subscription).region())
            .setExecutor(executor)
            .build();
    return numPartitions(subscription, client);
  }

  public static int numPartitions(SubscriptionPath subscription, AdminClient client)
      throws StatusException {
    ApiFuture<Subscription> subscriptionFuture = client.getSubscription(subscription);
    try {
      return numPartitions(TopicPath.of(subscriptionFuture.get().getTopic()), client);
    } catch (ExecutionException e) {
      throw ExtractStatus.toCanonical(e.getCause());
    } catch (InterruptedException t) {
      throw ExtractStatus.toCanonical(t);
    } catch (Throwable t) {
      throw ExtractStatus.toCanonical(t);
    }
  }
}

@@ -1257,32 +1257,44 @@
 uledExecutor();%0A
+    try %7B%0A  
     AdminClient 
@@ -1294,32 +1294,34 @@
 Client client =%0A
+  
         AdminCli
@@ -1345,32 +1345,34 @@
 r()%0A            
+ 
+ 
 .setRegion(Topic
@@ -1406,32 +1406,34 @@
 ())%0A            
+  
 .setExecutor(exe
@@ -1443,32 +1443,34 @@
 or)%0A            
+ 
+ 
 .build();%0A    re
@@ -1455,32 +1455,34 @@
       .build();%0A
+  
     return numPa
@@ -1498,32 +1498,84 @@
 topic, client);%0A
+    %7D finally %7B%0A      executor.shutdownNow();%0A    %7D%0A
   %7D%0A%0A  public st

Modify PartitionLookupUtils to tear down executor service.