From dc260f196ae01eaebfbaaad766372a2adf23bb9d Mon Sep 17 00:00:00 2001 From: chen quan Date: Sat, 7 Jan 2023 13:32:56 +0800 Subject: [PATCH] refactor: simplify the code (#2763) * refactor: simplify the code * fix: fix data race * refactor: simplify the code * refactor: simplify the code --- core/logx/writer.go | 5 ++- rest/handler/loghandler.go | 60 +++++++-------------------------- rest/handler/loghandler_test.go | 34 +++++-------------- 3 files changed, 24 insertions(+), 75 deletions(-) diff --git a/core/logx/writer.go b/core/logx/writer.go index 8a7774b2..1117e155 100644 --- a/core/logx/writer.go +++ b/core/logx/writer.go @@ -269,7 +269,10 @@ func combineGlobalFields(fields []LogField) []LogField { return fields } - return append(globals.([]LogField), fields...) + originLogFields := globals.([]LogField) + logFields := append(make([]LogField, 0, len(originLogFields)+len(fields)), originLogFields...) + + return append(logFields, fields...) } func output(writer io.Writer, level string, val interface{}, fields ...LogField) { diff --git a/rest/handler/loghandler.go b/rest/handler/loghandler.go index d3d5cddc..68552c4c 100644 --- a/rest/handler/loghandler.go +++ b/rest/handler/loghandler.go @@ -22,6 +22,7 @@ import ( "github.com/zeromicro/go-zero/core/utils" "github.com/zeromicro/go-zero/rest/httpx" "github.com/zeromicro/go-zero/rest/internal" + "github.com/zeromicro/go-zero/rest/internal/response" ) const ( @@ -31,66 +32,30 @@ const ( var slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold) -type loggedResponseWriter struct { - w http.ResponseWriter - r *http.Request - code int -} - -func (w *loggedResponseWriter) Flush() { - if flusher, ok := w.w.(http.Flusher); ok { - flusher.Flush() - } -} - -func (w *loggedResponseWriter) Header() http.Header { - return w.w.Header() -} - -// Hijack implements the http.Hijacker interface. -// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it. -func (w *loggedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { - if hijacked, ok := w.w.(http.Hijacker); ok { - return hijacked.Hijack() - } - - return nil, nil, errors.New("server doesn't support hijacking") -} - -func (w *loggedResponseWriter) Write(bytes []byte) (int, error) { - return w.w.Write(bytes) -} - -func (w *loggedResponseWriter) WriteHeader(code int) { - w.w.WriteHeader(code) - w.code = code -} - // LogHandler returns a middleware that logs http request and response. func LogHandler(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { timer := utils.NewElapsedTimer() logs := new(internal.LogCollector) - lrw := loggedResponseWriter{ - w: w, - r: r, - code: http.StatusOK, + lrw := response.WithCodeResponseWriter{ + Writer: w, + Code: http.StatusOK, } var dup io.ReadCloser r.Body, dup = iox.DupReadCloser(r.Body) next.ServeHTTP(&lrw, r.WithContext(context.WithValue(r.Context(), internal.LogContext, logs))) r.Body = dup - logBrief(r, lrw.code, timer, logs) + logBrief(r, lrw.Code, timer, logs) }) } type detailLoggedResponseWriter struct { - writer *loggedResponseWriter + writer *response.WithCodeResponseWriter buf *bytes.Buffer } -func newDetailLoggedResponseWriter(writer *loggedResponseWriter, buf *bytes.Buffer) *detailLoggedResponseWriter { +func newDetailLoggedResponseWriter(writer *response.WithCodeResponseWriter, buf *bytes.Buffer) *detailLoggedResponseWriter { return &detailLoggedResponseWriter{ writer: writer, buf: buf, @@ -108,7 +73,7 @@ func (w *detailLoggedResponseWriter) Header() http.Header { // Hijack implements the http.Hijacker interface. // This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it. func (w *detailLoggedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { - if hijacked, ok := w.writer.w.(http.Hijacker); ok { + if hijacked, ok := w.writer.Writer.(http.Hijacker); ok { return hijacked.Hijack() } @@ -129,10 +94,9 @@ func DetailedLogHandler(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { timer := utils.NewElapsedTimer() var buf bytes.Buffer - lrw := newDetailLoggedResponseWriter(&loggedResponseWriter{ - w: w, - r: r, - code: http.StatusOK, + lrw := newDetailLoggedResponseWriter(&response.WithCodeResponseWriter{ + Writer: w, + Code: http.StatusOK, }, &buf) var dup io.ReadCloser @@ -203,7 +167,7 @@ func logDetails(r *http.Request, response *detailLoggedResponseWriter, timer *ut logs *internal.LogCollector) { var buf bytes.Buffer duration := timer.Duration() - code := response.writer.code + code := response.writer.Code logger := logx.WithContext(r.Context()) buf.WriteString(fmt.Sprintf("[HTTP] %s - %d - %s - %s\n=> %s\n", r.Method, code, r.RemoteAddr, timex.ReprOfDuration(duration), dumpRequest(r))) diff --git a/rest/handler/loghandler_test.go b/rest/handler/loghandler_test.go index a84b43d0..65cd0b88 100644 --- a/rest/handler/loghandler_test.go +++ b/rest/handler/loghandler_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/zeromicro/go-zero/rest/internal" + "github.com/zeromicro/go-zero/rest/internal/response" ) func init() { @@ -54,7 +55,7 @@ func TestLogHandlerVeryLong(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "http://localhost", &buf) handler := LogHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { r.Context().Value(internal.LogContext).(*internal.LogCollector).Append("anything") - io.Copy(io.Discard, r.Body) + _, _ = io.Copy(io.Discard, r.Body) w.Header().Set("X-Test", "test") w.WriteHeader(http.StatusServiceUnavailable) _, err := w.Write([]byte("content")) @@ -89,45 +90,26 @@ func TestLogHandlerSlow(t *testing.T) { assert.Equal(t, http.StatusOK, resp.Code) } } - -func TestLogHandler_Hijack(t *testing.T) { - resp := httptest.NewRecorder() - writer := &loggedResponseWriter{ - w: resp, - } - assert.NotPanics(t, func() { - writer.Hijack() - }) - - writer = &loggedResponseWriter{ - w: mockedHijackable{resp}, - } - assert.NotPanics(t, func() { - writer.Hijack() - }) -} - func TestDetailedLogHandler_Hijack(t *testing.T) { resp := httptest.NewRecorder() writer := &detailLoggedResponseWriter{ - writer: &loggedResponseWriter{ - w: resp, + writer: &response.WithCodeResponseWriter{ + Writer: resp, }, } assert.NotPanics(t, func() { - writer.Hijack() + _, _, _ = writer.Hijack() }) writer = &detailLoggedResponseWriter{ - writer: &loggedResponseWriter{ - w: mockedHijackable{resp}, + writer: &response.WithCodeResponseWriter{ + Writer: mockedHijackable{resp}, }, } assert.NotPanics(t, func() { - writer.Hijack() + _, _, _ = writer.Hijack() }) } - func TestSetSlowThreshold(t *testing.T) { assert.Equal(t, defaultSlowThreshold, slowThreshold.Load()) SetSlowThreshold(time.Second)