prome

pkg/common/prometheus/prometheus.go

@@ -2,6 +2,7 @@ package prometheus
 
 import (
 	"Open_IM/pkg/common/config"
+	"bytes"
 	"net/http"
 	"strconv"
 
@@ -26,6 +27,28 @@ func PrometheusHandler() gin.HandlerFunc {
 	}
 }
 
+type responseBodyWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+func (r responseBodyWriter) Write(b []byte) (int, error) {
+	r.body.Write(b)
+	return r.ResponseWriter.Write(b)
+}
+
+func PromeTheusMiddleware(c *gin.Context) {
+	PromeInc(ApiRequestCounter)
+	w := &responseBodyWriter{body: &bytes.Buffer{}, ResponseWriter: c.Writer}
+	c.Writer = w
+	c.Next()
+	if c.Writer.Status() == http.StatusOK {
+		PromeInc(ApiRequestSuccessCounter)
+	} else {
+		PromeInc(ApiRequestFailedCounter)
+	}
+}
+
 func PromeInc(counter prometheus.Counter) {
 	if config.Config.Prometheus.Enable {
 		counter.Inc()