From ebb00f754e619a800966b30fbd3223acc33a4af5 Mon Sep 17 00:00:00 2001
From: mohanli-ml <mohanli@google.com>
Date: Fri, 30 Oct 2020 06:04:14 +0000
Subject: [PATCH] feat(bigtable): add an experimental attempt DirectPath flag

---
 bigtable/bigtable.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/bigtable/bigtable.go b/bigtable/bigtable.go
index 594c0d2c0340..e1d47a0c8454 100644
--- a/bigtable/bigtable.go
+++ b/bigtable/bigtable.go
@@ -24,12 +24,14 @@ import (
 	"net/url"
 	"strconv"
 	"time"
+  "flag"
 
 	btopt "cloud.google.com/go/bigtable/internal/option"
 	"cloud.google.com/go/internal/trace"
 	"github.com/golang/protobuf/proto"
 	gax "github.com/googleapis/gax-go/v2"
 	"google.golang.org/api/option"
+  "google.golang.org/api/option/internaloption"
 	gtransport "google.golang.org/api/transport/grpc"
 	btpb "google.golang.org/genproto/googleapis/bigtable/v2"
 	"google.golang.org/grpc"
@@ -38,7 +40,17 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-const prodAddr = "bigtable.googleapis.com:443"
+const (
+  prodAddr = "bigtable.googleapis.com:443"
+)
+
+// Add an experimental flag to control the attempt to directpath.
+// Once e2e tests are finished, this flag will be removed, and client
+// will attempt by default
+var attemptDirectPath bool
+func init() {
+  flag.BoolVar(&attemptDirectPath, "it.attempt-directpath", false, "flag to attempt directpath")
+}
 
 // Client is a client for reading and writing data to tables in an instance.
 //
@@ -79,7 +91,8 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
 		option.WithGRPCDialOption(grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(1<<28), grpc.MaxCallRecvMsgSize(1<<28))),
 		// TODO(grpc/grpc-go#1388) using connection pool without WithBlock
 		// can cause RPCs to fail randomly. We can delete this after the issue is fixed.
-		option.WithGRPCDialOption(grpc.WithBlock()))
+		option.WithGRPCDialOption(grpc.WithBlock()),
+    internaloption.EnableDirectPath(attemptDirectPath))
 	o = append(o, opts...)
 	connPool, err := gtransport.DialPool(ctx, o...)
 	if err != nil {