cleanup(bigtable): strictly enforce bigtable limits in mutation batcher (#7036)

dbolduc · web-flow · commit 46c9e78d02ac · 2021-07-21T12:23:33.000-04:00
diff --git a/google/cloud/bigtable/mutation_batcher.cc b/google/cloud/bigtable/mutation_batcher.cc
@@ -25,6 +25,7 @@ inline namespace BIGTABLE_CLIENT_NS {
 // Cloud Bigtable doesn't accept more than this in a single request.
 auto constexpr kBigtableMutationLimit = 100000;
 // Maximum mutations that can be outstanding in the Cloud Bigtable front end.
+// NOTE: this is a system-wide limit, but it is only enforced per process.
 auto constexpr kBigtableOutstandingMutationLimit = 300000;
 // Let's make the default slightly smaller, so that overheads or
 // miscalculations don't tip us over.
@@ -41,6 +42,20 @@ MutationBatcher::Options::Options()
       max_outstanding_size(kDefaultMaxOutstandingSize),
       max_outstanding_mutations(kBigtableOutstandingMutationLimit) {}
 
+MutationBatcher::Options& MutationBatcher::Options::SetMaxMutationsPerBatch(
+    size_t max_mutations_per_batch_arg) {
+  max_mutations_per_batch = std::min<std::size_t>(max_mutations_per_batch_arg,
+                                                  kBigtableMutationLimit);
+  return *this;
+}
+
+MutationBatcher::Options& MutationBatcher::Options::SetMaxOutstandingMutations(
+    size_t max_outstanding_mutations_arg) {
+  max_outstanding_mutations = std::min<std::size_t>(
+      max_outstanding_mutations_arg, kBigtableOutstandingMutationLimit);
+  return *this;
+}
+
 std::pair<future<void>, future<Status>> MutationBatcher::AsyncApply(
     CompletionQueue& cq, SingleRowMutation mut) {
   AdmissionPromise admission_promise;
diff --git a/google/cloud/bigtable/mutation_batcher.h b/google/cloud/bigtable/mutation_batcher.h
@@ -61,10 +61,7 @@ class MutationBatcher {
     Options();
 
     /// A single RPC will not have more mutations than this.
-    Options& SetMaxMutationsPerBatch(size_t max_mutations_per_batch_arg) {
-      max_mutations_per_batch = max_mutations_per_batch_arg;
-      return *this;
-    }
+    Options& SetMaxMutationsPerBatch(size_t max_mutations_per_batch_arg);
 
     /// Sum of mutations' sizes in a single RPC will not be larger than this.
     Options& SetMaxSizePerBatch(size_t max_size_per_batch_arg) {
@@ -85,10 +82,7 @@ class MutationBatcher {
     }
 
     /// MutationBatcher will at most admit this many mutations.
-    Options& SetMaxOutstandingMutations(size_t max_outstanding_mutations_arg) {
-      max_outstanding_mutations = max_outstanding_mutations_arg;
-      return *this;
-    }
+    Options& SetMaxOutstandingMutations(size_t max_outstanding_mutations_arg);
 
     std::size_t max_mutations_per_batch;
     std::size_t max_size_per_batch;
diff --git a/google/cloud/bigtable/mutation_batcher_test.cc b/google/cloud/bigtable/mutation_batcher_test.cc
@@ -254,11 +254,23 @@ TEST(OptionsTest, Trivial) {
                                      .SetMaxMutationsPerBatch(1)
                                      .SetMaxSizePerBatch(2)
                                      .SetMaxBatches(3)
-                                     .SetMaxOutstandingSize(4);
+                                     .SetMaxOutstandingSize(4)
+                                     .SetMaxOutstandingMutations(5);
   ASSERT_EQ(1, opt.max_mutations_per_batch);
   ASSERT_EQ(2, opt.max_size_per_batch);
   ASSERT_EQ(3, opt.max_batches);
   ASSERT_EQ(4, opt.max_outstanding_size);
+  ASSERT_EQ(5, opt.max_outstanding_mutations);
+}
+
+TEST(OptionsTest, StrictLimits) {
+  MutationBatcher::Options opt = MutationBatcher::Options()
+                                     .SetMaxMutationsPerBatch(200000)
+                                     .SetMaxOutstandingMutations(400000);
+  // See `kBigtableMutationLimit`
+  ASSERT_EQ(100000, opt.max_mutations_per_batch);
+  // See `kBigtableOutstandingMutationLimit`
+  ASSERT_EQ(300000, opt.max_outstanding_mutations);
 }
 
 TEST_F(MutationBatcherTest, TrivialTest) {
