impl(bigtable): add RandomTwoLeastUsed stub decorator

google/cloud/bigtable/CMakeLists.txt

@@ -155,6 +155,8 @@ add_library(
     internal/bigtable_logging_decorator.h
     internal/bigtable_metadata_decorator.cc
     internal/bigtable_metadata_decorator.h
+    internal/bigtable_random_two_least_used_decorator.cc
+    internal/bigtable_random_two_least_used_decorator.h
     internal/bigtable_round_robin_decorator.cc
     internal/bigtable_round_robin_decorator.h
     internal/bigtable_stub.cc
@@ -444,6 +446,7 @@ if (BUILD_TESTING)
         internal/async_row_sampler_test.cc
         internal/async_streaming_read_test.cc
         internal/bigtable_channel_refresh_test.cc
+        internal/bigtable_random_two_least_used_decorator_test.cc
         internal/bigtable_stub_factory_test.cc
         internal/bulk_mutator_test.cc
         internal/channel_usage_test.cc

google/cloud/bigtable/bigtable_client_unit_tests.bzl

@@ -42,6 +42,7 @@ bigtable_client_unit_tests = [
     "internal/async_row_sampler_test.cc",
     "internal/async_streaming_read_test.cc",
     "internal/bigtable_channel_refresh_test.cc",
+    "internal/bigtable_random_two_least_used_decorator_test.cc",
     "internal/bigtable_stub_factory_test.cc",
     "internal/bulk_mutator_test.cc",
     "internal/channel_usage_test.cc",

google/cloud/bigtable/google_cloud_cpp_bigtable.bzl

@@ -74,6 +74,7 @@ google_cloud_cpp_bigtable_hdrs = [
     "internal/bigtable_channel_refresh.h",
     "internal/bigtable_logging_decorator.h",
     "internal/bigtable_metadata_decorator.h",
+    "internal/bigtable_random_two_least_used_decorator.h",
     "internal/bigtable_round_robin_decorator.h",
     "internal/bigtable_stub.h",
     "internal/bigtable_stub_factory.h",
@@ -190,6 +191,7 @@ google_cloud_cpp_bigtable_srcs = [
     "internal/bigtable_channel_refresh.cc",
     "internal/bigtable_logging_decorator.cc",
     "internal/bigtable_metadata_decorator.cc",
+    "internal/bigtable_random_two_least_used_decorator.cc",
     "internal/bigtable_round_robin_decorator.cc",
     "internal/bigtable_stub.cc",
     "internal/bigtable_stub_factory.cc",

google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc

@@ -0,0 +1,343 @@
+// Copyright 2026 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/bigtable_random_two_least_used_decorator.h"
+#include "google/cloud/internal/async_streaming_read_rpc.h"
+#include "google/cloud/internal/streaming_read_rpc.h"
+#include <functional>
+#include <memory>
+#include <optional>
+#include <utility>
+
+namespace google {
+namespace cloud {
+namespace bigtable_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+namespace {
+
+template <typename T>
+class StreamingReadRpcTracking
+    : public google::cloud::internal::StreamingReadRpc<T> {
+ public:
+  StreamingReadRpcTracking(
+      std::unique_ptr<google::cloud::internal::StreamingReadRpc<T>> child,
+      std::function<void(void)> on_destruction)
+      : child_(std::move(child)), on_destruction_(std::move(on_destruction)) {}
+
+  ~StreamingReadRpcTracking() override { on_destruction_(); }
+
+  void Cancel() override { child_->Cancel(); }
+  std::optional<Status> Read(T* response) override {
+    return child_->Read(response);
+  }
+  RpcMetadata GetRequestMetadata() const override {
+    return child_->GetRequestMetadata();
+  }
+
+ private:
+  std::unique_ptr<google::cloud::internal::StreamingReadRpc<T>> child_;
+  std::function<void(void)> on_destruction_;
+};
+
+template <typename T>
+class AsyncStreamingReadRpcTracking
+    : public google::cloud::internal::AsyncStreamingReadRpc<T> {
+ public:
+  AsyncStreamingReadRpcTracking(
+      std::unique_ptr<google::cloud::internal::AsyncStreamingReadRpc<T>> child,
+      std::function<void(void)> on_destruction)
+      : child_(std::move(child)), on_destruction_(std::move(on_destruction)) {}
+
+  ~AsyncStreamingReadRpcTracking() override { on_destruction_(); }
+
+  void Cancel() override { child_->Cancel(); }
+  future<bool> Start() override { return child_->Start(); }
+  future<std::optional<T>> Read() override { return child_->Read(); }
+  future<Status> Finish() override { return child_->Finish(); }
+  RpcMetadata GetRequestMetadata() const override {
+    return child_->GetRequestMetadata();
+  }
+
+ private:
+  std::unique_ptr<google::cloud::internal::AsyncStreamingReadRpc<T>> child_;
+  std::function<void(void)> on_destruction_;
+};
+
+template <typename Response>
+Response UnaryHelper(std::shared_ptr<DynamicChannelPool<BigtableStub>>& pool,
+                     std::function<Response(BigtableStub&)> fn) {
+  auto child = pool->GetChannelRandomTwoLeastUsed();
+  auto stub = child->AcquireStub();
+  auto result = fn(*stub);
+  child->ReleaseStub();
+  return result;
+}
+
+template <typename Response>
+std::unique_ptr<google::cloud::internal::StreamingReadRpc<Response>>
+StreamingHelper(
+    std::shared_ptr<DynamicChannelPool<BigtableStub>>& pool,
+    std::function<std::unique_ptr<
+        google::cloud::internal::StreamingReadRpc<Response>>(BigtableStub&)>
+        fn) {
+  auto child = pool->GetChannelRandomTwoLeastUsed();
+  auto stub = child->AcquireStub();
+  auto result = fn(*stub);
+  auto release_fn = [weak = child->MakeWeak()] {
+    auto child = weak.lock();
+    if (child) child->ReleaseStub();
+  };
+  return std::make_unique<StreamingReadRpcTracking<Response>>(
+      std::move(result), std::move(release_fn));
+}
+
+template <typename Response>
+std::unique_ptr<google::cloud::internal::AsyncStreamingReadRpc<Response>>
+AsyncStreamingHelper(
+    std::shared_ptr<DynamicChannelPool<BigtableStub>>& pool,
+    std::function<std::unique_ptr<
+        google::cloud::internal::AsyncStreamingReadRpc<Response>>(
+        BigtableStub&)>
+        fn) {
+  auto child = pool->GetChannelRandomTwoLeastUsed();
+  auto stub = child->AcquireStub();
+  auto result = fn(*stub);
+  auto release_fn = [weak = child->MakeWeak()] {
+    auto child = weak.lock();
+    if (child) child->ReleaseStub();
+  };
+  return std::make_unique<AsyncStreamingReadRpcTracking<Response>>(
+      std::move(result), std::move(release_fn));
+}
+
+}  // namespace
+
+std::unique_ptr<google::cloud::internal::StreamingReadRpc<
+    google::bigtable::v2::ReadRowsResponse>>
+BigtableRandomTwoLeastUsed::ReadRows(
+    std::shared_ptr<grpc::ClientContext> context, Options const& options,
+    google::bigtable::v2::ReadRowsRequest const& request) {
+  return StreamingHelper<google::bigtable::v2::ReadRowsResponse>(
+      pool_, [&, context = std::move(context)](BigtableStub& stub) mutable {
+        return stub.ReadRows(std::move(context), options, request);
+      });
+}
+
+std::unique_ptr<google::cloud::internal::StreamingReadRpc<
+    google::bigtable::v2::SampleRowKeysResponse>>
+BigtableRandomTwoLeastUsed::SampleRowKeys(
+    std::shared_ptr<grpc::ClientContext> context, Options const& options,
+    google::bigtable::v2::SampleRowKeysRequest const& request) {
+  return StreamingHelper<google::bigtable::v2::SampleRowKeysResponse>(
+      pool_, [&, context = std::move(context)](BigtableStub& stub) mutable {
+        return stub.SampleRowKeys(std::move(context), options, request);
+      });
+}
+
+StatusOr<google::bigtable::v2::MutateRowResponse>
+BigtableRandomTwoLeastUsed::MutateRow(
+    grpc::ClientContext& context, Options const& options,
+    google::bigtable::v2::MutateRowRequest const& request) {
+  return UnaryHelper<StatusOr<google::bigtable::v2::MutateRowResponse>>(
+      pool_, [&](BigtableStub& stub) {
+        return stub.MutateRow(context, options, request);
+      });
+}
+
+std::unique_ptr<google::cloud::internal::StreamingReadRpc<
+    google::bigtable::v2::MutateRowsResponse>>
+BigtableRandomTwoLeastUsed::MutateRows(
+    std::shared_ptr<grpc::ClientContext> context, Options const& options,
+    google::bigtable::v2::MutateRowsRequest const& request) {
+  return StreamingHelper<google::bigtable::v2::MutateRowsResponse>(
+      pool_, [&, context = std::move(context)](BigtableStub& stub) mutable {
+        return stub.MutateRows(std::move(context), options, request);
+      });
+}
+
+StatusOr<google::bigtable::v2::CheckAndMutateRowResponse>
+BigtableRandomTwoLeastUsed::CheckAndMutateRow(
+    grpc::ClientContext& context, Options const& options,
+    google::bigtable::v2::CheckAndMutateRowRequest const& request) {
+  return UnaryHelper<StatusOr<google::bigtable::v2::CheckAndMutateRowResponse>>(
+      pool_, [&](BigtableStub& stub) {
+        return stub.CheckAndMutateRow(context, options, request);
+      });
+}
+
+StatusOr<google::bigtable::v2::PingAndWarmResponse>
+BigtableRandomTwoLeastUsed::PingAndWarm(
+    grpc::ClientContext& context, Options const& options,
+    google::bigtable::v2::PingAndWarmRequest const& request) {
+  return UnaryHelper<StatusOr<google::bigtable::v2::PingAndWarmResponse>>(
+      pool_, [&](BigtableStub& stub) {
+        return stub.PingAndWarm(context, options, request);
+      });
+}
+
+StatusOr<google::bigtable::v2::ReadModifyWriteRowResponse>
+BigtableRandomTwoLeastUsed::ReadModifyWriteRow(
+    grpc::ClientContext& context, Options const& options,
+    google::bigtable::v2::ReadModifyWriteRowRequest const& request) {
+  return UnaryHelper<
+      StatusOr<google::bigtable::v2::ReadModifyWriteRowResponse>>(
+      pool_, [&](BigtableStub& stub) {
+        return stub.ReadModifyWriteRow(context, options, request);
+      });
+}
+
+StatusOr<google::bigtable::v2::PrepareQueryResponse>
+BigtableRandomTwoLeastUsed::PrepareQuery(
+    grpc::ClientContext& context, Options const& options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  return UnaryHelper<StatusOr<google::bigtable::v2::PrepareQueryResponse>>(
+      pool_, [&](BigtableStub& stub) {
+        return stub.PrepareQuery(context, options, request);
+      });
+}
+
+std::unique_ptr<google::cloud::internal::StreamingReadRpc<
+    google::bigtable::v2::ExecuteQueryResponse>>
+BigtableRandomTwoLeastUsed::ExecuteQuery(
+    std::shared_ptr<grpc::ClientContext> context, Options const& options,
+    google::bigtable::v2::ExecuteQueryRequest const& request) {
+  return StreamingHelper<google::bigtable::v2::ExecuteQueryResponse>(
+      pool_, [&, context = std::move(context)](BigtableStub& stub) mutable {
+        return stub.ExecuteQuery(std::move(context), options, request);
+      });
+}
+
+std::unique_ptr<google::cloud::internal::AsyncStreamingReadRpc<
+    google::bigtable::v2::ReadRowsResponse>>
+BigtableRandomTwoLeastUsed::AsyncReadRows(
+    google::cloud::CompletionQueue const& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::ReadRowsRequest const& request) {
+  return AsyncStreamingHelper<google::bigtable::v2::ReadRowsResponse>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncReadRows(cq, std::move(context), std::move(options),
+                                  request);
+      });
+}
+
+std::unique_ptr<google::cloud::internal::AsyncStreamingReadRpc<
+    google::bigtable::v2::SampleRowKeysResponse>>
+BigtableRandomTwoLeastUsed::AsyncSampleRowKeys(
+    google::cloud::CompletionQueue const& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::SampleRowKeysRequest const& request) {
+  return AsyncStreamingHelper<google::bigtable::v2::SampleRowKeysResponse>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncSampleRowKeys(cq, std::move(context),
+                                       std::move(options), request);
+      });
+}
+
+future<StatusOr<google::bigtable::v2::MutateRowResponse>>
+BigtableRandomTwoLeastUsed::AsyncMutateRow(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::MutateRowRequest const& request) {
+  return UnaryHelper<future<StatusOr<google::bigtable::v2::MutateRowResponse>>>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncMutateRow(cq, std::move(context), std::move(options),
+                                   request);
+      });
+}
+
+std::unique_ptr<google::cloud::internal::AsyncStreamingReadRpc<
+    google::bigtable::v2::MutateRowsResponse>>
+BigtableRandomTwoLeastUsed::AsyncMutateRows(
+    google::cloud::CompletionQueue const& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::MutateRowsRequest const& request) {
+  return AsyncStreamingHelper<google::bigtable::v2::MutateRowsResponse>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncMutateRows(cq, std::move(context), std::move(options),
+                                    request);
+      });
+}
+
+future<StatusOr<google::bigtable::v2::CheckAndMutateRowResponse>>
+BigtableRandomTwoLeastUsed::AsyncCheckAndMutateRow(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::CheckAndMutateRowRequest const& request) {
+  return UnaryHelper<
+      future<StatusOr<google::bigtable::v2::CheckAndMutateRowResponse>>>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncCheckAndMutateRow(cq, std::move(context),
+                                           std::move(options), request);
+      });
+}
+
+future<StatusOr<google::bigtable::v2::PingAndWarmResponse>>
+BigtableRandomTwoLeastUsed::AsyncPingAndWarm(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PingAndWarmRequest const& request) {
+  return UnaryHelper<
+      future<StatusOr<google::bigtable::v2::PingAndWarmResponse>>>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncPingAndWarm(cq, std::move(context), std::move(options),
+                                     request);
+      });
+}
+
+future<StatusOr<google::bigtable::v2::ReadModifyWriteRowResponse>>
+BigtableRandomTwoLeastUsed::AsyncReadModifyWriteRow(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::ReadModifyWriteRowRequest const& request) {
+  return UnaryHelper<
+      future<StatusOr<google::bigtable::v2::ReadModifyWriteRowResponse>>>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncReadModifyWriteRow(cq, std::move(context),
+                                            std::move(options), request);
+      });
+}
+
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+BigtableRandomTwoLeastUsed::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  return UnaryHelper<
+      future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>>(
+      pool_, [&, context = std::move(context),
+              options = std::move(options)](BigtableStub& stub) mutable {
+        return stub.AsyncPrepareQuery(cq, std::move(context),
+                                      std::move(options), request);
+      });
+}
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_internal
+}  // namespace cloud
+}  // namespace google