feat: Add APIs to enable request priorities (#1959)

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Rollback plan is reviewed and LGTMed Fixes #<issue_number_goes_here> ☕️ If you write sample code, please follow the [samples format]( https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
googleapis · Oct 26, 2023 · befd140 · befd140
1 parent bf5a9b7
commit befd140
Show file tree

Hide file tree

Showing 8 changed files with 393 additions and 1 deletion.
diff --git a/...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/AppProfile.java b/...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/AppProfile.java
@@ -17,6 +17,8 @@
 
 import com.google.api.core.InternalApi;
 import com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny;
+import com.google.bigtable.admin.v2.AppProfile.Priority;
+import com.google.bigtable.admin.v2.AppProfile.StandardIsolation;
 import com.google.bigtable.admin.v2.AppProfileName;
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
@@ -76,6 +78,15 @@ public RoutingPolicy getPolicy() {
     }
   }
 
+  public IsolationPolicy getIsolationPolicy() {
+    if (proto.hasStandardIsolation()) {
+      return new StandardIsolationPolicy(proto.getStandardIsolation());
+    } else {
+      // Should never happen because the constructor verifies that one must exist.
+      throw new IllegalStateException();
+    }
+  }
+
   /** Gets the id of this AppProfile. */
   @SuppressWarnings("WeakerAccess")
   public String getId() {
@@ -292,4 +303,110 @@ public int hashCode() {
       return Objects.hashCode(proto);
     }
   }
+
+  /** Represents the options for isolating this app profile's traffic from other use cases. */
+  @SuppressWarnings("WeakerAccess")
+  public interface IsolationPolicy {}
+
+  /**
+   * The possible priorities for an app profile. Note that higher priority writes can sometimes
+   * queue behind lower priority writes to the same tablet, as writes must be strictly sequenced in
+   * the durability log.
+   */
+  public static enum Priority {
+    LOW(com.google.bigtable.admin.v2.AppProfile.Priority.PRIORITY_LOW),
+    MEDIUM(com.google.bigtable.admin.v2.AppProfile.Priority.PRIORITY_MEDIUM),
+    HIGH(com.google.bigtable.admin.v2.AppProfile.Priority.PRIORITY_HIGH);
+
+    private final com.google.bigtable.admin.v2.AppProfile.Priority proto;
+
+    /**
+     * Wraps the protobuf. This method is considered an internal implementation detail and not meant
+     * to be used by applications.
+     */
+    @InternalApi
+    public static Priority fromProto(com.google.bigtable.admin.v2.AppProfile.Priority proto) {
+      Preconditions.checkNotNull(proto);
+
+      for (Priority priority : values()) {
+        if (priority.proto.equals(proto)) {
+          return priority;
+        }
+      }
+
+      throw new IllegalArgumentException("Unknown priority: " + proto);
+    }
+
+    Priority(com.google.bigtable.admin.v2.AppProfile.Priority proto) {
+      this.proto = proto;
+    }
+
+    /**
+     * Creates the request protobuf. This method is considered an internal implementation detail and
+     * not meant to be used by applications.
+     */
+    @InternalApi
+    public com.google.bigtable.admin.v2.AppProfile.Priority toProto() {
+      return proto;
+    }
+  }
+
+  /**
+   * A standard {@link IsolationPolicy} for isolating this app profile's traffic from other use
+   * cases. This accomplished by assigning different priorities to app profiles. A request that uses
+   * an app profile with a StandardIsolationPolicy with a HIGH priority will likely run before a
+   * request with a LOW priority.
+   */
+  public static class StandardIsolationPolicy implements IsolationPolicy {
+    private final StandardIsolation proto;
+
+    /** Creates a new instance of {@link StandardIsolationPolicy}. */
+    public static StandardIsolationPolicy of() {
+      return new StandardIsolationPolicy(StandardIsolation.getDefaultInstance());
+    }
+
+    /** Creates a new instance of {@link StandardIsolationPolicy} with the specified priority. */
+    public static StandardIsolationPolicy of(Priority priority) {
+      return new StandardIsolationPolicy(
+          StandardIsolation.newBuilder().setPriority(priority.toProto()).build());
+    }
+
+    /*
+     * Returns the priority for this app profile.
+     */
+    public Priority getPriority() {
+      return Priority.fromProto(proto.getPriority());
+    }
+
+    private StandardIsolationPolicy(
+        com.google.bigtable.admin.v2.AppProfile.StandardIsolation proto) {
+      this.proto = proto;
+    }
+
+    /**
+     * Creates the request protobuf. This method is considered an internal implementation detail and
+     * not meant to be used by applications.
+     */
+    @InternalApi
+    com.google.bigtable.admin.v2.AppProfile.StandardIsolation toProto() {
+      return proto;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      StandardIsolationPolicy that = (StandardIsolationPolicy) o;
+      return Objects.equal(proto, that.proto);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(proto);
+    }
+  }
 }
