feat: cluster autoscaling (#1092)

* feat: cluster autoscaling

* lint

* test

* test

* test

* test

* fix tests

* fix test

* use new instance for test

* lint

* add nonnull

* add nonnull

* add nonnull

Author: kolea2

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableInstanceAdminClient.java

@@ -24,13 +24,15 @@
 import com.google.bigtable.admin.v2.DeleteAppProfileRequest;
 import com.google.bigtable.admin.v2.GetAppProfileRequest;
 import com.google.bigtable.admin.v2.ListAppProfilesRequest;
+import com.google.bigtable.admin.v2.PartialUpdateClusterRequest;
 import com.google.cloud.Policy;
 import com.google.cloud.Policy.DefaultMarshaller;
 import com.google.cloud.bigtable.admin.v2.BaseBigtableInstanceAdminClient.ListAppProfilesPage;
 import com.google.cloud.bigtable.admin.v2.BaseBigtableInstanceAdminClient.ListAppProfilesPagedResponse;
 import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
 import com.google.cloud.bigtable.admin.v2.models.AppProfile;
 import com.google.cloud.bigtable.admin.v2.models.Cluster;
+import com.google.cloud.bigtable.admin.v2.models.ClusterAutoscalingConfig;
 import com.google.cloud.bigtable.admin.v2.models.CreateAppProfileRequest;
 import com.google.cloud.bigtable.admin.v2.models.CreateClusterRequest;
 import com.google.cloud.bigtable.admin.v2.models.CreateInstanceRequest;
@@ -49,6 +51,7 @@
 import com.google.iam.v1.TestIamPermissionsRequest;
 import com.google.iam.v1.TestIamPermissionsResponse;
 import com.google.protobuf.Empty;
+import com.google.protobuf.util.FieldMaskUtil;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
@@ -660,8 +663,10 @@ public List<Cluster> apply(com.google.bigtable.admin.v2.ListClustersResponse pro
   }
 
   /**
-   * Modifies the cluster's node count. Please note that only clusters that belong to a production
-   * instance can be resized.
+   * Modifies the cluster's node count for manual scaling. If autoscaling is already enabled, manual
+   * scaling will be silently ignored. If you wish to disable autoscaling and enable manual scaling,
+   * please use {@link BigtableInstanceAdminClient#disableClusterAutoscaling(String, String, int)}
+   * instead. Please note that only clusters that belong to a production instance can be resized.
    *
    * <p>Sample code:
    *
@@ -676,8 +681,11 @@ public Cluster resizeCluster(String instanceId, String clusterId, int numServeNo
   }
 
   /**
-   * Asynchronously modifies the cluster's node count. Please note that only clusters that belong to
-   * a production instance can be resized.
+   * Asynchronously modifies the cluster's node count for manual scaling. If autoscaling is already
+   * enabled, manual scaling will be silently ignored. If you wish to disable autoscaling and enable
+   * manual scaling, please use {@link BigtableInstanceAdminClient#disableClusterAutoscaling(String,
+   * String, int)} instead. Please note that only clusters that belong to a production instance can
+   * be resized.
    *
    * <pre>{@code
    * ApiFuture<Cluster> clusterFuture = client.resizeCluster("my-instance", "my-cluster", 30);
@@ -707,6 +715,109 @@ public Cluster apply(com.google.bigtable.admin.v2.Cluster proto) {
         MoreExecutors.directExecutor());
   }
 
+  /**
+   * Modifies the cluster's autoscaling config. This will enable autoscaling and disable manual
+   * scaling if the cluster is manually scaled. Please note that only clusters that belong to a
+   * production instance can enable autoscaling.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * ClusterAutoscalingConfig clusterAutoscalingConfig =
+   *      ClusterAutoscalingConfig.of("my-instance", "my-cluster")
+   *          .setMinNodes(1)
+   *          .setMaxNodes(4)
+   *          .setCpuUtilizationTargetPercent(40);
+   * Cluster cluster = client.updateClusterAutoscalingConfig(clusterAutoscalingConfig);
+   * }</pre>
+   */
+  public Cluster updateClusterAutoscalingConfig(
+      @Nonnull ClusterAutoscalingConfig clusterAutoscalingConfig) {
+    return ApiExceptions.callAndTranslateApiException(
+        updateClusterAutoscalingConfigAsync(clusterAutoscalingConfig));
+  }
+
+  /**
+   * Asynchronously modifies the cluster's autoscaling config. This will enable autoscaling and
+   * disable manual scaling if the cluster is manually scaled. Please note that only clusters that
+   * belong to a production instance can enable autoscaling.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * ClusterAutoscalingConfig clusterAutoscalingConfig =
+   *      ClusterAutoscalingConfig.of(targetInstanceId, targetClusterId)
+   *          .setMinNodes(1)
+   *          .setMaxNodes(4)
+   *          .setCpuUtilizationTargetPercent(40);
+   *
+   *  ApiFuture<Cluster> clusterApiFuture = client.updateClusterAutoscalingConfigAsync(clusterAutoscalingConfig);
+   *  Cluster cluster = clusterApiFuture.get();
+   * }</pre>
+   */
+  public ApiFuture<Cluster> updateClusterAutoscalingConfigAsync(
+      @Nonnull ClusterAutoscalingConfig clusterAutoscalingConfig) {
+    PartialUpdateClusterRequest proto = clusterAutoscalingConfig.toProto(projectId);
+
+    return ApiFutures.transform(
+        stub.partialUpdateClusterOperationCallable().futureCall(proto),
+        Cluster::fromProto,
+        MoreExecutors.directExecutor());
+  }
+
+  /**
+   * Disables autoscaling and enables manual scaling by setting a static node count for the cluster.
+   * Please note that only clusters that belong to a production instance can be resized.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * Cluster cluster = client.disableClusterAutoscaling("my-instance", "my-cluster", 3);
+   * }</pre>
+   */
+  public Cluster disableClusterAutoscaling(String instanceId, String clusterId, int staticSize) {
+    return ApiExceptions.callAndTranslateApiException(
+        disableClusterAutoscalingAsync(instanceId, clusterId, staticSize));
+  }
+
+  /**
+   * Asynchronously disables autoscaling and enables manual scaling by setting a static node count
+   * for the cluster. Please note that only clusters that belong to a production instance can be
+   * resized.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * ApiFuture<Cluster> clusterApiFuture = client.disableClusterAutoscalingAsync("my-instance", "my-cluster", 3);
+   * Cluster cluster = clusterApiFuture.get();
+   * }</pre>
+   */
+  public ApiFuture<Cluster> disableClusterAutoscalingAsync(
+      String instanceId, String clusterId, int staticSize) {
+    String name = NameUtil.formatClusterName(projectId, instanceId, clusterId);
+
+    com.google.bigtable.admin.v2.Cluster request =
+        com.google.bigtable.admin.v2.Cluster.newBuilder()
+            .setName(name)
+            .setServeNodes(staticSize)
+            .setClusterConfig(
+                com.google.bigtable.admin.v2.Cluster.ClusterConfig.getDefaultInstance())
+            .build();
+
+    PartialUpdateClusterRequest partialUpdateClusterRequest =
+        PartialUpdateClusterRequest.newBuilder()
+            .setUpdateMask(
+                FieldMaskUtil.fromStringList(
+                    com.google.bigtable.admin.v2.Cluster.class,
+                    Lists.newArrayList("cluster_config.cluster_autoscaling_config", "serve_nodes")))
+            .setCluster(request)
+            .build();
+    return ApiFutures.transform(
+        stub.partialUpdateClusterOperationCallable().futureCall(partialUpdateClusterRequest),
+        Cluster::fromProto,
+        MoreExecutors.directExecutor());
+  }
+
   /**
    * Deletes the specified cluster. Please note that an instance must have at least 1 cluster. To
    * remove the last cluster, please use {@link BigtableInstanceAdminClient#deleteInstance(String)}.

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/Cluster.java

@@ -148,6 +148,36 @@ public int getServeNodes() {
     return stateProto.getServeNodes();
   }
 
+  /** Get the minimum number of nodes to scale down to. */
+  public int getAutoscalingMinServeNodes() {
+    return stateProto
+        .getClusterConfig()
+        .getClusterAutoscalingConfig()
+        .getAutoscalingLimits()
+        .getMinServeNodes();
+  }
+
+  /** Get the maximum number of nodes to scale up to. */
+  public int getAutoscalingMaxServeNodes() {
+    return stateProto
+        .getClusterConfig()
+        .getClusterAutoscalingConfig()
+        .getAutoscalingLimits()
+        .getMaxServeNodes();
+  }
+
+  /**
+   * Get the cpu utilization that the Autoscaler should be trying to achieve. This number is on a
+   * scale from 0 (no utilization) to 100 (total utilization).
+   */
+  public int getAutoscalingCpuPercentageTarget() {
+    return stateProto
+        .getClusterConfig()
+        .getClusterAutoscalingConfig()
+        .getAutoscalingTargets()
+        .getCpuUtilizationPercent();
+  }
+
   /**
    * The type of storage used by this cluster to serve its parent instance's tables, unless
    * explicitly overridden.