From 87620aa8014d667030a788d108a2b993badcec74 Mon Sep 17 00:00:00 2001
From: Hossein Zolfi <hossein.zolfi@gmail.com>
Date: Mon, 22 Apr 2024 00:31:49 +0330
Subject: [PATCH] Enhance Logging: Include TraceID and SpanID only for
 Recording Spans (#20)

Logging non-recording spans is uninformative for troubleshooting and absent in tracing platforms like Jaeger.

Co-authored-by: Samuel Berthe <dev@samuel-berthe.fr>
---
 middleware.go | 38 ++++++++++++++++++++++++++++----------
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/middleware.go b/middleware.go
index 1747c4a..6d86ac4 100644
--- a/middleware.go
+++ b/middleware.go
@@ -1,12 +1,11 @@
 package sloggin
 
 import (
+	"log/slog"
 	"net/http"
 	"strings"
 	"time"
 
-	"log/slog"
-
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 	"go.opentelemetry.io/otel/trace"
@@ -169,14 +168,7 @@ func NewWithConfig(logger *slog.Logger, config Config) gin.HandlerFunc {
 		}
 
 		// otel
-		if config.WithTraceID {
-			traceID := trace.SpanFromContext(c.Request.Context()).SpanContext().TraceID().String()
-			baseAttributes = append(baseAttributes, slog.String(TraceIDKey, traceID))
-		}
-		if config.WithSpanID {
-			spanID := trace.SpanFromContext(c.Request.Context()).SpanContext().SpanID().String()
-			baseAttributes = append(baseAttributes, slog.String(SpanIDKey, spanID))
-		}
+		baseAttributes = extractTraceSpanID(c, config, baseAttributes)
 
 		// request body
 		requestAttributes = append(requestAttributes, slog.Int("length", br.bytes))
@@ -290,3 +282,29 @@ func AddCustomAttributes(c *gin.Context, attr slog.Attr) {
 		c.Set(customAttributesCtxKey, append(attrs, attr))
 	}
 }
+
+func extractTraceSpanID(c *gin.Context, config Config, baseAttributes []slog.Attr) []slog.Attr {
+	if !(config.WithTraceID || config.WithSpanID) {
+		return baseAttributes
+	}
+
+	ctx := c.Request.Context()
+	span := trace.SpanFromContext(ctx)
+	if !span.IsRecording() {
+		return baseAttributes
+	}
+
+	spanCtx := span.SpanContext()
+
+	if config.WithTraceID && spanCtx.HasTraceID() {
+		traceID := trace.SpanFromContext(ctx).SpanContext().TraceID().String()
+		baseAttributes = append(baseAttributes, slog.String(TraceIDKey, traceID))
+	}
+
+	if config.WithSpanID && spanCtx.HasSpanID() {
+		spanID := spanCtx.SpanID().String()
+		baseAttributes = append(baseAttributes, slog.String(SpanIDKey, spanID))
+	}
+
+	return baseAttributes
+}