From 1189ee3f3dd90ffce8c4abfd52557ca156cbdd62 Mon Sep 17 00:00:00 2001
From: Marek Siarkowicz <siarkowicz@google.com>
Date: Wed, 12 May 2021 14:55:22 +0200
Subject: [PATCH] client: Add logger argument to NewCtxClient

---
 client/v3/client.go                        | 23 +++++++++++++++++++---
 client/v3/client_test.go                   | 11 +++++++++++
 server/etcdserver/api/v3client/v3client.go |  6 +-----
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/client/v3/client.go b/client/v3/client.go
index 2c3f3fede..75e3c97e2 100644
--- a/client/v3/client.go
+++ b/client/v3/client.go
@@ -83,11 +83,21 @@ func New(cfg Config) (*Client, error) {
 // NewCtxClient creates a client with a context but no underlying grpc
 // connection. This is useful for embedded cases that override the
 // service interface implementations and do not need connection management.
-func NewCtxClient(ctx context.Context) *Client {
+func NewCtxClient(ctx context.Context, opts ...Option) *Client {
 	cctx, cancel := context.WithCancel(ctx)
-	return &Client{ctx: cctx, cancel: cancel, lgMu: new(sync.RWMutex), lg: zap.NewNop()}
+	c := &Client{ctx: cctx, cancel: cancel, lgMu: new(sync.RWMutex)}
+	for _, opt := range opts {
+		opt(c)
+	}
+	if c.lg == nil {
+		c.lg = zap.NewNop()
+	}
+	return c
 }
 
+// Option is a function type that can be passed as argument to NewCtxClient to configure client
+type Option func(*Client)
+
 // NewFromURL creates a new etcdv3 client from a URL.
 func NewFromURL(url string) (*Client, error) {
 	return New(Config{Endpoints: []string{url}})
@@ -98,9 +108,16 @@ func NewFromURLs(urls []string) (*Client, error) {
 	return New(Config{Endpoints: urls})
 }
 
+// WithZapLogger is a NewCtxClient option that overrides the logger
+func WithZapLogger(lg *zap.Logger) Option {
+	return func(c *Client) {
+		c.lg = lg
+	}
+}
+
 // WithLogger overrides the logger.
 //
-// Deprecated: Please use Logger field in clientv3.Config
+// Deprecated: Please use WithZapLogger or Logger field in clientv3.Config
 //
 // Does not changes grpcLogger, that can be explicitly configured
 // using grpc_zap.ReplaceGrpcLoggerV2(..) method.
diff --git a/client/v3/client_test.go b/client/v3/client_test.go
index a15dbbff7..b2ff7ef17 100644
--- a/client/v3/client_test.go
+++ b/client/v3/client_test.go
@@ -17,6 +17,7 @@ package clientv3
 import (
 	"context"
 	"fmt"
+	"go.uber.org/zap"
 	"net"
 	"testing"
 	"time"
@@ -187,3 +188,13 @@ func TestWithLogger(t *testing.T) {
 		t.Errorf("WithLogger should modify *zap.Logger")
 	}
 }
+
+func TestZapWithLogger(t *testing.T) {
+	ctx := context.Background()
+	lg := zap.NewNop()
+	c := NewCtxClient(ctx, WithZapLogger(lg))
+
+	if c.lg != lg {
+		t.Errorf("WithZapLogger should modify *zap.Logger")
+	}
+}
diff --git a/server/etcdserver/api/v3client/v3client.go b/server/etcdserver/api/v3client/v3client.go
index ff4d91b5c..8342dc434 100644
--- a/server/etcdserver/api/v3client/v3client.go
+++ b/server/etcdserver/api/v3client/v3client.go
@@ -28,11 +28,7 @@ import (
 // of making gRPC calls through sockets, the client makes direct function calls
 // to the etcd server through its api/v3rpc function interfaces.
 func New(s *etcdserver.EtcdServer) *clientv3.Client {
-	c := clientv3.NewCtxClient(context.Background())
-	lg := s.Logger()
-	if lg != nil {
-		c.WithLogger(lg)
-	}
+	c := clientv3.NewCtxClient(context.Background(), clientv3.WithZapLogger(s.Logger()))
 
 	kvc := adapter.KvServerToKvClient(v3rpc.NewQuotaKVServer(s))
 	c.KV = clientv3.NewKVFromKVClient(kvc, c)