From c16dae8b4d76c55037a56d9382d359fa91dc27f6 Mon Sep 17 00:00:00 2001
From: Dmitrii Stepanov <d.stepanov@yadro.com>
Date: Fri, 18 Oct 2024 13:31:53 +0300
Subject: [PATCH] [#1437] logger: Use context to log trace id

Signed-off-by: Dmitrii Stepanov
---
 pkg/util/logger/log.go         | 40 ++++++++++++++++++++++++++++++++++
 pkg/util/logger/logger.go      | 16 +++++++++++---
 pkg/util/logger/test/logger.go | 11 +++++-----
 3 files changed, 59 insertions(+), 8 deletions(-)
 create mode 100644 pkg/util/logger/log.go

diff --git a/pkg/util/logger/log.go b/pkg/util/logger/log.go
new file mode 100644
index 000000000..269e07d90
--- /dev/null
+++ b/pkg/util/logger/log.go
@@ -0,0 +1,40 @@
+package logger
+
+import (
+	"context"
+
+	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/tracing"
+	"go.uber.org/zap"
+)
+
+func (l *Logger) Debug(ctx context.Context, msg string, fields ...zap.Field) {
+	if traceID := tracing.GetTraceID(ctx); traceID != "" {
+		l.z.Debug(msg, append(fields, zap.String("trace_id", traceID))...)
+		return
+	}
+	l.z.Debug(msg, fields...)
+}
+
+func (l *Logger) Info(ctx context.Context, msg string, fields ...zap.Field) {
+	if traceID := tracing.GetTraceID(ctx); traceID != "" {
+		l.z.Info(msg, append(fields, zap.String("trace_id", traceID))...)
+		return
+	}
+	l.z.Info(msg, fields...)
+}
+
+func (l *Logger) Warn(ctx context.Context, msg string, fields ...zap.Field) {
+	if traceID := tracing.GetTraceID(ctx); traceID != "" {
+		l.z.Warn(msg, append(fields, zap.String("trace_id", traceID))...)
+		return
+	}
+	l.z.Warn(msg, fields...)
+}
+
+func (l *Logger) Error(ctx context.Context, msg string, fields ...zap.Field) {
+	if traceID := tracing.GetTraceID(ctx); traceID != "" {
+		l.z.Error(msg, append(fields, zap.String("trace_id", traceID))...)
+		return
+	}
+	l.z.Error(msg, fields...)
+}
diff --git a/pkg/util/logger/logger.go b/pkg/util/logger/logger.go
index 4b60f02de..d4ac2ab02 100644
--- a/pkg/util/logger/logger.go
+++ b/pkg/util/logger/logger.go
@@ -12,7 +12,7 @@ import (
 // Logger represents a component
 // for writing messages to log.
 type Logger struct {
-	*zap.Logger
+	z   *zap.Logger
 	lvl zap.AtomicLevel
 }
 
@@ -136,7 +136,7 @@ func newConsoleLogger(prm *Prm) (*Logger, error) {
 		return nil, err
 	}
 
-	l := &Logger{Logger: lZap, lvl: lvl}
+	l := &Logger{z: lZap, lvl: lvl}
 	prm._log = l
 
 	return l, nil
@@ -169,7 +169,7 @@ func newJournaldLogger(prm *Prm) (*Logger, error) {
 
 	lZap := zap.New(coreWithContext, zap.AddStacktrace(zap.NewAtomicLevelAt(zap.FatalLevel)))
 
-	l := &Logger{Logger: lZap, lvl: lvl}
+	l := &Logger{z: lZap, lvl: lvl}
 	prm._log = l
 
 	return l, nil
@@ -179,3 +179,13 @@ func (l *Logger) reload(prm Prm) error {
 	l.lvl.SetLevel(prm.level)
 	return nil
 }
+
+func (l *Logger) WithOptions(options ...zap.Option) {
+	l.z = l.z.WithOptions(options...)
+}
+
+func NewLoggerWrapper(z *zap.Logger) *Logger {
+	return &Logger{
+		z: z,
+	}
+}
diff --git a/pkg/util/logger/test/logger.go b/pkg/util/logger/test/logger.go
index f93756d17..b5b0a31eb 100644
--- a/pkg/util/logger/test/logger.go
+++ b/pkg/util/logger/test/logger.go
@@ -11,9 +11,10 @@ import (
 
 // NewLogger creates a new logger.
 func NewLogger(t testing.TB) *logger.Logger {
-	var l logger.Logger
-	l.Logger = zaptest.NewLogger(t,
-		zaptest.Level(zapcore.DebugLevel),
-		zaptest.WrapOptions(zap.Development(), zap.AddCaller()))
-	return &l
+	return logger.NewLoggerWrapper(
+		zaptest.NewLogger(t,
+			zaptest.Level(zapcore.DebugLevel),
+			zaptest.WrapOptions(zap.Development(), zap.AddCaller()),
+		),
+	)
 }