impl(bigtable): add MutateRowsLimiter

google/cloud/bigtable/CMakeLists.txt

@@ -193,6 +193,8 @@ add_library(
     internal/legacy_row_reader.h
     internal/logging_data_client.cc
     internal/logging_data_client.h
+    internal/mutate_rows_limiter.cc
+    internal/mutate_rows_limiter.h
     internal/prefix_range_end.cc
     internal/prefix_range_end.h
     internal/rate_limiter.cc
@@ -386,6 +388,7 @@ if (BUILD_TESTING)
         internal/legacy_bulk_mutator_test.cc
         internal/legacy_row_reader_test.cc
         internal/logging_data_client_test.cc
+        internal/mutate_rows_limiter_test.cc
         internal/prefix_range_end_test.cc
         internal/rate_limiter_test.cc
         internal/retry_traits_test.cc

google/cloud/bigtable/bigtable_client_unit_tests.bzl

@@ -59,6 +59,7 @@ bigtable_client_unit_tests = [
     "internal/legacy_bulk_mutator_test.cc",
     "internal/legacy_row_reader_test.cc",
     "internal/logging_data_client_test.cc",
+    "internal/mutate_rows_limiter_test.cc",
     "internal/prefix_range_end_test.cc",
     "internal/rate_limiter_test.cc",
     "internal/retry_traits_test.cc",

google/cloud/bigtable/google_cloud_cpp_bigtable.bzl

@@ -96,6 +96,7 @@ google_cloud_cpp_bigtable_hdrs = [
     "internal/legacy_async_row_sampler.h",
     "internal/legacy_row_reader.h",
     "internal/logging_data_client.h",
+    "internal/mutate_rows_limiter.h",
     "internal/prefix_range_end.h",
     "internal/rate_limiter.h",
     "internal/readrowsparser.h",
@@ -196,6 +197,7 @@ google_cloud_cpp_bigtable_srcs = [
     "internal/legacy_async_row_sampler.cc",
     "internal/legacy_row_reader.cc",
     "internal/logging_data_client.cc",
+    "internal/mutate_rows_limiter.cc",
     "internal/prefix_range_end.cc",
     "internal/rate_limiter.cc",
     "internal/readrowsparser.cc",

google/cloud/bigtable/internal/mutate_rows_limiter.cc

@@ -0,0 +1,68 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "google/cloud/bigtable/internal/mutate_rows_limiter.h"
+#include "google/cloud/bigtable/options.h"
+#include "google/cloud/internal/opentelemetry.h"
+#include <algorithm>
+#include <thread>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+namespace {
+
+// clamp value to the range [min, max]
+template <typename T>
+T Clamp(T value, T min, T max) {
+  return (std::min)(max, (std::max)(min, value));
+}
+
+}  // namespace
+
+void ThrottlingMutateRowsLimiter::Acquire() {
+  auto wait = limiter_.acquire(1);
+  throttled_since_last_update_ =
+      throttled_since_last_update_ || wait != Clock::duration::zero();
+  on_wait_(wait);
+}
+
+void ThrottlingMutateRowsLimiter::Update(
+    google::bigtable::v2::MutateRowsResponse const& response) {
+  if (!response.has_rate_limit_info()) return;
+  auto const now = clock_->Now();
+  if (now < next_update_) return;
+  auto const& info = response.rate_limit_info();
+  next_update_ = now + std::chrono::duration_cast<Clock::duration>(
+                           std::chrono::seconds(info.period().seconds()) +
+                           std::chrono::nanoseconds(info.period().nanos()));
+
+  // The effective QPS can lag behind the max QPS allowed by the rate limiter.
+  // In such a case, we should not keep increasing the max QPS allowed. We
+  // should only increase the ceiling if we are actually hitting that ceiling.
+  if (info.factor() > 1 && !throttled_since_last_update_) return;
+  throttled_since_last_update_ = false;
+
+  auto factor = Clamp(info.factor(), min_factor_, max_factor_);
+  auto const period = Clamp(
+      std::chrono::duration_cast<Clock::duration>(limiter_.period() / factor),
+      min_period_, max_period_);
+  limiter_.set_period(period);
+}
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google

google/cloud/bigtable/internal/mutate_rows_limiter.h

@@ -0,0 +1,95 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_MUTATE_ROWS_LIMITER_H
+#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_MUTATE_ROWS_LIMITER_H
+
+#include "google/cloud/bigtable/internal/rate_limiter.h"
+#include "google/cloud/internal/clock.h"
+#include "google/cloud/version.h"
+#include <google/bigtable/v2/bigtable.pb.h>
+#include <chrono>
+#include <functional>
+#include <memory>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+/// A Bigtable-specific wrapper over the more generic `RateLimiter`
+class MutateRowsLimiter {
+ public:
+  virtual void Acquire() = 0;
+  virtual void Update(
+      google::bigtable::v2::MutateRowsResponse const& response) = 0;
+};
+
+class NoopMutateRowsLimiter : public MutateRowsLimiter {
+ public:
+  void Acquire() override {}
+  void Update(google::bigtable::v2::MutateRowsResponse const&) override {}
+};
+
+class ThrottlingMutateRowsLimiter : public MutateRowsLimiter {
+ public:
+  using Clock = RateLimiter::Clock;
+  template <typename Rep1, typename Period1, typename Rep2, typename Period2,
+            typename Rep3, typename Period3>
+  explicit ThrottlingMutateRowsLimiter(
+      std::shared_ptr<Clock> clock,
+      std::function<void(Clock::duration)> on_wait,
+      std::chrono::duration<Rep1, Period1> initial_period,
+      std::chrono::duration<Rep2, Period2> min_period,
+      std::chrono::duration<Rep3, Period3> max_period, double min_factor,
+      double max_factor)
+      : clock_(std::move(clock)),
+        limiter_(clock_, initial_period),
+        on_wait_(std::move(on_wait)),
+        next_update_(clock_->Now()),
+        min_period_(std::chrono::duration_cast<Clock::duration>(min_period)),
+        max_period_(std::chrono::duration_cast<Clock::duration>(max_period)),
+        min_factor_(min_factor),
+        max_factor_(max_factor) {}
+
+  void Acquire() override;
+
+  /**
+   * As specified in:
+   * https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#google.bigtable.v2.RateLimitInfo
+   */
+  void Update(
+      google::bigtable::v2::MutateRowsResponse const& response) override;
+
+  // Exposed for testing
+  Clock::duration period() const { return limiter_.period(); }
+
+ private:
+  std::shared_ptr<Clock> clock_;
+  RateLimiter limiter_;
+  std::function<void(Clock::duration)> on_wait_;
+  bool throttled_since_last_update_ = false;
+  Clock::time_point next_update_;
+  Clock::duration min_period_;
+  Clock::duration max_period_;
+  double min_factor_;
+  double max_factor_;
+};
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google
+
+#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_MUTATE_ROWS_LIMITER_H