Merge pull request #16 from weedge/main

fix: print span info with ctx  && add  CtxKVLog method

logging/zap/logger.go

@@ -17,6 +17,7 @@ package zap
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 
 	"github.com/cloudwego/kitex/pkg/klog"
@@ -43,7 +44,7 @@ var (
 )
 
 type Logger struct {
-	l      *zap.SugaredLogger
+	*zap.SugaredLogger
 	config *config
 }
 
@@ -60,30 +61,31 @@ func NewLogger(opts ...Option) *Logger {
 		config.zapOpts...)
 
 	return &Logger{
-		l:      logger.Sugar(),
-		config: config,
+		SugaredLogger: logger.Sugar(),
+		config:        config,
 	}
 }
 
 func (l *Logger) Log(level klog.Level, kvs ...interface{}) {
+	logger := l.With()
 	switch level {
 	case klog.LevelTrace, klog.LevelDebug:
-		l.l.Debug(kvs...)
+		logger.Debug(kvs...)
 	case klog.LevelInfo:
-		l.l.Info(kvs...)
+		logger.Info(kvs...)
 	case klog.LevelNotice, klog.LevelWarn:
-		l.l.Warn(kvs...)
+		logger.Warn(kvs...)
 	case klog.LevelError:
-		l.l.Error(kvs...)
+		logger.Error(kvs...)
 	case klog.LevelFatal:
-		l.l.Fatal(kvs...)
+		logger.Fatal(kvs...)
 	default:
-		l.l.Warn(kvs...)
+		logger.Warn(kvs...)
 	}
 }
 
 func (l *Logger) Logf(level klog.Level, format string, kvs ...interface{}) {
-	logger := l.l.With()
+	logger := l.With()
 	switch level {
 	case klog.LevelTrace, klog.LevelDebug:
 		logger.Debugf(format, kvs...)
@@ -102,10 +104,25 @@ func (l *Logger) Logf(level klog.Level, format string, kvs ...interface{}) {
 
 func (l *Logger) CtxLogf(level klog.Level, ctx context.Context, format string, kvs ...interface{}) {
 	var zlevel zapcore.Level
+	var sl *zap.SugaredLogger
+
 	span := trace.SpanFromContext(ctx)
+	var traceKVs []interface{}
+	if span.SpanContext().TraceID().IsValid() {
+		traceKVs = append(traceKVs, traceIDKey, span.SpanContext().TraceID())
+	}
+	if span.SpanContext().SpanID().IsValid() {
+		traceKVs = append(traceKVs, spanIDKey, span.SpanContext().SpanID())
+	}
+	if span.SpanContext().TraceFlags().IsSampled() {
+		traceKVs = append(traceKVs, traceFlagsKey, span.SpanContext().TraceFlags())
+	}
+	if len(traceKVs) > 0 {
+		sl = l.With(traceKVs...)
+	} else {
+		sl = l.With()
+	}
 
-	sl := l.l.With(
-		traceIDKey, span.SpanContext().TraceID(), spanIDKey, span.SpanContext().SpanID(), traceFlagsKey, span.SpanContext().TraceFlags())
 	switch level {
 	case klog.LevelDebug, klog.LevelTrace:
 		zlevel = zap.DebugLevel
@@ -127,12 +144,12 @@ func (l *Logger) CtxLogf(level klog.Level, ctx context.Context, format string, k
 		sl.Warnf(format, kvs...)
 	}
 
-	msg := getMessage(format, kvs)
-
 	if !span.IsRecording() {
-		l.Logf(level, format, kvs...)
 		return
 	}
+
+	msg := getMessage(format, kvs)
+
 	attrs := []attribute.KeyValue{
 		logMessageKey.String(msg),
 		logSeverityTextKey.String(OtelSeverityText(zlevel)),
@@ -254,11 +271,67 @@ func (l *Logger) SetOutput(writer io.Writer) {
 	log := zap.New(
 		zapcore.NewCore(l.config.coreConfig.enc, ws, l.config.coreConfig.lvl),
 		l.config.zapOpts...,
-	).Sugar()
+	)
 	l.config.coreConfig.ws = ws
-	l.l = log
+	l.SugaredLogger = log.Sugar()
 }
 
-func (l *Logger) Sync() {
-	_ = l.l.Sync()
+func (l *Logger) CtxKVLog(ctx context.Context, level klog.Level, format string, kvs ...interface{}) {
+	if len(kvs) == 0 || len(kvs)%2 != 0 {
+		l.Warn(fmt.Sprint("Keyvalues must appear in pairs:", kvs))
+		return
+	}
+
+	span := trace.SpanFromContext(ctx)
+	if span.SpanContext().TraceID().IsValid() {
+		kvs = append(kvs, traceIDKey, span.SpanContext().TraceID())
+	}
+	if span.SpanContext().SpanID().IsValid() {
+		kvs = append(kvs, spanIDKey, span.SpanContext().SpanID())
+	}
+	if span.SpanContext().TraceFlags().IsSampled() {
+		kvs = append(kvs, traceFlagsKey, span.SpanContext().TraceFlags())
+	}
+
+	var zlevel zapcore.Level
+	zl := l.With()
+	switch level {
+	case klog.LevelDebug, klog.LevelTrace:
+		zlevel = zap.DebugLevel
+		zl.Debugw(format, kvs...)
+	case klog.LevelInfo:
+		zlevel = zap.InfoLevel
+		zl.Infow(format, kvs...)
+	case klog.LevelNotice, klog.LevelWarn:
+		zlevel = zap.WarnLevel
+		zl.Warnw(format, kvs...)
+	case klog.LevelError:
+		zlevel = zap.ErrorLevel
+		zl.Errorw(format, kvs...)
+	case klog.LevelFatal:
+		zlevel = zap.FatalLevel
+		zl.Fatalw(format, kvs...)
+	default:
+		zlevel = zap.WarnLevel
+		zl.Warnw(format, kvs...)
+	}
+
+	if !span.IsRecording() {
+		return
+	}
+
+	msg := getMessage(format, kvs)
+	attrs := []attribute.KeyValue{
+		logMessageKey.String(msg),
+		logSeverityTextKey.String(OtelSeverityText(zlevel)),
+	}
+
+	// notice: AddEvent,SetStatus,RecordError all have check span.IsRecording
+	span.AddEvent(logEventKey, trace.WithAttributes(attrs...))
+
+	// set span status
+	if zlevel <= l.config.traceConfig.errorSpanLevel {
+		span.SetStatus(codes.Error, msg)
+		span.RecordError(errors.New(msg), trace.WithStackTrace(l.config.traceConfig.recordStackTraceInSpan))
+	}
 }

logging/zap/logger_test.go

@@ -88,7 +88,12 @@ func TestLogger(t *testing.T) {
 		WithTraceErrorSpanLevel(zap.WarnLevel),
 		WithRecordStackTraceInSpan(true),
 	)
-	defer logger.Sync()
+	defer func() {
+		err := logger.Sync()
+		if err != nil {
+			return
+		}
+	}()
 
 	klog.SetLogger(logger)
 	klog.SetOutput(buf)
@@ -138,7 +143,12 @@ func TestLogLevel(t *testing.T) {
 		WithTraceErrorSpanLevel(zap.WarnLevel),
 		WithRecordStackTraceInSpan(true),
 	)
-	defer logger.Sync()
+	defer func() {
+		err := logger.Sync()
+		if err != nil {
+			return
+		}
+	}()
 
 	// output to buffer
 	logger.SetOutput(buf)
@@ -161,7 +171,12 @@ func TestCoreOption(t *testing.T) {
 		WithCoreLevel(zap.NewAtomicLevelAt(zapcore.WarnLevel)),
 		WithCoreWs(zapcore.AddSync(buf)),
 	)
-	defer logger.Sync()
+	defer func() {
+		err := logger.Sync()
+		if err != nil {
+			return
+		}
+	}()
 
 	logger.SetOutput(buf)
 
@@ -183,7 +198,12 @@ func TestZapOption(t *testing.T) {
 	logger := NewLogger(
 		WithZapOptions(zap.AddCaller()),
 	)
-	defer logger.Sync()
+	defer func() {
+		err := logger.Sync()
+		if err != nil {
+			return
+		}
+	}()
 
 	logger.SetOutput(buf)
 
@@ -194,3 +214,62 @@ func TestZapOption(t *testing.T) {
 	// test caller in log result
 	assert.True(t, strings.Contains(buf.String(), "caller"))
 }
+
+// TestCtxLogger test kv logger work with ctx
+func TestCtxKVLogger(t *testing.T) {
+	ctx := context.Background()
+
+	buf := new(bytes.Buffer)
+
+	shutdown := stdoutProvider(ctx)
+	defer shutdown()
+
+	logger := NewLogger(
+		WithTraceErrorSpanLevel(zap.WarnLevel),
+		WithRecordStackTraceInSpan(true),
+	)
+	defer func() {
+		err := logger.Sync()
+		if err != nil {
+			return
+		}
+	}()
+
+	klog.SetLogger(logger)
+	klog.SetOutput(buf)
+	klog.SetLevel(klog.LevelTrace)
+
+	for _, level := range []klog.Level{
+		klog.LevelTrace,
+		klog.LevelDebug,
+		klog.LevelInfo,
+		klog.LevelNotice,
+		klog.LevelWarn,
+		klog.LevelError,
+		// klog.LevelFatal,
+	} {
+		logger.CtxLogf(level, context.Background(), "log from origin zap %s=%s", "k1", "v1")
+		println(buf.String())
+		assert.True(t, strings.Contains(buf.String(), "log from origin zap"))
+		assert.True(t, strings.Contains(buf.String(), "k1"))
+		assert.True(t, strings.Contains(buf.String(), "v1"))
+		buf.Reset()
+	}
+
+	for _, level := range []klog.Level{
+		klog.LevelTrace,
+		klog.LevelDebug,
+		klog.LevelInfo,
+		klog.LevelNotice,
+		klog.LevelWarn,
+		klog.LevelError,
+		// klog.LevelFatal,
+	} {
+		logger.CtxKVLog(context.Background(), level, "log from origin zap", "k1", "v1")
+		println(buf.String())
+		assert.True(t, strings.Contains(buf.String(), "log from origin zap"))
+		assert.True(t, strings.Contains(buf.String(), "k1"))
+		assert.True(t, strings.Contains(buf.String(), "v1"))
+		buf.Reset()
+	}
+}