refactor: simplify the code (#2763)

* refactor: simplify the code

* fix: fix data race

* refactor: simplify the code

* refactor: simplify the code
This commit is contained in:
chen quan
2023-01-07 13:32:56 +08:00
committed by GitHub
parent 559726112c
commit dc260f196a
3 changed files with 24 additions and 75 deletions

View File

@@ -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)))