From 29a996e70ef5596b1bfc29111f2789117b756783 Mon Sep 17 00:00:00 2001
From: John Guo <john@johng.cn>
Date: Sat, 18 Dec 2021 21:54:12 +0800
Subject: [PATCH] add Trace-Id response header for package ghttp

---
 net/ghttp/ghttp.go                          | 27 ++++++------
 net/ghttp/ghttp_response.go                 |  2 +
 net/ghttp/ghttp_server_handler.go           |  4 +-
 net/ghttp/ghttp_z_unit_feature_otel_test.go | 47 +++++++++++++++++++++
 4 files changed, 66 insertions(+), 14 deletions(-)
 create mode 100644 net/ghttp/ghttp_z_unit_feature_otel_test.go

diff --git a/net/ghttp/ghttp.go b/net/ghttp/ghttp.go
index ce21dcd4b..0dd27cda4 100644
--- a/net/ghttp/ghttp.go
+++ b/net/ghttp/ghttp.go
@@ -119,19 +119,20 @@ const (
 )
 
 const (
-	supportedHttpMethods = "GET,PUT,POST,DELETE,PATCH,HEAD,CONNECT,OPTIONS,TRACE"
-	defaultMethod        = "ALL"
-	exceptionExit        = "exit"
-	exceptionExitAll     = "exit_all"
-	exceptionExitHook    = "exit_hook"
-	routeCacheDuration   = time.Hour
-	methodNameInit       = "Init"
-	methodNameShut       = "Shut"
-	ctxKeyForRequest     = "gHttpRequestObject"
-	contentTypeXml       = "text/xml"
-	contentTypeHtml      = "text/html"
-	contentTypeJson      = "application/json"
-	swaggerUIPackedPath  = "/goframe/swaggerui"
+	supportedHttpMethods  = "GET,PUT,POST,DELETE,PATCH,HEAD,CONNECT,OPTIONS,TRACE"
+	defaultMethod         = "ALL"
+	exceptionExit         = "exit"
+	exceptionExitAll      = "exit_all"
+	exceptionExitHook     = "exit_hook"
+	routeCacheDuration    = time.Hour
+	methodNameInit        = "Init"
+	methodNameShut        = "Shut"
+	ctxKeyForRequest      = "gHttpRequestObject"
+	contentTypeXml        = "text/xml"
+	contentTypeHtml       = "text/html"
+	contentTypeJson       = "application/json"
+	swaggerUIPackedPath   = "/goframe/swaggerui"
+	responseTraceIdHeader = "Trace-Id"
 )
 
 var (
diff --git a/net/ghttp/ghttp_response.go b/net/ghttp/ghttp_response.go
index 533a40472..d1f0d041c 100644
--- a/net/ghttp/ghttp_response.go
+++ b/net/ghttp/ghttp_response.go
@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"net/http"
 
+	"github.com/gogf/gf/v2/net/gtrace"
 	"github.com/gogf/gf/v2/os/gfile"
 	"github.com/gogf/gf/v2/os/gres"
 )
@@ -142,6 +143,7 @@ func (r *Response) ClearBuffer() {
 
 // Flush outputs the buffer content to the client and clears the buffer.
 func (r *Response) Flush() {
+	r.Header().Set(responseTraceIdHeader, gtrace.GetTraceID(r.Request.Context()))
 	if r.Server.config.ServerAgent != "" {
 		r.Header().Set("Server", r.Server.config.ServerAgent)
 	}
diff --git a/net/ghttp/ghttp_server_handler.go b/net/ghttp/ghttp_server_handler.go
index b8e8e2359..0bf94b7a4 100644
--- a/net/ghttp/ghttp_server_handler.go
+++ b/net/ghttp/ghttp_server_handler.go
@@ -67,7 +67,9 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		s.handleAccessLog(request)
 		// Close the session, which automatically update the TTL
 		// of the session if it exists.
-		request.Session.Close()
+		if err := request.Session.Close(); err != nil {
+			intlog.Error(request.Context(), err)
+		}
 
 		// Close the request and response body
 		// to release the file descriptor in time.
diff --git a/net/ghttp/ghttp_z_unit_feature_otel_test.go b/net/ghttp/ghttp_z_unit_feature_otel_test.go
new file mode 100644
index 000000000..e7fb400ec
--- /dev/null
+++ b/net/ghttp/ghttp_z_unit_feature_otel_test.go
@@ -0,0 +1,47 @@
+// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
+//
+// This Source Code Form is subject to the terms of the MIT License.
+// If a copy of the MIT was not distributed with this file,
+// You can obtain one at https://github.com/gogf/gf.
+
+package ghttp_test
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/gogf/gf/v2/frame/g"
+	"github.com/gogf/gf/v2/net/ghttp"
+	"github.com/gogf/gf/v2/net/gtrace"
+	"github.com/gogf/gf/v2/test/gtest"
+)
+
+func Test_OTEL_TraceID(t *testing.T) {
+	var (
+		traceId string
+	)
+	p, _ := ports.PopRand()
+	s := g.Server(p)
+	s.BindHandler("/", func(r *ghttp.Request) {
+		traceId = gtrace.GetTraceID(r.Context())
+		r.Response.Write(r.GetUrl())
+	})
+	s.SetPort(p)
+	s.SetDumpRouterMap(false)
+	s.Start()
+	defer s.Shutdown()
+
+	time.Sleep(100 * time.Millisecond)
+	gtest.C(t, func(t *gtest.T) {
+		prefix := fmt.Sprintf("http://127.0.0.1:%d", p)
+		client := g.Client()
+		client.SetBrowserMode(true)
+		client.SetPrefix(prefix)
+		res, err := client.Get(ctx, "/")
+		t.AssertNil(err)
+		defer res.Close()
+
+		t.Assert(res.Header.Get("Trace-Id"), traceId)
+	})
+}