Skip to content

Commit

Permalink
feat: add support for instance processing units (#665)
Browse files Browse the repository at this point in the history
* 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 <thiagotnunes@google.com>
  • Loading branch information
olavloite and thiagotnunes committed Jun 21, 2021
1 parent 9641bc1 commit 9c1c8e9
Show file tree
Hide file tree
Showing 7 changed files with 345 additions and 12 deletions.
Expand Up @@ -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);
Expand Down Expand Up @@ -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:
Expand Down
Expand Up @@ -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 =
Expand Down Expand Up @@ -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>
Expand Down
Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand All @@ -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;

Expand All @@ -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);
}
Expand All @@ -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;
Expand Down Expand Up @@ -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;

Expand All @@ -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);
}
Expand All @@ -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;
Expand All @@ -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();
Expand All @@ -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());
Expand Down
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
}

Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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
Expand Down

0 comments on commit 9c1c8e9

Please sign in to comment.