From ac090fe3265a78948a4c4bb9c53ef7f31f355e4d Mon Sep 17 00:00:00 2001
From: nolouch <nolouch@gmail.com>
Date: Tue, 15 Jan 2019 17:42:24 +0800
Subject: [PATCH] embed: add zap logger builder

Signed-off-by: nolouch <nolouch@gmail.com>
---
 embed/config.go         |  3 ++
 embed/config_logging.go | 87 +++++++++++++++++++++++++++--------------
 2 files changed, 60 insertions(+), 30 deletions(-)

diff --git a/embed/config.go b/embed/config.go
index 10cab6e66..d9317b956 100644
--- a/embed/config.go
+++ b/embed/config.go
@@ -301,6 +301,9 @@ type Config struct {
 	// Debug is true, to enable debug level logging.
 	Debug bool `json:"debug"`
 
+	// ZapLoggerBuilder is used to build the zap logger.
+	ZapLoggerBuilder func(*Config) error
+
 	// logger logs server-side operations. The default is nil,
 	// and "setupLogging" must be called before starting server.
 	// Do not set logger directly.
diff --git a/embed/config_logging.go b/embed/config_logging.go
index 2a9ce77f8..fc6406697 100644
--- a/embed/config_logging.go
+++ b/embed/config_logging.go
@@ -186,28 +186,29 @@ func (cfg *Config) setupLogging() error {
 				lcfg.Level = zap.NewAtomicLevelAt(zap.DebugLevel)
 				grpc.EnableTracing = true
 			}
-
-			var err error
-			cfg.logger, err = lcfg.Build()
-			if err != nil {
-				return err
-			}
-
-			cfg.loggerConfig = &lcfg
-			cfg.loggerCore = nil
-			cfg.loggerWriteSyncer = nil
-
-			grpcLogOnce.Do(func() {
-				// debug true, enable info, warning, error
-				// debug false, only discard info
-				var gl grpclog.LoggerV2
-				gl, err = logutil.NewGRPCLoggerV2(lcfg)
-				if err == nil {
-					grpclog.SetLoggerV2(gl)
+			if cfg.ZapLoggerBuilder == nil {
+				cfg.ZapLoggerBuilder = func(c *Config) error {
+					var err error
+					c.logger, err = lcfg.Build()
+					if err != nil {
+						return err
+					}
+					c.loggerMu.Lock()
+					defer c.loggerMu.Unlock()
+					c.loggerConfig = &lcfg
+					c.loggerCore = nil
+					c.loggerWriteSyncer = nil
+					grpcLogOnce.Do(func() {
+						// debug true, enable info, warning, error
+						// debug false, only discard info
+						var gl grpclog.LoggerV2
+						gl, err = logutil.NewGRPCLoggerV2(lcfg)
+						if err == nil {
+							grpclog.SetLoggerV2(gl)
+						}
+					})
+					return nil
 				}
-			})
-			if err != nil {
-				return err
 			}
 		} else {
 			if len(cfg.LogOutputs) > 1 {
@@ -237,17 +238,26 @@ func (cfg *Config) setupLogging() error {
 				syncer,
 				lvl,
 			)
-			cfg.logger = zap.New(cr, zap.AddCaller(), zap.ErrorOutput(syncer))
+			if cfg.ZapLoggerBuilder == nil {
+				cfg.ZapLoggerBuilder = func(c *Config) error {
+					c.logger = zap.New(cr, zap.AddCaller(), zap.ErrorOutput(syncer))
+					c.loggerMu.Lock()
+					defer c.loggerMu.Unlock()
+					c.loggerConfig = nil
+					c.loggerCore = cr
+					c.loggerWriteSyncer = syncer
 
-			cfg.loggerConfig = nil
-			cfg.loggerCore = cr
-			cfg.loggerWriteSyncer = syncer
-
-			grpcLogOnce.Do(func() {
-				grpclog.SetLoggerV2(logutil.NewGRPCLoggerV2FromZapCore(cr, syncer))
-			})
+					grpcLogOnce.Do(func() {
+						grpclog.SetLoggerV2(logutil.NewGRPCLoggerV2FromZapCore(cr, syncer))
+					})
+					return nil
+				}
+			}
+		}
+		err := cfg.ZapLoggerBuilder(cfg)
+		if err != nil {
+			return err
 		}
-
 		logTLSHandshakeFailure := func(conn *tls.Conn, err error) {
 			state := conn.ConnectionState()
 			remoteAddr := conn.RemoteAddr().String()
@@ -284,3 +294,20 @@ func (cfg *Config) setupLogging() error {
 
 	return nil
 }
+
+// NewZapCoreLoggerBuilder generates a zap core logger builder.
+func NewZapCoreLoggerBuilder(lg *zap.Logger, cr zapcore.Core, syncer zapcore.WriteSyncer) func(*Config) error {
+	return func(cfg *Config) error {
+		cfg.loggerMu.Lock()
+		defer cfg.loggerMu.Unlock()
+		cfg.logger = lg
+		cfg.loggerConfig = nil
+		cfg.loggerCore = cr
+		cfg.loggerWriteSyncer = syncer
+
+		grpcLogOnce.Do(func() {
+			grpclog.SetLoggerV2(logutil.NewGRPCLoggerV2FromZapCore(cr, syncer))
+		})
+		return nil
+	}
+}