add http handler response interceptor

Author: leopoldxx

middleware/recover.go

@@ -1,36 +1,30 @@
 package middleware
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/leopoldxx/go-utils/trace"
 )
 
-// Recover middleware will response a default 500 error
-// when panic occurred inside of the handler process.
-// NOTIC: if there is no a trace handler out the recover handler,
-// prepend a default trace handler
-func Recover() Middleware {
-	return func(next http.HandlerFunc) http.HandlerFunc {
-		return func(w http.ResponseWriter, r *http.Request) {
-			tracer := trace.GetTraceFromRequest(r)
-			defer func() {
-				if err := recover(); err != nil {
-					tracer.Error("panic:", tracer.Stack())
-					http.Error(w, "internal error, plz check log!", http.StatusInternalServerError)
-				}
-			}()
-			next(w, r)
-		}
-	}
-}
-
 // RecoverWithTrace middleware is a RecoverMiddleware wraps with a trace handler
 func RecoverWithTrace(name string) Middleware {
 	return func(next http.HandlerFunc) http.HandlerFunc {
 		return func(w http.ResponseWriter, r *http.Request) {
+			var rw *responseWriter
+			if defaultResponseInterceptor != nil {
+				rw = &responseWriter{
+					ResponseWriter: w,
+					status:         http.StatusOK,
+				}
+			}
 			recoverHandler := func(w http.ResponseWriter, r *http.Request) {
 				tracer := trace.GetTraceFromRequest(r)
+				if rw, ok := w.(interface {
+					Record(ctx context.Context, recorder Recorder)
+				}); ok {
+					defer rw.Record(r.Context(), defaultResponseInterceptor)
+				}
 				defer func() {
 					if err := recover(); err != nil {
 						tracer.Error("panic:", tracer.Stack())
@@ -39,6 +33,10 @@ func RecoverWithTrace(name string) Middleware {
 				}()
 				next(w, r)
 			}
+			if rw != nil {
+				trace.HandleFunc(name, recoverHandler)(rw, r)
+				return
+			}
 
 			trace.HandleFunc(name, recoverHandler)(w, r)
 		}

middleware/recover_test.go

@@ -18,11 +18,12 @@ func TestRecover(t *testing.T) {
 		panic("+_+!")
 	}
 
-	tmw := Trace("recover_test")
-	rmw := Recover()
+	rmw := RecoverWithTrace("rwt")
 
-	newhh := Chain(tmw, rmw).HandlerFunc(ph)
-	newhh2 := Chain(tmw, rmw).HandlerFunc(ph2)
+	SetDefaultResponseInterceptor(NewMultiRecorder(NewLogRecorder(), NewLogRecorder(), NewLogRecorder()))
+
+	newhh := Chain(rmw).HandlerFunc(ph)
+	newhh2 := Chain(rmw).HandlerFunc(ph2)
 
 	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
 	w := httptest.NewRecorder()
@@ -32,6 +33,6 @@ func TestRecover(t *testing.T) {
 	if w.Code != 500 {
 		t.Fatal("fail chain handler:", w)
 	}
-	t.Log(string(w.Body.Bytes()))
+	t.Log(len(w.Body.Bytes()), string(w.Body.Bytes()))
 
 }

middleware/response_interceptor.go

@@ -0,0 +1,97 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"sync"
+
+	"github.com/leopoldxx/go-utils/trace"
+)
+
+var defaultResponseInterceptor Recorder
+
+// SetDefaultResponseInterceptor for http handler response
+func SetDefaultResponseInterceptor(r Recorder) {
+	defaultResponseInterceptor = r
+}
+
+type responseWriter struct {
+	http.ResponseWriter
+	sync.Mutex
+
+	status int
+	size   int
+}
+
+func (rs *responseWriter) Header() http.Header {
+	return rs.ResponseWriter.Header()
+}
+
+func (rs *responseWriter) Write(data []byte) (int, error) {
+	rs.Lock()
+	rs.size += len(data)
+	rs.Unlock()
+	return rs.ResponseWriter.Write(data)
+}
+
+func (rs *responseWriter) WriteHeader(status int) {
+	rs.Lock()
+	rs.status = status
+	rs.Unlock()
+	rs.ResponseWriter.WriteHeader(status)
+}
+
+// Recorder for http handler response status & body size
+type Recorder interface {
+	Record(ctx context.Context, statistics Statistics)
+}
+
+// NewLogRecorder for log purpose
+func NewLogRecorder() Recorder {
+	return &logRecorder{}
+}
+
+type logRecorder struct{}
+
+func (lr logRecorder) Record(ctx context.Context, statistics Statistics) {
+	tracer := trace.GetTraceFromContext(ctx)
+	tracer.Infof("%+v", statistics)
+}
+
+// NewMultiRecorder will chain MultiRecorder
+func NewMultiRecorder(recorders ...Recorder) Recorder {
+	return &multiRecorder{recorders: recorders}
+}
+
+type multiRecorder struct {
+	recorders []Recorder
+}
+
+func (mr multiRecorder) Record(ctx context.Context, statistics Statistics) {
+	var wg sync.WaitGroup
+	for i := range mr.recorders {
+		wg.Add(1)
+		go func(r Recorder, ctx context.Context, statistics Statistics) {
+			defer wg.Done()
+			r.Record(ctx, statistics)
+		}(mr.recorders[i], ctx, statistics)
+	}
+	wg.Wait()
+}
+
+// Statistics for http handler response
+type Statistics struct {
+	Status   int
+	BodySize int
+}
+
+func (rs *responseWriter) Record(ctx context.Context, recorder Recorder) {
+	var s Statistics
+	rs.Lock()
+	s.Status = rs.status
+	s.BodySize = rs.size
+	rs.Unlock()
+	if recorder != nil {
+		recorder.Record(ctx, s)
+	}
+}

server/example/main.go

@@ -12,6 +12,7 @@ import (
 )
 
 func main() {
+	middleware.SetDefaultResponseInterceptor(middleware.NewLogRecorder())
 	s := server.New(server.ListenAddr(":8001"), server.APIPrefix("/example"))
 	s.Register(new(filesvr))
 	s.ListenAndServe()