clean up tests and formatting

Author: mutianf

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableBatchingCallSettings.java

@@ -199,8 +199,8 @@ public Builder setRetrySettings(@Nonnull RetrySettings retrySettings) {
  * reach the target rpc latency.
  */
  public Builder enableLatencyBasedThrottling(long targetRpcLatency) {
- Preconditions.checkArgument(targetRpcLatency > 0,
- "target RPC latency must be greater than 0");
+ Preconditions.checkArgument(
+ targetRpcLatency > 0, "target RPC latency must be greater than 0");
  this.isLatencyBasedThrottlingEnabled = true;
  this.targetRpcLatencyMs = targetRpcLatency;
  return this;

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlCallable.java

@@ -95,6 +95,7 @@ public void run() {
  dynamicFlowControlStats.updateLatency(timer.elapsed(TimeUnit.MILLISECONDS));
  long lastAdjustedTimestamp = dynamicFlowControlStats.getLastAdjustedTimestampMs();
  long now = System.currentTimeMillis();
+ // Avoid adjusting the thresholds too frequently
  if (now - lastAdjustedTimestamp < adjustingIntervalMs) {
  return;
  }
@@ -105,17 +106,25 @@ public void run() {
  : (now - flowControlEvents.getLastFlowControlEvent().getTimestampMs()
  <= TimeUnit.MINUTES.toMillis(5));
  if (meanLatency > targetLatency * 3) {
+ // Decrease at 30% of the maximum
  decrease(
  lastAdjustedTimestamp, now, flowController.getMaxOutstandingElementCount() * 3 / 10);
  } else if (meanLatency > targetLatency * 1.2) {
+ // Decrease at 10% of the maximum
  decrease(lastAdjustedTimestamp, now, flowController.getMaxOutstandingElementCount() / 10);
  } else if (throttled && meanLatency < targetLatency * 0.8) {
+ // if latency is low, and there was throttling, then increase the parallelism so that new
+ // calls will not be throttled.
+
+ // Increase parallelism at a slower than we decrease. The lower rate should help the
+ // system maintain stability.
  increase(
  lastAdjustedTimestamp, now, flowController.getMaxOutstandingElementCount() * 5 / 100);
  } else if (throttled
  && flowController.getCurrentOutstandingElementCount()
  < flowController.getMaxOutstandingElementCount() * 0.05
  && meanLatency < 2 * targetLatency) {
+ // If latency is low and threshold is low, increase slowly.
  increase(
  lastAdjustedTimestamp, now, flowController.getMaxOutstandingElementCount() * 2 / 100);
  }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlStats.java

@@ -27,12 +27,11 @@
  * last timestamp when the thresholds in {@link FlowController} are updated.
  */
 @InternalApi
-public class DynamicFlowControlStats {
+public final class DynamicFlowControlStats {
 
  private static final double DEFAULT_DECAY_CONSTANT = 0.015; // Biased to the past 5 minutes
 
  private AtomicLong lastAdjustedTimestampMs;
- private AtomicLong adjustedCounter;
  private DecayingAverage meanLatency;
 
  DynamicFlowControlStats() {
@@ -41,7 +40,6 @@ public class DynamicFlowControlStats {
 
  DynamicFlowControlStats(double decayConstant) {
  this.lastAdjustedTimestampMs = new AtomicLong(0);
- this.adjustedCounter = new AtomicLong(0);
  this.meanLatency = new DecayingAverage(decayConstant);
  }
 
@@ -67,13 +65,8 @@ public long getLastAdjustedTimestampMs() {
  return lastAdjustedTimestampMs.get();
  }
 
- public long getAdjustedCounter() {
- return adjustedCounter.get();
- }
-
  boolean setLastAdjustedTimestampMs(long last, long now) {
  if (lastAdjustedTimestampMs.compareAndSet(last, now)) {
- adjustedCounter.addAndGet(1);
  return true;
  }
  return false;

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -19,8 +19,8 @@
 import com.google.api.core.InternalApi;
 import com.google.api.gax.batching.Batcher;
 import com.google.api.gax.batching.BatcherImpl;
-import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.batching.FlowControlEventStats;
+import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.core.BackgroundResource;
 import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.core.GaxProperties;

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/BulkMutateIT.java

@@ -17,28 +17,15 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import com.google.api.core.ApiFuture;
-import com.google.api.gax.batching.Batcher;
 import com.google.api.gax.batching.BatcherImpl;
-import com.google.api.gax.batching.BatchingSettings;
-import com.google.api.gax.batching.DynamicFlowControlSettings;
 import com.google.api.gax.batching.FlowControlEventStats;
-import com.google.api.gax.batching.FlowControlSettings;
-import com.google.api.gax.batching.FlowController;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
 import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
 import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.Row;
 import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
-import com.google.cloud.bigtable.data.v2.stub.DynamicFlowControlStats;
-import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
 import com.google.cloud.bigtable.test_helpers.env.TestEnvRule;
-import com.google.protobuf.ByteString;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import java.util.concurrent.TimeUnit;
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -47,10 +34,9 @@
 @RunWith(JUnit4.class)
 public class BulkMutateIT {
 
- @ClassRule
- public static TestEnvRule testEnvRule = new TestEnvRule();
+ @ClassRule public static TestEnvRule testEnvRule = new TestEnvRule();
 
- @Test (timeout = 20 * 1000)
+ @Test(timeout = 10 * 1000)
  public void test() throws IOException, InterruptedException {
  BigtableDataSettings settings = testEnvRule.env().getDataClientSettings();
  // Set target latency really low so it'll trigger adjusting thresholds
@@ -61,26 +47,24 @@ public void test() throws IOException, InterruptedException {
  (BatcherImpl) client.newBulkMutationBatcher(testEnvRule.env().getTableId());
  try {
  FlowControlEventStats events = batcher.getFlowControlEventStats();
- FlowController flowController = batcher.getFlowController();
- // DynamicFlowControlStats dynamicFlowControlStats =
- // settings.getStubSettings().bulkMutateRowsSettings().toBuilder()
- // .getDynamicFlowControlStats();
-
- assertThat(batcher.getFlowController().getCurrentOutstandingElementCount()).isNotEqualTo(
- batcher.getFlowController().getMinOutstandingElementCount());
- assertThat(batcher.getFlowController().getCurrentOutstandingElementCount()).isNotEqualTo(
- batcher.getFlowController().getMaxOutstandingElementCount());
+ long initialThreashold = batcher.getFlowController().getCurrentOutstandingElementCount();
+ assertThat(batcher.getFlowController().getCurrentOutstandingElementCount())
+ .isNotEqualTo(batcher.getFlowController().getMinOutstandingElementCount());
+ assertThat(batcher.getFlowController().getCurrentOutstandingElementCount())
+ .isNotEqualTo(batcher.getFlowController().getMaxOutstandingElementCount());
 
  String familyId = testEnvRule.env().getFamilyId();
  long initial = batcher.getFlowController().getCurrentOutstandingElementCount();
  for (long i = 0; i < initial * 3; i++) {
  String key = "test-key" + i;
- batcher.add(
- RowMutationEntry.create(key).setCell(familyId, "qualifier", i));
+ batcher.add(RowMutationEntry.create(key).setCell(familyId, "qualifier", i));
  }
  batcher.flush();
  assertThat(events.getLastFlowControlEvent()).isNotNull();
- //assertThat(dynamicFlowControlStats.getAdjustedCounter()).isEqualTo(1);
+ // Verify that the threshold is adjusted
+ assertThat(batcher.getFlowController().getCurrentOutstandingElementCount())
+ .isNotEqualTo(initialThreashold);
+ // Query a key to make sure the write succeeded
  Row row =
  testEnvRule
  .env()

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java

@@ -78,8 +78,6 @@ public void settingsAreNotLostTest() {
  WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class);
  Duration watchdogInterval = Duration.ofSeconds(12);
  HeaderTracer headerTracer = Mockito.mock(HeaderTracer.class);
- boolean isBatchMutationLatencyBasedThrottlingEnabled = true;
- long batchMutationTargetLatency = 10;
 
  EnhancedBigtableStubSettings.Builder builder =
  EnhancedBigtableStubSettings.newBuilder()
@@ -103,9 +101,7 @@ public void settingsAreNotLostTest() {
  credentialsProvider,
  watchdogProvider,
  watchdogInterval,
- headerTracer,
- isBatchMutationLatencyBasedThrottlingEnabled,
- batchMutationTargetLatency);
+ headerTracer);
  verifySettings(
  builder.build(),
  projectId,
@@ -116,9 +112,7 @@ public void settingsAreNotLostTest() {
  credentialsProvider,
  watchdogProvider,
  watchdogInterval,
- headerTracer,
- isBatchMutationLatencyBasedThrottlingEnabled,
- batchMutationTargetLatency);
+ headerTracer);
  verifyBuilder(
  builder.build().toBuilder(),
  projectId,
@@ -129,9 +123,7 @@ public void settingsAreNotLostTest() {
  credentialsProvider,
  watchdogProvider,
  watchdogInterval,
- headerTracer,
- isBatchMutationLatencyBasedThrottlingEnabled,
- batchMutationTargetLatency);
+ headerTracer);
  }
 
  private void verifyBuilder(
@@ -144,9 +136,7 @@ private void verifyBuilder(
  CredentialsProvider credentialsProvider,
  WatchdogProvider watchdogProvider,
  Duration watchdogInterval,
- HeaderTracer headerTracer,
- boolean isBatchMutationLatencyBasedThrottlingEnabled,
- long batchMutationTargetLatency) {
+ HeaderTracer headerTracer) {
  assertThat(builder.getProjectId()).isEqualTo(projectId);
  assertThat(builder.getInstanceId()).isEqualTo(instanceId);
  assertThat(builder.getAppProfileId()).isEqualTo(appProfileId);
@@ -156,12 +146,6 @@ private void verifyBuilder(
  assertThat(builder.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
  assertThat(builder.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
  assertThat(builder.getHeaderTracer()).isEqualTo(headerTracer);
- assertThat(builder.isLatencyBasedThrottlingForBatchMutationEnabled())
- .isEqualTo(isBatchMutationLatencyBasedThrottlingEnabled);
- assertThat(builder.bulkMutateRowsSettings().isLatencyBasedThrottlingEnabled())
- .isEqualTo(isBatchMutationLatencyBasedThrottlingEnabled);
- assertThat(builder.bulkMutateRowsSettings().getTargetRpcLatencyMs())
- .isEqualTo(batchMutationTargetLatency);
  }
 
  private void verifySettings(
@@ -174,9 +158,7 @@ private void verifySettings(
  CredentialsProvider credentialsProvider,
  WatchdogProvider watchdogProvider,
  Duration watchdogInterval,
- HeaderTracer headerTracer,
- boolean isBatchMutationLatencyBasedThrottlingEnabled,
- long batchMutationTargetLatency) {
+ HeaderTracer headerTracer) {
  assertThat(settings.getProjectId()).isEqualTo(projectId);
  assertThat(settings.getInstanceId()).isEqualTo(instanceId);
  assertThat(settings.getAppProfileId()).isEqualTo(appProfileId);
@@ -186,10 +168,6 @@ private void verifySettings(
  assertThat(settings.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
  assertThat(settings.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
  assertThat(settings.getHeaderTracer()).isEqualTo(headerTracer);
- assertThat(settings.bulkMutateRowsSettings().isLatencyBasedThrottlingEnabled())
- .isEqualTo(isBatchMutationLatencyBasedThrottlingEnabled);
- assertThat(settings.bulkMutateRowsSettings().getTargetRpcLatencyMs())
- .isEqualTo(batchMutationTargetLatency);
  }
 
  @Test
@@ -493,7 +471,6 @@ public void bulkMutateRowsSettingsAreNotLostTest() {
  .isSameInstanceAs(batchingSettings);
  assertThat(builder.bulkMutateRowsSettings().isLatencyBasedThrottlingEnabled()).isTrue();
  assertThat(builder.bulkMutateRowsSettings().getTargetRpcLatencyMs()).isEqualTo(targetLatency);
-
  assertThat(
  builder
  .bulkMutateRowsSettings()

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java

@@ -183,11 +183,16 @@ public void testBulkMutationFlowControllerConfigured() throws Exception {
  EnhancedBigtableStub stub1 = EnhancedBigtableStub.create(settings.build().getStubSettings());
  EnhancedBigtableStub stub2 = EnhancedBigtableStub.create(settings.build().getStubSettings());
 
+ // Creating 2 batchers from the same stub, they should share the same FlowController and
+ // FlowControlEventStats
  try (BatcherImpl batcher1 = (BatcherImpl) stub1.newMutateRowsBatcher("my-table1");
  BatcherImpl batcher2 = (BatcherImpl) stub1.newMutateRowsBatcher("my-table2")) {
  assertThat(batcher1.getFlowController()).isNotNull();
+ assertThat(batcher1.getFlowControlEventStats()).isNotNull();
  assertThat(batcher1).isNotSameInstanceAs(batcher2);
  assertThat(batcher1.getFlowController()).isSameInstanceAs(batcher2.getFlowController());
+ assertThat(batcher1.getFlowControlEventStats())
+ .isSameInstanceAs(batcher2.getFlowControlEventStats());
  // Verify flow controller settings
  assertThat(batcher1.getFlowController().getLimitExceededBehavior())
  .isEqualTo(LimitExceededBehavior.Block);
@@ -205,9 +210,15 @@ public void testBulkMutationFlowControllerConfigured() throws Exception {
  assertThat(batcher1.getFlowController().getMinOutstandingRequestBytes()).isEqualTo(1000L);
  }
 
+ // Creating 2 batchers from different stubs, they should not share the same FlowController and
+ // FlowControlEventStats
  try (BatcherImpl batcher1 = (BatcherImpl) stub1.newMutateRowsBatcher("my-table1");
  BatcherImpl batcher2 = (BatcherImpl) stub2.newMutateRowsBatcher("my-table2")) {
+ assertThat(batcher1.getFlowController()).isNotNull();
+ assertThat(batcher1.getFlowControlEventStats()).isNotNull();
  assertThat(batcher1.getFlowController()).isNotSameInstanceAs(batcher2.getFlowController());
+ assertThat(batcher1.getFlowControlEventStats())
+ .isNotSameInstanceAs(batcher2.getFlowControlEventStats());
  }
 
  stub2 =