Count senddata bytes in prometheus

internal/helper/countingresponsewriter.go

+package helper
+
+import (
+	"net/http"
+)
+
+type CountingResponseWriter interface {
+	http.ResponseWriter
+	Count() int64
+}
+
+type countingResponseWriter struct {
+	rw     http.ResponseWriter
+	status int
+	count  int64
+}
+
+func NewCountingResponseWriter(rw http.ResponseWriter) CountingResponseWriter {
+	return &countingResponseWriter{rw: rw}
+}
+
+func (c *countingResponseWriter) Header() http.Header {
+	return c.rw.Header()
+}
+
+func (c *countingResponseWriter) Write(data []byte) (n int, err error) {
+	if c.status == 0 {
+		c.WriteHeader(http.StatusOK)
+	}
+
+	n, err = c.rw.Write(data)
+	c.count += int64(n)
+	return n, err
+}
+
+func (c *countingResponseWriter) WriteHeader(status int) {
+	if c.status != 0 {
+		return
+	}
+
+	c.status = status
+	c.rw.WriteHeader(status)
+}
+
+func (c *countingResponseWriter) Count() int64 {
+	return c.count
+}

internal/senddata/senddata.go

@@ -16,10 +16,18 @@ var (
 		},
 		[]string{"name"},
 	)
+	sendDataResponseBytes = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "gitlab_workhorse_senddata_response_bytes",
+			Help: "How many bytes have been written by workhorse senddata response hijackers",
+		},
+		[]string{"name"},
+	)
 )
 
 func init() {
 	prometheus.MustRegister(sendDataResponses)
+	prometheus.MustRegister(sendDataResponseBytes)
 }
 
 type sendDataResponseWriter struct {
@@ -82,7 +90,9 @@ func (s *sendDataResponseWriter) tryInject() bool {
 			s.hijacked = true
 			sendDataResponses.WithLabelValues(injecter.Name()).Inc()
 			helper.DisableResponseBuffering(s.rw)
-			injecter.Inject(s.rw, s.req, header)
+			crw := helper.NewCountingResponseWriter(s.rw)
+			injecter.Inject(crw, s.req, header)
+			sendDataResponseBytes.WithLabelValues(injecter.Name()).Add(float64(crw.Count()))
 			return true
 		}
 	}