Browse Source

clientv3: let user provide a client context through Config

Anthony Romano 8 years ago
parent
commit
c9452c6ad4
2 changed files with 12 additions and 2 deletions
  1. 7 2
      clientv3/client.go
  2. 5 0
      clientv3/config.go

+ 7 - 2
clientv3/client.go

@@ -272,7 +272,7 @@ func (c *Client) dial(endpoint string, dopts ...grpc.DialOption) (*grpc.ClientCo
 			tokenMu: &sync.RWMutex{},
 		}
 
-		err := c.getToken(context.TODO())
+		err := c.getToken(c.ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -307,7 +307,12 @@ func newClient(cfg *Config) (*Client, error) {
 	}
 
 	// use a temporary skeleton client to bootstrap first connection
-	ctx, cancel := context.WithCancel(context.TODO())
+	baseCtx := context.TODO()
+	if cfg.Context != nil {
+		baseCtx = cfg.Context
+	}
+
+	ctx, cancel := context.WithCancel(baseCtx)
 	client := &Client{
 		conn:   nil,
 		cfg:    *cfg,

+ 5 - 0
clientv3/config.go

@@ -18,6 +18,7 @@ import (
 	"crypto/tls"
 	"time"
 
+	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 )
 
@@ -43,4 +44,8 @@ type Config struct {
 
 	// DialOptions is a list of dial options for the grpc client (e.g., for interceptors).
 	DialOptions []grpc.DialOption
+
+	// Context is the default client context; it can be used to cancel grpc dial out and
+	// other operations that do not have an explicit context.
+	Context context.Context
 }