From ffa1dca7926a2f2d431242547526b0233ec2236f Mon Sep 17 00:00:00 2001
From: kvii <56432636+kvii@users.noreply.github.com>
Date: Fri, 15 Mar 2024 11:15:31 +0800
Subject: [PATCH] fix: logging middleware caller issue #3133 (#3134)

* fix: logging middleware caller issue #3133

* fix: fix lint issue
---
 middleware/logging/logging.go      |  4 +--
 middleware/logging/logging_test.go | 56 ++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/middleware/logging/logging.go b/middleware/logging/logging.go
index 9573d1985..8595750bd 100644
--- a/middleware/logging/logging.go
+++ b/middleware/logging/logging.go
@@ -37,7 +37,7 @@ func Server(logger log.Logger) middleware.Middleware {
 				reason = se.Reason
 			}
 			level, stack := extractError(err)
-			_ = log.WithContext(ctx, logger).Log(level,
+			log.NewHelper(log.WithContext(ctx, logger)).Log(level,
 				"kind", "server",
 				"component", kind,
 				"operation", operation,
@@ -73,7 +73,7 @@ func Client(logger log.Logger) middleware.Middleware {
 				reason = se.Reason
 			}
 			level, stack := extractError(err)
-			_ = log.WithContext(ctx, logger).Log(level,
+			log.NewHelper(log.WithContext(ctx, logger)).Log(level,
 				"kind", "client",
 				"component", kind,
 				"operation", operation,
diff --git a/middleware/logging/logging_test.go b/middleware/logging/logging_test.go
index 76cad19b1..1e9495188 100644
--- a/middleware/logging/logging_test.go
+++ b/middleware/logging/logging_test.go
@@ -177,3 +177,59 @@ func TestExtractError(t *testing.T) {
 		})
 	}
 }
+
+type extractKeyValues [][]any
+
+func (l *extractKeyValues) Log(_ log.Level, kv ...any) error { *l = append(*l, kv); return nil }
+
+func TestServer_CallerPath(t *testing.T) {
+	var a extractKeyValues
+	logger := log.With(&a, "caller", log.Caller(5)) // report where the helper was called
+
+	// make sure the caller is same
+	sameCaller := func(fn middleware.Handler) { _, _ = fn(context.Background(), nil) }
+
+	// caller: [... log inside middleware, fn(context.Background(), nil)]
+	h := func(context.Context, any) (a any, e error) { return }
+	h = Server(logger)(h)
+	sameCaller(h)
+
+	// caller: [... helper.Info("foo"), fn(context.Background(), nil)]
+	helper := log.NewHelper(logger)
+	sameCaller(func(context.Context, any) (a any, e error) { helper.Info("foo"); return })
+
+	t.Log(a[0])
+	t.Log(a[1])
+	if a[0][0] != "caller" || a[1][0] != "caller" {
+		t.Fatal("caller not found")
+	}
+	if a[0][1] != a[1][1] {
+		t.Fatalf("middleware should have the same caller as log.Helper. middleware: %s, helper: %s", a[0][1], a[1][1])
+	}
+}
+
+func TestClient_CallerPath(t *testing.T) {
+	var a extractKeyValues
+	logger := log.With(&a, "caller", log.Caller(5)) // report where the helper was called
+
+	// make sure the caller is same
+	sameCaller := func(fn middleware.Handler) { _, _ = fn(context.Background(), nil) }
+
+	// caller: [... log inside middleware, fn(context.Background(), nil)]
+	h := func(context.Context, any) (a any, e error) { return }
+	h = Client(logger)(h)
+	sameCaller(h)
+
+	// caller: [... helper.Info("foo"), fn(context.Background(), nil)]
+	helper := log.NewHelper(logger)
+	sameCaller(func(context.Context, any) (a any, e error) { helper.Info("foo"); return })
+
+	t.Log(a[0])
+	t.Log(a[1])
+	if a[0][0] != "caller" || a[1][0] != "caller" {
+		t.Fatal("caller not found")
+	}
+	if a[0][1] != a[1][1] {
+		t.Fatalf("middleware should have the same caller as log.Helper. middleware: %s, helper: %s", a[0][1], a[1][1])
+	}
+}