From ebc90720eaf05c14567e25e6d2ac2a5cd8bcd44b Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Mon, 1 Nov 2021 14:48:26 +0800
Subject: [PATCH] feat: slow threshold customizable in rest (#1189)

* feat: slow threshold customizable in rest

* feat: slow threshold customizable in rest
---
 rest/handler/loghandler.go      | 16 ++++++++++++----
 rest/handler/loghandler_test.go |  8 +++++++-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/rest/handler/loghandler.go b/rest/handler/loghandler.go
index fd8adb2f..57d147b8 100644
--- a/rest/handler/loghandler.go
+++ b/rest/handler/loghandler.go
@@ -16,6 +16,7 @@ import (
 
 	"github.com/tal-tech/go-zero/core/iox"
 	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/syncx"
 	"github.com/tal-tech/go-zero/core/timex"
 	"github.com/tal-tech/go-zero/core/utils"
 	"github.com/tal-tech/go-zero/rest/httpx"
@@ -23,10 +24,12 @@ import (
 )
 
 const (
-	limitBodyBytes = 1024
-	slowThreshold  = time.Millisecond * 500
+	limitBodyBytes       = 1024
+	defaultSlowThreshold = time.Millisecond * 500
 )
 
+var slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
+
 type loggedResponseWriter struct {
 	w    http.ResponseWriter
 	r    *http.Request
@@ -140,6 +143,11 @@ func DetailedLogHandler(next http.Handler) http.Handler {
 	})
 }
 
+// SetSlowThreshold sets the slow threshold.
+func SetSlowThreshold(threshold time.Duration) {
+	slowThreshold.Set(threshold)
+}
+
 func dumpRequest(r *http.Request) string {
 	reqContent, err := httputil.DumpRequest(r, true)
 	if err != nil {
@@ -155,7 +163,7 @@ func logBrief(r *http.Request, code int, timer *utils.ElapsedTimer, logs *intern
 	logger := logx.WithContext(r.Context())
 	buf.WriteString(fmt.Sprintf("[HTTP] %s - %d - %s - %s - %s - %s",
 		r.Method, code, r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent(), timex.ReprOfDuration(duration)))
-	if duration > slowThreshold {
+	if duration > slowThreshold.Load() {
 		logger.Slowf("[HTTP] %s - %d - %s - %s - %s - slowcall(%s)",
 			r.Method, code, r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent(), timex.ReprOfDuration(duration))
 	}
@@ -191,7 +199,7 @@ func logDetails(r *http.Request, response *detailLoggedResponseWriter, timer *ut
 	logger := logx.WithContext(r.Context())
 	buf.WriteString(fmt.Sprintf("[HTTP] %s - %d - %s - %s\n=> %s\n",
 		r.Method, response.writer.code, r.RemoteAddr, timex.ReprOfDuration(duration), dumpRequest(r)))
-	if duration > slowThreshold {
+	if duration > defaultSlowThreshold {
 		logger.Slowf("[HTTP] %s - %d - %s - slowcall(%s)\n=> %s\n",
 			r.Method, response.writer.code, r.RemoteAddr, timex.ReprOfDuration(duration), dumpRequest(r))
 	}
diff --git a/rest/handler/loghandler_test.go b/rest/handler/loghandler_test.go
index f94576fc..b5b75eaf 100644
--- a/rest/handler/loghandler_test.go
+++ b/rest/handler/loghandler_test.go
@@ -53,7 +53,7 @@ func TestLogHandlerSlow(t *testing.T) {
 	for _, logHandler := range handlers {
 		req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
 		handler := logHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			time.Sleep(slowThreshold + time.Millisecond*50)
+			time.Sleep(defaultSlowThreshold + time.Millisecond*50)
 		}))
 
 		resp := httptest.NewRecorder()
@@ -100,6 +100,12 @@ func TestDetailedLogHandler_Hijack(t *testing.T) {
 	})
 }
 
+func TestSetSlowThreshold(t *testing.T) {
+	assert.Equal(t, defaultSlowThreshold, slowThreshold.Load())
+	SetSlowThreshold(time.Second)
+	assert.Equal(t, time.Second, slowThreshold.Load())
+}
+
 func BenchmarkLogHandler(b *testing.B) {
 	b.ReportAllocs()