From 1a5b7349bc06de23f75051d6f9dc7b6154e29ac6 Mon Sep 17 00:00:00 2001
From: mohanli-ml <mohanli@google.com>
Date: Fri, 30 Oct 2020 06:04:16 +0000
Subject: [PATCH] feat(option): add AttemptDirectPath option

---
 internal/settings.go   |  1 +
 option/option.go       | 12 ++++++++++++
 transport/grpc/dial.go | 20 +++++---------------
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/internal/settings.go b/internal/settings.go
index 26259b82abb..bb473eabf04 100644
--- a/internal/settings.go
+++ b/internal/settings.go
@@ -41,6 +41,7 @@ type DialSettings struct {
 	CustomClaims        map[string]interface{}
 	SkipValidation      bool
 	ImpersonationConfig *impersonate.Config
+	AttemptDirectPath   bool
 
 	// Google API system parameters. For more information please read:
 	// https://cloud.google.com/apis/docs/system-parameters
diff --git a/option/option.go b/option/option.go
index 686476f9cbb..7ae6e99ceb0 100644
--- a/option/option.go
+++ b/option/option.go
@@ -324,3 +324,15 @@ func (i impersonateServiceAccount) Apply(o *internal.DialSettings) {
 	o.ImpersonationConfig.Delegates = make([]string, len(i.delegates))
 	copy(o.ImpersonationConfig.Delegates, i.delegates)
 }
+
+// WithAttemptDirectPath returns a ClientOption that overrides the default 
+// attempt to use DirectPath
+func WithAttemptDirectPath(dp bool) ClientOption {
+  return withAttemptDirectPath(dp)
+}
+
+type withAttemptDirectPath bool
+
+func (w withAttemptDirectPath) Apply(o *internal.DialSettings) {
+  o.AttemptDirectPath = bool(w)
+}
diff --git a/transport/grpc/dial.go b/transport/grpc/dial.go
index 55c04a5af06..7153cf40207 100644
--- a/transport/grpc/dial.go
+++ b/transport/grpc/dial.go
@@ -12,7 +12,6 @@ import (
 	"crypto/tls"
 	"errors"
 	"log"
-	"os"
 	"strings"
 
 	"go.opencensus.io/plugin/ocgrpc"
@@ -135,12 +134,11 @@ func dial(ctx context.Context, insecure bool, o *internal.DialSettings) (*grpc.C
 		}
 
 		// Attempt Direct Path only if:
+    // * Yoshi library set AttemptDirectPath option
 		// * The endpoint is a host:port (or dns:///host:port).
 		// * Credentials are obtained via GCE metadata server, using the default
 		//   service account.
-		// * Opted in via GOOGLE_CLOUD_ENABLE_DIRECT_PATH environment variable.
-		//   For example, GOOGLE_CLOUD_ENABLE_DIRECT_PATH=spanner,pubsub
-		if isDirectPathEnabled(endpoint) && isTokenSourceDirectPathCompatible(creds.TokenSource) {
+		if o.AttemptDirectPath && checkDirectPathEndPoint(endpoint) && isTokenSourceDirectPathCompatible(creds.TokenSource) {
 			if !strings.HasPrefix(endpoint, "dns:///") {
 				endpoint = "dns:///" + endpoint
 			}
@@ -189,7 +187,7 @@ func dial(ctx context.Context, insecure bool, o *internal.DialSettings) (*grpc.C
 	// point when isDirectPathEnabled will default to true, we guard it by
 	// the Directpath env var for now once we can introspect user defined
 	// dialer (https://github.com/grpc/grpc-go/issues/2795).
-	if timeoutDialerOption != nil && isDirectPathEnabled(endpoint) {
+	if timeoutDialerOption != nil && o.AttemptDirectPath && checkDirectPathEndPoint(endpoint) {
 		grpcOpts = append(grpcOpts, timeoutDialerOption)
 	}
 
@@ -250,7 +248,7 @@ func isTokenSourceDirectPathCompatible(ts oauth2.TokenSource) bool {
 	return true
 }
 
-func isDirectPathEnabled(endpoint string) bool {
+func checkDirectPathEndPoint(endpoint string) bool {
 	// Only host:port is supported, not other schemes (e.g., "tcp://" or "unix://").
 	// Also don't try direct path if the user has chosen an alternate name resolver
 	// (i.e., via ":///" prefix).
@@ -261,15 +259,7 @@ func isDirectPathEnabled(endpoint string) bool {
 		return false
 	}
 
-	// Only try direct path if the user has opted in via the environment variable.
-	directPathAPIs := strings.Split(os.Getenv("GOOGLE_CLOUD_ENABLE_DIRECT_PATH"), ",")
-	for _, api := range directPathAPIs {
-		// Ignore empty string since an empty env variable splits into [""]
-		if api != "" && strings.Contains(endpoint, api) {
-			return true
-		}
-	}
-	return false
+	return true 
 }
 
 func processAndValidateOpts(opts []option.ClientOption) (*internal.DialSettings, error) {