feat: Add support for Managed Autoscaler (googleapis#2624)

- Add the ability to create / update an instance with autoscaling config - Note that now user can only specify one of the node_count, processing_units or autoscaling_config as the compute capacity for an instance. - Add the ability to see autoscaling config with an instance
claire921 · Oct 19, 2023 · e5e6923 · e5e6923
1 parent 712c65f
commit e5e6923
Show file tree

Hide file tree

Showing 7 changed files with 472 additions and 34 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Instance.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Instance.java
@@ -23,6 +23,7 @@
 import com.google.cloud.spanner.Options.ListOption;
 import com.google.longrunning.Operation;
 import com.google.spanner.admin.database.v1.CreateDatabaseMetadata;
+import com.google.spanner.admin.instance.v1.AutoscalingConfig;
 import com.google.spanner.admin.instance.v1.UpdateInstanceMetadata;
 import java.util.Map;
 
@@ -86,6 +87,12 @@ public Builder setProcessingUnits(int processingUnits) {
       return this;
     }
 
+    @Override
+    public Builder setAutoscalingConfig(AutoscalingConfig autoscalingConfig) {
+      infoBuilder.setAutoscalingConfig(autoscalingConfig);
+      return this;
+    }
+
     @Override
     public Builder setState(State state) {
       infoBuilder.setState(state);
@@ -220,6 +227,7 @@ static Instance fromProto(
             .setNodeCount(proto.getNodeCount())
             .setCreateTime(Timestamp.fromProto(proto.getCreateTime()))
             .setUpdateTime(Timestamp.fromProto(proto.getUpdateTime()))
+            .setAutoscalingConfig(proto.getAutoscalingConfig())
             .setProcessingUnits(proto.getProcessingUnits());
     State state;
     switch (proto.getState()) {

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceAdminClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceAdminClientImpl.java
@@ -193,9 +193,6 @@ public Operation fromProto(Operation proto) {
   @Override
   public OperationFuture<Instance, CreateInstanceMetadata> createInstance(InstanceInfo instance)
       throws SpannerException {
-    Preconditions.checkArgument(
-        instance.getNodeCount() == 0 || instance.getProcessingUnits() == 0,
-        "Only one of nodeCount and processingUnits can be set when creating a new instance");
     String projectName = PROJECT_NAME_TEMPLATE.instantiate("project", projectId);
     OperationFuture<com.google.spanner.admin.instance.v1.Instance, CreateInstanceMetadata>
         rawOperationFuture =

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceInfo.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceInfo.java
@@ -23,6 +23,7 @@
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.FieldMask;
+import com.google.spanner.admin.instance.v1.AutoscalingConfig;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -35,13 +36,16 @@ public enum InstanceField implements FieldSelector {
     DISPLAY_NAME("display_name"),
     NODE_COUNT("node_count"),
     PROCESSING_UNITS("processing_units"),
+    AUTOSCALING_CONFIG("autoscaling_config"),
     LABELS("labels");
 
     static InstanceField[] defaultFieldsToUpdate(InstanceInfo info) {
-      if (info.getNodeCount() > 0) {
-        return new InstanceField[] {DISPLAY_NAME, NODE_COUNT, LABELS};
+      if (info.getAutoscalingConfig() != null) {
+        return new InstanceField[] {DISPLAY_NAME, AUTOSCALING_CONFIG, LABELS};
+      } else if (info.getNodeCount() > 0) {
+        return new InstanceField[] {DISPLAY_NAME, AUTOSCALING_CONFIG, NODE_COUNT, LABELS};
       } else {
-        return new InstanceField[] {DISPLAY_NAME, PROCESSING_UNITS, LABELS};
+        return new InstanceField[] {DISPLAY_NAME, AUTOSCALING_CONFIG, PROCESSING_UNITS, LABELS};
       }
     }
 
@@ -87,21 +91,31 @@ Builder setCreateTime(Timestamp createTime) {
     }
 
     /**
-     * Sets the number of nodes for the instance. Exactly one of processing units or node count must
-     * be set when creating a new instance.
+     * Sets the number of nodes for the instance. Exactly one of processing units, node count or
+     * autoscaling config must be set when creating a new instance.
      */
     public abstract Builder setNodeCount(int nodeCount);
 
     /**
-     * Sets the number of processing units for the instance. Exactly one of processing units or node
-     * count must be set when creating a new instance. Processing units must be between 1 and 999
-     * (inclusive) when creating a new instance with node count = 0. Processing units from 1000 and
-     * up must always be a multiple of 1000 (that is equal to an integer number of nodes).
+     * Sets the number of processing units for the instance. Exactly one of processing units, node
+     * count, or autoscaling config must be set when creating a new instance. Processing units must
+     * be between 1 and 999 (inclusive) when creating a new instance with node count = 0. Processing
+     * units from 1000 and up must always be a multiple of 1000 (that is equal to an integer number
+     * of nodes).
      */
     public Builder setProcessingUnits(int processingUnits) {
       throw new UnsupportedOperationException("Unimplemented");
     }
 
+    /**
+     * Sets the autoscaling config for the instance, which will enable the autoscaling for this
+     * instance. Exactly one of processing units, node count, or autoscaling config must be set when
+     * creating a new instance.
+     */
+    public Builder setAutoscalingConfig(AutoscalingConfig autoscalingConfig) {
+      throw new UnsupportedOperationException("Unimplemented");
+    }
+
     public abstract Builder setState(State state);
 
     public abstract Builder addLabel(String key, String value);
@@ -117,6 +131,7 @@ static class BuilderImpl extends Builder {
     private String displayName;
     private int nodeCount;
     private int processingUnits;
+    private AutoscalingConfig autoscalingConfig;
     private State state;
     private Map<String, String> labels;
     private Timestamp updateTime;
@@ -133,6 +148,7 @@ static class BuilderImpl extends Builder {
       this.displayName = instance.displayName;
       this.nodeCount = instance.nodeCount;
       this.processingUnits = instance.processingUnits;
+      this.autoscalingConfig = instance.autoscalingConfig;
       this.state = instance.state;
       this.labels = new HashMap<>(instance.labels);
       this.updateTime = instance.updateTime;
@@ -175,6 +191,12 @@ public BuilderImpl setProcessingUnits(int processingUnits) {
       return this;
     }
 
+    @Override
+    public BuilderImpl setAutoscalingConfig(AutoscalingConfig autoscalingConfig) {
+      this.autoscalingConfig = autoscalingConfig;
+      return this;
+    }
+
     @Override
     public BuilderImpl setState(State state) {
       this.state = state;
@@ -204,6 +226,7 @@ public InstanceInfo build() {
   private final String displayName;
   private final int nodeCount;
   private final int processingUnits;
+  private final AutoscalingConfig autoscalingConfig;
   private final State state;
   private final ImmutableMap<String, String> labels;
   private final Timestamp updateTime;
@@ -215,6 +238,7 @@ public InstanceInfo build() {
     this.displayName = builder.displayName;
     this.nodeCount = builder.nodeCount;
     this.processingUnits = builder.processingUnits;
+    this.autoscalingConfig = builder.autoscalingConfig;
     this.state = builder.state;
     this.labels = ImmutableMap.copyOf(builder.labels);
     this.updateTime = builder.updateTime;
@@ -254,6 +278,11 @@ public int getProcessingUnits() {
     return processingUnits;
   }
 
+  /** Returns the autoscaling config of the instance. */
+  public AutoscalingConfig getAutoscalingConfig() {
+    return autoscalingConfig;
+  }
+
   /** Returns the current state of the instance. */
   public State getState() {
     return state;
@@ -276,6 +305,7 @@ public String toString() {
         .add("displayName", displayName)
         .add("nodeCount", nodeCount)
         .add("processingUnits", processingUnits)
+        .add("autoscaling_config", autoscalingConfig)
         .add("state", state)
         .add("labels", labels)
         .add("createTime", createTime)
@@ -297,6 +327,7 @@ public boolean equals(Object o) {
         && Objects.equals(displayName, that.displayName)
         && nodeCount == that.nodeCount
         && processingUnits == that.processingUnits
+        && Objects.equals(autoscalingConfig, that.autoscalingConfig)
         && state == that.state
         && Objects.equals(labels, that.labels)
         && Objects.equals(updateTime, that.updateTime)
@@ -311,6 +342,7 @@ public int hashCode() {
         displayName,
         nodeCount,
         processingUnits,
+        autoscalingConfig,
         state,
         labels,
         updateTime,
@@ -330,6 +362,9 @@ com.google.spanner.admin.instance.v1.Instance toProto() {
     if (getInstanceConfigId() != null) {
       builder.setConfig(getInstanceConfigId().getName());
     }
+    if (getAutoscalingConfig() != null) {
+      builder.setAutoscalingConfig(getAutoscalingConfig());
+    }
     return builder.build();
   }