feat: add support for instance processing units (#665)

* feat: add support for instance processing units

* docs: add sample for instance processing units

* fix: remove commented code + use junit asserts

* fix: remove Spanner snapshot reference

* test: add tests for InstanceInfo

* cleanup: fix formatting + tests

* chore: fix formatting

* test: fixes instanceinfo test

Due to master merge the compilation was failing. Fixes the builder
construction here.

* refactor: adds default impl for processing units

Adds default implementation for the InstanceInfo set processing units in
order to avoid a breaking change.

* samples: remove LCI samples for now

Removes the LCI samples for now, because they won't compile. We will
re-add them once the main implementation is released.

* fix: addresses PR comments

Co-authored-by: Thiago Tasca Nunes <[email protected]>

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/spanner/Instance.java

@@ -67,6 +67,12 @@ public Builder setNodeCount(int nodeCount) {
       return this;
     }
 
+    @Override
+    public Builder setProcessingUnits(int processingUnits) {
+      infoBuilder.setProcessingUnits(processingUnits);
+      return this;
+    }
+
     @Override
     public Builder setState(State state) {
       infoBuilder.setState(state);
@@ -198,7 +204,8 @@ static Instance fromProto(
         new Builder(instanceClient, dbClient, id)
             .setInstanceConfigId(InstanceConfigId.of(proto.getConfig()))
             .setDisplayName(proto.getDisplayName())
-            .setNodeCount(proto.getNodeCount());
+            .setNodeCount(proto.getNodeCount())
+            .setProcessingUnits(proto.getProcessingUnits());
     State state;
     switch (proto.getState()) {
       case STATE_UNSPECIFIED:

google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceAdminClientImpl.java

@@ -96,6 +96,9 @@ public InstanceConfig fromProto(
   @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 =
@@ -158,7 +161,8 @@ public OperationFuture<Instance, UpdateInstanceMetadata> updateInstance(
       InstanceInfo instance, InstanceInfo.InstanceField... fieldsToUpdate) {
     FieldMask fieldMask =
         fieldsToUpdate.length == 0
-            ? InstanceInfo.InstanceField.toFieldMask(InstanceInfo.InstanceField.values())
+            ? InstanceInfo.InstanceField.toFieldMask(
+                InstanceInfo.InstanceField.defaultFieldsToUpdate(instance))
             : InstanceInfo.InstanceField.toFieldMask(fieldsToUpdate);
 
     OperationFuture<com.google.spanner.admin.instance.v1.Instance, UpdateInstanceMetadata>

google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceInfo.java

@@ -33,8 +33,17 @@ public class InstanceInfo {
   public enum InstanceField implements FieldSelector {
     DISPLAY_NAME("display_name"),
     NODE_COUNT("node_count"),
+    PROCESSING_UNITS("processing_units"),
     LABELS("labels");
 
+    static InstanceField[] defaultFieldsToUpdate(InstanceInfo info) {
+      if (info.getNodeCount() > 0) {
+        return new InstanceField[] {DISPLAY_NAME, NODE_COUNT, LABELS};
+      } else {
+        return new InstanceField[] {DISPLAY_NAME, PROCESSING_UNITS, LABELS};
+      }
+    }
+
     private final String selector;
 
     InstanceField(String selector) {
@@ -68,8 +77,22 @@ public abstract static class Builder {
 
     public abstract Builder setDisplayName(String displayName);
 
+    /**
+     * Sets the number of nodes for the instance. Exactly one of processing units or node count 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).
+     */
+    public Builder setProcessingUnits(int processingUnits) {
+      throw new UnsupportedOperationException("Unimplemented");
+    }
+
     public abstract Builder setState(State state);
 
     public abstract Builder addLabel(String key, String value);
@@ -84,6 +107,7 @@ static class BuilderImpl extends Builder {
     private InstanceConfigId configId;
     private String displayName;
     private int nodeCount;
+    private int processingUnits;
     private State state;
     private Map<String, String> labels;
 
@@ -97,6 +121,7 @@ static class BuilderImpl extends Builder {
       this.configId = instance.configId;
       this.displayName = instance.displayName;
       this.nodeCount = instance.nodeCount;
+      this.processingUnits = instance.processingUnits;
       this.state = instance.state;
       this.labels = new HashMap<>(instance.labels);
     }
@@ -119,6 +144,12 @@ public BuilderImpl setNodeCount(int nodeCount) {
       return this;
     }
 
+    @Override
+    public BuilderImpl setProcessingUnits(int processingUnits) {
+      this.processingUnits = processingUnits;
+      return this;
+    }
+
     @Override
     public BuilderImpl setState(State state) {
       this.state = state;
@@ -147,6 +178,7 @@ public InstanceInfo build() {
   private final InstanceConfigId configId;
   private final String displayName;
   private final int nodeCount;
+  private final int processingUnits;
   private final State state;
   private final ImmutableMap<String, String> labels;
 
@@ -155,6 +187,7 @@ public InstanceInfo build() {
     this.configId = builder.configId;
     this.displayName = builder.displayName;
     this.nodeCount = builder.nodeCount;
+    this.processingUnits = builder.processingUnits;
     this.state = builder.state;
     this.labels = ImmutableMap.copyOf(builder.labels);
   }
@@ -179,6 +212,11 @@ public int getNodeCount() {
     return nodeCount;
   }
 
+  /** Returns the number of processing units of the instance. */
+  public int getProcessingUnits() {
+    return processingUnits;
+  }
+
   /** Returns the current state of the instance. */
   public State getState() {
     return state;
@@ -200,6 +238,7 @@ public String toString() {
         .add("configName", configId == null ? null : configId.getName())
         .add("displayName", displayName)
         .add("nodeCount", nodeCount)
+        .add("processingUnits", processingUnits)
         .add("state", state)
         .add("labels", labels)
         .toString();
@@ -218,20 +257,22 @@ public boolean equals(Object o) {
         && Objects.equals(configId, that.configId)
         && Objects.equals(displayName, that.displayName)
         && nodeCount == that.nodeCount
+        && processingUnits == that.processingUnits
         && state == that.state
         && Objects.equals(labels, that.labels);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(id, configId, displayName, nodeCount, state, labels);
+    return Objects.hash(id, configId, displayName, nodeCount, processingUnits, state, labels);
   }
 
   com.google.spanner.admin.instance.v1.Instance toProto() {
     com.google.spanner.admin.instance.v1.Instance.Builder builder =
         com.google.spanner.admin.instance.v1.Instance.newBuilder()
             .setName(getId().getName())
             .setNodeCount(getNodeCount())
+            .setProcessingUnits(getProcessingUnits())
             .putAllLabels(getLabels());
     if (getDisplayName() != null) {
       builder.setDisplayName(getDisplayName());

google-cloud-spanner/src/test/java/com/google/cloud/spanner/InstanceAdminClientImplTest.java

@@ -17,6 +17,9 @@
 package com.google.cloud.spanner;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.initMocks;
@@ -98,14 +101,24 @@ private com.google.spanner.admin.instance.v1.Instance getInstanceProto() {
         .setConfig(CONFIG_NAME)
         .setName(INSTANCE_NAME)
         .setNodeCount(1)
+        .setProcessingUnits(1000)
+        .build();
+  }
+
+  private com.google.spanner.admin.instance.v1.Instance getInstanceProtoWithProcessingUnits() {
+    return com.google.spanner.admin.instance.v1.Instance.newBuilder()
+        .setConfig(CONFIG_NAME)
+        .setName(INSTANCE_NAME)
+        .setProcessingUnits(10)
         .build();
   }
 
   private com.google.spanner.admin.instance.v1.Instance getAnotherInstanceProto() {
     return com.google.spanner.admin.instance.v1.Instance.newBuilder()
         .setConfig(CONFIG_NAME)
         .setName(INSTANCE_NAME2)
-        .setNodeCount(1)
+        .setNodeCount(2)
+        .setProcessingUnits(2000)
         .build();
   }
 
@@ -115,7 +128,10 @@ public void createInstance() throws Exception {
         rawOperationFuture =
             OperationFutureUtil.immediateOperationFuture(
                 "createInstance", getInstanceProto(), CreateInstanceMetadata.getDefaultInstance());
-    when(rpc.createInstance("projects/" + PROJECT_ID, INSTANCE_ID, getInstanceProto()))
+    when(rpc.createInstance(
+            "projects/" + PROJECT_ID,
+            INSTANCE_ID,
+            getInstanceProto().toBuilder().setProcessingUnits(0).build()))
         .thenReturn(rawOperationFuture);
     OperationFuture<Instance, CreateInstanceMetadata> op =
         client.createInstance(
@@ -128,9 +144,50 @@ public void createInstance() throws Exception {
   }
 
   @Test
-  public void getInstance() {
+  public void testCreateInstanceWithProcessingUnits() throws Exception {
+    OperationFuture<com.google.spanner.admin.instance.v1.Instance, CreateInstanceMetadata>
+        rawOperationFuture =
+            OperationFutureUtil.immediateOperationFuture(
+                "createInstance",
+                getInstanceProtoWithProcessingUnits(),
+                CreateInstanceMetadata.getDefaultInstance());
+    when(rpc.createInstance(
+            "projects/" + PROJECT_ID, INSTANCE_ID, getInstanceProtoWithProcessingUnits()))
+        .thenReturn(rawOperationFuture);
+    OperationFuture<Instance, CreateInstanceMetadata> operation =
+        client.createInstance(
+            InstanceInfo.newBuilder(InstanceId.of(PROJECT_ID, INSTANCE_ID))
+                .setInstanceConfigId(InstanceConfigId.of(PROJECT_ID, CONFIG_ID))
+                .setProcessingUnits(10)
+                .build());
+    assertTrue(operation.isDone());
+    assertEquals(INSTANCE_NAME, operation.get().getId().getName());
+  }
+
+  @Test
+  public void testCreateInstanceWithBothNodeCountAndProcessingUnits() throws Exception {
+    try {
+      client.createInstance(
+          InstanceInfo.newBuilder(InstanceId.of(PROJECT_ID, INSTANCE_ID))
+              .setInstanceConfigId(InstanceConfigId.of(PROJECT_ID, CONFIG_ID))
+              .setNodeCount(1)
+              .setProcessingUnits(100)
+              .build());
+      fail("missing expected exception");
+    } catch (IllegalArgumentException e) {
+      assertTrue(
+          e.getMessage()
+              .contains(
+                  "Only one of nodeCount and processingUnits can be set when creating a new instance"));
+    }
+  }
+
+  @Test
+  public void testGetInstance() {
     when(rpc.getInstance(INSTANCE_NAME)).thenReturn(getInstanceProto());
-    assertThat(client.getInstance(INSTANCE_ID).getId().getName()).isEqualTo(INSTANCE_NAME);
+    Instance instance = client.getInstance(INSTANCE_ID);
+    assertEquals(INSTANCE_NAME, instance.getId().getName());
+    assertEquals(1000, instance.getProcessingUnits());
   }
 
   @Test
@@ -165,7 +222,80 @@ public void updateInstanceMetadata() throws Exception {
   }
 
   @Test
-  public void listInstances() {
+  public void testUpdateInstanceProcessingUnits() throws Exception {
+    com.google.spanner.admin.instance.v1.Instance instance =
+        com.google.spanner.admin.instance.v1.Instance.newBuilder()
+            .setName(INSTANCE_NAME)
+            .setConfig(CONFIG_NAME)
+            .setProcessingUnits(10)
+            .build();
+    OperationFuture<com.google.spanner.admin.instance.v1.Instance, UpdateInstanceMetadata>
+        rawOperationFuture =
+            OperationFutureUtil.immediateOperationFuture(
+                "updateInstance",
+                getInstanceProtoWithProcessingUnits(),
+                UpdateInstanceMetadata.getDefaultInstance());
+    when(rpc.updateInstance(instance, FieldMask.newBuilder().addPaths("processing_units").build()))
+        .thenReturn(rawOperationFuture);
+    InstanceInfo instanceInfo =
+        InstanceInfo.newBuilder(InstanceId.of(INSTANCE_NAME))
+            .setInstanceConfigId(InstanceConfigId.of(CONFIG_NAME))
+            .setProcessingUnits(10)
+            .build();
+    OperationFuture<Instance, UpdateInstanceMetadata> operationWithFieldMask =
+        client.updateInstance(instanceInfo, InstanceInfo.InstanceField.PROCESSING_UNITS);
+    assertTrue(operationWithFieldMask.isDone());
+    assertEquals(INSTANCE_NAME, operationWithFieldMask.get().getId().getName());
+
+    when(rpc.updateInstance(
+            instance,
+            FieldMask.newBuilder()
+                .addAllPaths(Arrays.asList("display_name", "processing_units", "labels"))
+                .build()))
+        .thenReturn(rawOperationFuture);
+    OperationFuture<Instance, UpdateInstanceMetadata> operation =
+        client.updateInstance(instanceInfo);
+    assertTrue(operation.isDone());
+    assertEquals(INSTANCE_NAME, operation.get().getId().getName());
+  }
+
+  @Test
+  public void testUpdateInstanceWithNodeCountAndProcessingUnits() throws Exception {
+    com.google.spanner.admin.instance.v1.Instance instance =
+        com.google.spanner.admin.instance.v1.Instance.newBuilder()
+            .setName(INSTANCE_NAME)
+            .setConfig(CONFIG_NAME)
+            .setNodeCount(3)
+            .setProcessingUnits(3000)
+            .build();
+    OperationFuture<com.google.spanner.admin.instance.v1.Instance, UpdateInstanceMetadata>
+        rawOperationFuture =
+            OperationFutureUtil.immediateOperationFuture(
+                "updateInstance",
+                getInstanceProtoWithProcessingUnits(),
+                UpdateInstanceMetadata.getDefaultInstance());
+    // node_count should take precedence over processing_units when node_count>0 and no specific
+    // field mask is set by the caller.
+    when(rpc.updateInstance(
+            instance,
+            FieldMask.newBuilder()
+                .addAllPaths(Arrays.asList("display_name", "node_count", "labels"))
+                .build()))
+        .thenReturn(rawOperationFuture);
+    InstanceInfo instanceInfo =
+        InstanceInfo.newBuilder(InstanceId.of(INSTANCE_NAME))
+            .setInstanceConfigId(InstanceConfigId.of(CONFIG_NAME))
+            .setNodeCount(3)
+            .setProcessingUnits(3000)
+            .build();
+    OperationFuture<Instance, UpdateInstanceMetadata> operationWithFieldMask =
+        client.updateInstance(instanceInfo);
+    assertTrue(operationWithFieldMask.isDone());
+    assertEquals(INSTANCE_NAME, operationWithFieldMask.get().getId().getName());
+  }
+
+  @Test
+  public void testListInstances() {
     String nextToken = "token";
     String filter = "env:dev";
     when(rpc.listInstances(1, null, filter))
@@ -175,9 +305,11 @@ public void listInstances() {
     List<Instance> instances =
         Lists.newArrayList(
             client.listInstances(Options.pageSize(1), Options.filter(filter)).iterateAll());
-    assertThat(instances.get(0).getId().getName()).isEqualTo(INSTANCE_NAME);
-    assertThat(instances.get(1).getId().getName()).isEqualTo(INSTANCE_NAME2);
-    assertThat(instances.size()).isEqualTo(2);
+    assertEquals(INSTANCE_NAME, instances.get(0).getId().getName());
+    assertEquals(1000, instances.get(0).getProcessingUnits());
+    assertEquals(INSTANCE_NAME2, instances.get(1).getId().getName());
+    assertEquals(2000, instances.get(1).getProcessingUnits());
+    assertEquals(2, instances.size());
   }
 
   @Test