diff --git a/...able/src/main/java/com/google/cloud/bigtable/admin/v2/models/CreateAppProfileRequest.java b/...able/src/main/java/com/google/cloud/bigtable/admin/v2/models/CreateAppProfileRequest.java
@@ -17,9 +17,11 @@
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
+import com.google.cloud.bigtable.admin.v2.models.AppProfile.IsolationPolicy;
 import com.google.cloud.bigtable.admin.v2.models.AppProfile.MultiClusterRoutingPolicy;
 import com.google.cloud.bigtable.admin.v2.models.AppProfile.RoutingPolicy;
 import com.google.cloud.bigtable.admin.v2.models.AppProfile.SingleClusterRoutingPolicy;
+import com.google.cloud.bigtable.admin.v2.models.AppProfile.StandardIsolationPolicy;
 import com.google.common.base.Preconditions;
 import javax.annotation.Nonnull;
 
@@ -92,6 +94,21 @@ public CreateAppProfileRequest setRoutingPolicy(RoutingPolicy routingPolicy) {
     return this;
   }
 
+  /** Sets the isolation policy for all read/write requests that use this app profile. */
+  public CreateAppProfileRequest setIsolationPolicy(IsolationPolicy isolationPolicy) {
+    Preconditions.checkNotNull(isolationPolicy);
+
+    if (isolationPolicy instanceof StandardIsolationPolicy) {
+      proto
+          .getAppProfileBuilder()
+          .setStandardIsolation(((StandardIsolationPolicy) isolationPolicy).toProto());
+    } else {
+      throw new IllegalArgumentException("Unknown policy type: " + isolationPolicy);
+    }
+
+    return this;
+  }
+
   /**
    * Creates the request protobuf. This method is considered an internal implementation detail and
    * not meant to be used by applications.

diff --git a/...able/src/main/java/com/google/cloud/bigtable/admin/v2/models/UpdateAppProfileRequest.java b/...able/src/main/java/com/google/cloud/bigtable/admin/v2/models/UpdateAppProfileRequest.java
@@ -17,9 +17,11 @@
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
+import com.google.cloud.bigtable.admin.v2.models.AppProfile.IsolationPolicy;
 import com.google.cloud.bigtable.admin.v2.models.AppProfile.MultiClusterRoutingPolicy;
 import com.google.cloud.bigtable.admin.v2.models.AppProfile.RoutingPolicy;
 import com.google.cloud.bigtable.admin.v2.models.AppProfile.SingleClusterRoutingPolicy;
+import com.google.cloud.bigtable.admin.v2.models.AppProfile.StandardIsolationPolicy;
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.FieldMask;
@@ -121,6 +123,22 @@ public UpdateAppProfileRequest setRoutingPolicy(@Nonnull RoutingPolicy routingPo
     return this;
   }
 
+  /** Sets the isolation policy for all read/write requests that use this app profile. */
+  public UpdateAppProfileRequest setIsolationPolicy(@Nonnull IsolationPolicy isolationPolicy) {
+    Preconditions.checkNotNull(isolationPolicy);
+
+    if (isolationPolicy instanceof StandardIsolationPolicy) {
+      proto
+          .getAppProfileBuilder()
+          .setStandardIsolation(((StandardIsolationPolicy) isolationPolicy).toProto());
+      updateFieldMask(com.google.bigtable.admin.v2.AppProfile.STANDARD_ISOLATION_FIELD_NUMBER);
+    } else {
+      throw new IllegalArgumentException("Unknown policy type: " + isolationPolicy);
+    }
+
+    return this;
+  }
+
   private void updateFieldMask(int fieldNumber) {
     FieldMask newMask =
         FieldMaskUtil.fromFieldNumbers(com.google.bigtable.admin.v2.AppProfile.class, fieldNumber);

diff --git a/...le/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableInstanceAdminClientTests.java b/...le/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableInstanceAdminClientTests.java
@@ -43,6 +43,8 @@
 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.AppProfile.MultiClusterRoutingPolicy;
+import com.google.cloud.bigtable.admin.v2.models.AppProfile.Priority;
+import com.google.cloud.bigtable.admin.v2.models.AppProfile.StandardIsolationPolicy;
 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;
@@ -82,7 +84,7 @@
 
 @RunWith(JUnit4.class)
 /**
- * Tests for {@link BigtableTableAdminClient}. This test class uses Mockito so it has been
+ * Tests for {@link BigtableInstanceAdminClient}. This test class uses Mockito so it has been
  * explicitly excluded from Native Image testing by not following the naming convention of (IT* and
  * *ClientTest).
  */
@@ -983,6 +985,55 @@ public void testCreateAppProfileAddMultipleClusterIdsWithList() {
     assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
   }
 
+  @Test
+  public void testCreateAppProfileAddPriority() {
+    // Setup
+    Mockito.when(mockStub.createAppProfileCallable()).thenReturn(mockCreateAppProfileCallable);
+
+    com.google.bigtable.admin.v2.CreateAppProfileRequest expectedRequest =
+        com.google.bigtable.admin.v2.CreateAppProfileRequest.newBuilder()
+            .setParent(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID))
+            .setAppProfileId(APP_PROFILE_ID)
+            .setAppProfile(
+                com.google.bigtable.admin.v2.AppProfile.newBuilder()
+                    .setDescription("my description")
+                    .setMultiClusterRoutingUseAny(
+                        com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                            .newBuilder()
+                            .addClusterIds("cluster-id-1"))
+                    .setStandardIsolation(
+                        com.google.bigtable.admin.v2.AppProfile.StandardIsolation.newBuilder()
+                            .setPriority(
+                                com.google.bigtable.admin.v2.AppProfile.Priority.PRIORITY_MEDIUM)))
+            .build();
+
+    com.google.bigtable.admin.v2.AppProfile expectedResponse =
+        com.google.bigtable.admin.v2.AppProfile.newBuilder()
+            .setName(APP_PROFILE_NAME)
+            .setDescription("my description")
+            .setMultiClusterRoutingUseAny(
+                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny.newBuilder()
+                    .addClusterIds("cluster-id-1"))
+            .setStandardIsolation(
+                com.google.bigtable.admin.v2.AppProfile.StandardIsolation.newBuilder()
+                    .setPriority(com.google.bigtable.admin.v2.AppProfile.Priority.PRIORITY_MEDIUM))
+            .build();
+
+    Mockito.when(mockCreateAppProfileCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(expectedResponse));
+
+    // Execute
+    AppProfile actualResult =
+        adminClient.createAppProfile(
+            CreateAppProfileRequest.of(INSTANCE_ID, APP_PROFILE_ID)
+                .setDescription("my description")
+                .setRoutingPolicy(MultiClusterRoutingPolicy.of("cluster-id-1"))
+                .setIsolationPolicy(StandardIsolationPolicy.of(Priority.MEDIUM)));
+
+    // Verify
+    assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
+  }
+
   @Test
   public void testGetAppProfile() {
     // Setup
@@ -1101,6 +1152,47 @@ public void testUpdateAppProfile() {
     assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
   }
 
+  @Test
+  public void testUpdateAppProfileStandardIsolation() {
+    // Setup
+    Mockito.when(mockStub.updateAppProfileOperationCallable())
+        .thenReturn(mockUpdateAppProfileCallable);
+
+    com.google.bigtable.admin.v2.UpdateAppProfileRequest expectedRequest =
+        com.google.bigtable.admin.v2.UpdateAppProfileRequest.newBuilder()
+            .setAppProfile(
+                com.google.bigtable.admin.v2.AppProfile.newBuilder()
+                    .setName(APP_PROFILE_NAME)
+                    .setStandardIsolation(
+                        com.google.bigtable.admin.v2.AppProfile.StandardIsolation.newBuilder()
+                            .setPriority(
+                                com.google.bigtable.admin.v2.AppProfile.Priority.PRIORITY_LOW)))
+            .setUpdateMask(FieldMask.newBuilder().addPaths("standard_isolation"))
+            .build();
+
+    com.google.bigtable.admin.v2.AppProfile expectedResponse =
+        com.google.bigtable.admin.v2.AppProfile.newBuilder()
+            .setName(APP_PROFILE_NAME)
+            .setMultiClusterRoutingUseAny(
+                com.google.bigtable.admin.v2.AppProfile.MultiClusterRoutingUseAny
+                    .getDefaultInstance())
+            .setStandardIsolation(
+                com.google.bigtable.admin.v2.AppProfile.StandardIsolation.newBuilder()
+                    .setPriority(com.google.bigtable.admin.v2.AppProfile.Priority.PRIORITY_LOW))
+            .build();
+
+    mockOperationResult(mockUpdateAppProfileCallable, expectedRequest, expectedResponse);
+
+    // Execute
+    AppProfile actualResult =
+        adminClient.updateAppProfile(
+            UpdateAppProfileRequest.of(INSTANCE_ID, APP_PROFILE_ID)
+                .setIsolationPolicy(StandardIsolationPolicy.of(Priority.LOW)));
+
+    // Verify
+    assertThat(actualResult).isEqualTo(AppProfile.fromProto(expectedResponse));
+  }
+
   @Test
   public void testDeleteAppProfile() throws Exception {
     // Setup

diff --git a/...le/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableInstanceAdminClientIT.java b/...le/src/test/java/com/google/cloud/bigtable/admin/v2/it/BigtableInstanceAdminClientIT.java
@@ -191,6 +191,57 @@ public void appProfileTestMultiClusterWithIds() {
     }
   }
 
+  @Test
+  public void appProfileTestPriority() {
+    String newInstanceId = prefixGenerator.newPrefix();
+    String newClusterId = newInstanceId + "-c1";
+
+    client.createInstance(
+        CreateInstanceRequest.of(newInstanceId)
+            .addCluster(newClusterId, testEnvRule.env().getPrimaryZone(), 1, StorageType.SSD)
+            .setDisplayName("Priority-Instance-Test")
+            .addLabel("state", "readytodelete")
+            .setType(Type.PRODUCTION));
+
+    try {
+      assertThat(client.exists(newInstanceId)).isTrue();
+
+      String testAppProfile = prefixGenerator.newPrefix();
+
+      // This should be created with HIGH priority.
+      CreateAppProfileRequest request =
+          CreateAppProfileRequest.of(newInstanceId, testAppProfile)
+              .setRoutingPolicy(AppProfile.SingleClusterRoutingPolicy.of(newClusterId))
+              .setDescription("This is to test app profile");
+
+      AppProfile newlyCreatedAppProfile = client.createAppProfile(request);
+      AppProfile.StandardIsolationPolicy newlyCreatedAppProfilePolicy =
+          (AppProfile.StandardIsolationPolicy) newlyCreatedAppProfile.getIsolationPolicy();
+      assertThat(newlyCreatedAppProfilePolicy.getPriority()).isEqualTo(AppProfile.Priority.HIGH);
+
+      AppProfile updated =
+          client.updateAppProfile(
+              UpdateAppProfileRequest.of(newlyCreatedAppProfile)
+                  .setIsolationPolicy(
+                      AppProfile.StandardIsolationPolicy.of(AppProfile.Priority.LOW)));
+
+      AppProfile freshAppProfile = client.getAppProfile(newInstanceId, testAppProfile);
+      AppProfile.StandardIsolationPolicy freshAppProfilePolicy =
+          (AppProfile.StandardIsolationPolicy) freshAppProfile.getIsolationPolicy();
+      AppProfile.StandardIsolationPolicy updatedAppProfilePolicy =
+          (AppProfile.StandardIsolationPolicy) updated.getIsolationPolicy();
+
+      assertThat(freshAppProfilePolicy.getPriority()).isEqualTo(AppProfile.Priority.LOW);
+      assertThat(freshAppProfilePolicy).isEqualTo(updatedAppProfilePolicy);
+
+      assertThat(client.listAppProfiles(newInstanceId)).contains(freshAppProfile);
+    } finally {
+      if (client.exists(newInstanceId)) {
+        client.deleteInstance(newInstanceId);
+      }
+    }
+  }
+
   @Test
   public void iamUpdateTest() {
     Policy policy = client.getIamPolicy(instanceId);