[Bigtable] Allow the connection_pool_size to be configurable

If left unset, the default should be far larger than the default 4 the
client library uses.

PiperOrigin-RevId: 203139678

3 files changed, 36 insertions, 4 deletions

diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc b/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc
index 8a7309e870..94c0dec812 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc
@@ -31,6 +31,17 @@ class BigtableClientOp : public OpKernel {
                 errors::InvalidArgument("project_id must be non-empty"));
     OP_REQUIRES(ctx, !instance_id_.empty(),
                 errors::InvalidArgument("instance_id must be non-empty"));
+
+    OP_REQUIRES_OK(
+        ctx, ctx->GetAttr("connection_pool_size", &connection_pool_size_));
+    // If left unset by the client code, set it to a default of 100. Note: the
+    // cloud-cpp default of 4 concurrent connections is far too low for high
+    // performance streaming.
+    if (connection_pool_size_ == -1) {
+      connection_pool_size_ = 100;
+    }
+    OP_REQUIRES(ctx, connection_pool_size_ > 0,
+                errors::InvalidArgument("connection_pool_size must be > 0"));
   }
 
   ~BigtableClientOp() override {
@@ -56,10 +67,10 @@ class BigtableClientOp : public OpKernel {
               cinfo_.container(), cinfo_.name(), &resource,
               [this, ctx](
                   BigtableClientResource** ret) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+                auto client_options = google::cloud::bigtable::ClientOptions();
                 std::shared_ptr<google::cloud::bigtable::DataClient> client =
                     google::cloud::bigtable::CreateDefaultDataClient(
-                        project_id_, instance_id_,
-                        google::cloud::bigtable::ClientOptions());
+                        project_id_, instance_id_, std::move(client_options));
                 *ret = new BigtableClientResource(project_id_, instance_id_,
                                                   std::move(client));
                 return Status::OK();
@@ -75,6 +86,7 @@ class BigtableClientOp : public OpKernel {
  private:
   string project_id_;
   string instance_id_;
+  int64 connection_pool_size_;
 
   mutex mu_;
   ContainerInfo cinfo_ GUARDED_BY(mu_);
diff --git a/tensorflow/contrib/bigtable/ops/bigtable_ops.cc b/tensorflow/contrib/bigtable/ops/bigtable_ops.cc
index 17ecc3dcb2..c7ff012ec8 100644
--- a/tensorflow/contrib/bigtable/ops/bigtable_ops.cc
+++ b/tensorflow/contrib/bigtable/ops/bigtable_ops.cc
@@ -22,6 +22,7 @@ namespace tensorflow {
 REGISTER_OP("BigtableClient")
     .Attr("project_id: string")
     .Attr("instance_id: string")
+    .Attr("connection_pool_size: int")
     .Attr("container: string = ''")
     .Attr("shared_name: string = ''")
     .Output("client: resource")
diff --git a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
index a54e020ed7..39c58ba665 100644
--- a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
+++ b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
@@ -49,16 +49,35 @@ class BigtableClient(object):
   `table` method to open a Bigtable Table.
   """
 
-  def __init__(self, project_id, instance_id):
+  def __init__(self, project_id, instance_id, connection_pool_size=None):
     """Creates a BigtableClient that can be used to open connections to tables.
 
     Args:
       project_id: A string representing the GCP project id to connect to.
       instance_id: A string representing the Bigtable instance to connect to.
+      connection_pool_size: (Optional.) A number representing the number of
+        concurrent connections to the Cloud Bigtable service to make.
+
+    Raises:
+      ValueError: if the arguments are invalid (e.g. wrong type, or out of
+        expected ranges (e.g. negative).)
     """
+    if not isinstance(project_id, str):
+      raise ValueError("`project_id` must be a string")
     self._project_id = project_id
+
+    if not isinstance(instance_id, str):
+      raise ValueError("`instance_id` must be a string")
     self._instance_id = instance_id
-    self._resource = gen_bigtable_ops.bigtable_client(project_id, instance_id)
+
+    if connection_pool_size is None:
+      connection_pool_size = -1
+    elif connection_pool_size < 1:
+      raise ValueError("`connection_pool_size` must be positive")
+    self._connection_pool_size = connection_pool_size
+
+    self._resource = gen_bigtable_ops.bigtable_client(project_id, instance_id,
+                                                      connection_pool_size)
 
   def table(self, name, snapshot=None):
     """Opens a table and returns a `BigTable` object.