Commit 625b53a1 authored by Jacob Vosmaer's avatar Jacob Vosmaer

Count senddata bytes in prometheus

parent a29c725a
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
}
...@@ -16,10 +16,18 @@ var ( ...@@ -16,10 +16,18 @@ var (
}, },
[]string{"name"}, []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() { func init() {
prometheus.MustRegister(sendDataResponses) prometheus.MustRegister(sendDataResponses)
prometheus.MustRegister(sendDataResponseBytes)
} }
type sendDataResponseWriter struct { type sendDataResponseWriter struct {
...@@ -82,7 +90,9 @@ func (s *sendDataResponseWriter) tryInject() bool { ...@@ -82,7 +90,9 @@ func (s *sendDataResponseWriter) tryInject() bool {
s.hijacked = true s.hijacked = true
sendDataResponses.WithLabelValues(injecter.Name()).Inc() sendDataResponses.WithLabelValues(injecter.Name()).Inc()
helper.DisableResponseBuffering(s.rw) 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 return true
} }
} }
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment