Add traffic size metrics
Co-authored-by: OmarElawady <omarelawady1998@gmail.com> Co-authored-by: Mathieu Lonjaret <mathieu.lonjaret@gmail.com> Co-authored-by: Romain <rtribotte@users.noreply.github.com>
This commit is contained in:
Tom Moulard
GitHub
parent
10528c973a
commit
d578ed7327
26 changed files with 1125 additions and 339 deletions
198
pkg/middlewares/capture/capture.go
Normal file
198
pkg/middlewares/capture/capture.go
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
// Package capture is a middleware that captures requests/responses size, and status.
|
||||
//
|
||||
// For another middleware to get those attributes of a request/response, this middleware
|
||||
// should be added before in the middleware chain.
|
||||
//
|
||||
// handler, _ := NewHandler()
|
||||
// chain := alice.New().
|
||||
// Append(WrapHandler(handler)).
|
||||
// Append(myOtherMiddleware).
|
||||
// then(...)
|
||||
//
|
||||
// As this middleware stores those data in the request's context, the data can
|
||||
// be retrieved at anytime after the ServerHTTP.
|
||||
//
|
||||
// func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http.Handler) {
|
||||
// capt, err := capture.FromContext(req.Context())
|
||||
// if err != nil {
|
||||
// ...
|
||||
// }
|
||||
//
|
||||
// fmt.Println(capt.Status())
|
||||
// fmt.Println(capt.ResponseSize())
|
||||
// fmt.Println(capt.RequestSize())
|
||||
// }
|
||||
package capture
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
|
||||
"github.com/containous/alice"
|
||||
"github.com/traefik/traefik/v2/pkg/middlewares"
|
||||
)
|
||||
|
||||
type key string
|
||||
|
||||
const capturedData key = "capturedData"
|
||||
|
||||
// Handler will store each request data to its context.
|
||||
type Handler struct{}
|
||||
|
||||
// WrapHandler wraps capture handler into an Alice Constructor.
|
||||
func WrapHandler(handler *Handler) alice.Constructor {
|
||||
return func(next http.Handler) (http.Handler, error) {
|
||||
return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
|
||||
handler.ServeHTTP(rw, req, next)
|
||||
}), nil
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http.Handler) {
|
||||
c := Capture{}
|
||||
if req.Body != nil {
|
||||
readCounter := &readCounter{source: req.Body}
|
||||
c.rr = readCounter
|
||||
req.Body = readCounter
|
||||
}
|
||||
responseWriter := newResponseWriter(rw)
|
||||
c.rw = responseWriter
|
||||
ctx := context.WithValue(req.Context(), capturedData, &c)
|
||||
next.ServeHTTP(responseWriter, req.WithContext(ctx))
|
||||
}
|
||||
|
||||
// Capture is the object populated by the capture middleware,
|
||||
// allowing to gather information about the request and response.
|
||||
type Capture struct {
|
||||
rr *readCounter
|
||||
rw responseWriter
|
||||
}
|
||||
|
||||
// FromContext returns the Capture value found in ctx, or an empty Capture otherwise.
|
||||
func FromContext(ctx context.Context) (*Capture, error) {
|
||||
c := ctx.Value(capturedData)
|
||||
if c == nil {
|
||||
return nil, errors.New("value not found")
|
||||
}
|
||||
capt, ok := c.(*Capture)
|
||||
if !ok {
|
||||
return nil, errors.New("value stored in Context is not a *Capture")
|
||||
}
|
||||
return capt, nil
|
||||
}
|
||||
|
||||
func (c Capture) ResponseSize() int64 {
|
||||
return c.rw.Size()
|
||||
}
|
||||
|
||||
func (c Capture) StatusCode() int {
|
||||
return c.rw.Status()
|
||||
}
|
||||
|
||||
// RequestSize returns the size of the request's body if it applies,
|
||||
// zero otherwise.
|
||||
func (c Capture) RequestSize() int64 {
|
||||
if c.rr == nil {
|
||||
return 0
|
||||
}
|
||||
return c.rr.size
|
||||
}
|
||||
|
||||
type readCounter struct {
|
||||
// source ReadCloser from where the request body is read.
|
||||
source io.ReadCloser
|
||||
// size is total the number of bytes read.
|
||||
size int64
|
||||
}
|
||||
|
||||
func (r *readCounter) Read(p []byte) (int, error) {
|
||||
n, err := r.source.Read(p)
|
||||
r.size += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *readCounter) Close() error {
|
||||
return r.source.Close()
|
||||
}
|
||||
|
||||
var _ middlewares.Stateful = &responseWriterWithCloseNotify{}
|
||||
|
||||
type responseWriter interface {
|
||||
http.ResponseWriter
|
||||
Size() int64
|
||||
Status() int
|
||||
}
|
||||
|
||||
func newResponseWriter(rw http.ResponseWriter) responseWriter {
|
||||
capt := &captureResponseWriter{rw: rw}
|
||||
if _, ok := rw.(http.CloseNotifier); !ok {
|
||||
return capt
|
||||
}
|
||||
|
||||
return &responseWriterWithCloseNotify{capt}
|
||||
}
|
||||
|
||||
// captureResponseWriter is a wrapper of type http.ResponseWriter
|
||||
// that tracks response status and size.
|
||||
type captureResponseWriter struct {
|
||||
rw http.ResponseWriter
|
||||
status int
|
||||
size int64
|
||||
}
|
||||
|
||||
func (crw *captureResponseWriter) Header() http.Header {
|
||||
return crw.rw.Header()
|
||||
}
|
||||
|
||||
func (crw *captureResponseWriter) Size() int64 {
|
||||
return crw.size
|
||||
}
|
||||
|
||||
func (crw *captureResponseWriter) Status() int {
|
||||
return crw.status
|
||||
}
|
||||
|
||||
func (crw *captureResponseWriter) Write(b []byte) (int, error) {
|
||||
if crw.status == 0 {
|
||||
crw.status = http.StatusOK
|
||||
}
|
||||
|
||||
size, err := crw.rw.Write(b)
|
||||
crw.size += int64(size)
|
||||
|
||||
return size, err
|
||||
}
|
||||
|
||||
func (crw *captureResponseWriter) WriteHeader(s int) {
|
||||
crw.rw.WriteHeader(s)
|
||||
crw.status = s
|
||||
}
|
||||
|
||||
func (crw *captureResponseWriter) Flush() {
|
||||
if f, ok := crw.rw.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (crw *captureResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||||
if h, ok := crw.rw.(http.Hijacker); ok {
|
||||
return h.Hijack()
|
||||
}
|
||||
|
||||
return nil, nil, fmt.Errorf("not a hijacker: %T", crw.rw)
|
||||
}
|
||||
|
||||
type responseWriterWithCloseNotify struct {
|
||||
*captureResponseWriter
|
||||
}
|
||||
|
||||
// CloseNotify returns a channel that receives at most a
|
||||
// single value (true) when the client connection has gone away.
|
||||
func (r *responseWriterWithCloseNotify) CloseNotify() <-chan bool {
|
||||
return r.rw.(http.CloseNotifier).CloseNotify()
|
||||
}
|
||||
234
pkg/middlewares/capture/capture_test.go
Normal file
234
pkg/middlewares/capture/capture_test.go
Normal file
|
|
@ -0,0 +1,234 @@
|
|||
package capture
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/containous/alice"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCapture(t *testing.T) {
|
||||
wrapMiddleware := func(next http.Handler) (http.Handler, error) {
|
||||
return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
|
||||
capt, err := FromContext(req.Context())
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = fmt.Fprintf(rw, "%d,%d,%d,", capt.RequestSize(), capt.ResponseSize(), capt.StatusCode())
|
||||
require.NoError(t, err)
|
||||
|
||||
next.ServeHTTP(rw, req)
|
||||
|
||||
_, err = fmt.Fprintf(rw, ",%d,%d,%d", capt.RequestSize(), capt.ResponseSize(), capt.StatusCode())
|
||||
require.NoError(t, err)
|
||||
}), nil
|
||||
}
|
||||
|
||||
handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
|
||||
_, err := rw.Write([]byte("foo"))
|
||||
require.NoError(t, err)
|
||||
|
||||
all, err := io.ReadAll(req.Body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "bar", string(all))
|
||||
})
|
||||
|
||||
wrapped := WrapHandler(&Handler{})
|
||||
chain := alice.New()
|
||||
chain = chain.Append(wrapped)
|
||||
chain = chain.Append(wrapMiddleware)
|
||||
handlers, err := chain.Then(handler)
|
||||
require.NoError(t, err)
|
||||
|
||||
request, err := http.NewRequest(http.MethodGet, "/", bytes.NewReader([]byte("bar")))
|
||||
require.NoError(t, err)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
handlers.ServeHTTP(recorder, request)
|
||||
// 3 = len("bar")
|
||||
// 9 = len("0,0,0,toto")
|
||||
assert.Equal(t, "0,0,0,foo,3,9,200", recorder.Body.String())
|
||||
}
|
||||
|
||||
// BenchmarkCapture with response writer and request reader
|
||||
// $ go test -bench=. ./pkg/middlewares/capture/
|
||||
// goos: linux
|
||||
// goarch: amd64
|
||||
// pkg: github.com/traefik/traefik/v2/pkg/middlewares/capture
|
||||
// cpu: Intel(R) Core(TM) i7-10750H CPU @ 2.60GHz
|
||||
// BenchmarkCapture/2k-12 280507 4015 ns/op 510.03 MB/s 5072 B/op 14 allocs/op
|
||||
// BenchmarkCapture/20k-12 135726 8301 ns/op 2467.26 MB/s 41936 B/op 14 allocs/op
|
||||
// BenchmarkCapture/100k-12 45494 26059 ns/op 3929.54 MB/s 213968 B/op 14 allocs/op
|
||||
// BenchmarkCapture/2k_captured-12 263713 4356 ns/op 470.20 MB/s 5552 B/op 18 allocs/op
|
||||
// BenchmarkCapture/20k_captured-12 132243 8790 ns/op 2329.98 MB/s 42416 B/op 18 allocs/op
|
||||
// BenchmarkCapture/100k_captured-12 45650 26587 ns/op 3851.57 MB/s 214448 B/op 18 allocs/op
|
||||
// BenchmarkCapture/2k_body-12 274135 7471 ns/op 274.12 MB/s 5624 B/op 20 allocs/op
|
||||
// BenchmarkCapture/20k_body-12 130206 21149 ns/op 968.36 MB/s 42488 B/op 20 allocs/op
|
||||
// BenchmarkCapture/100k_body-12 41600 51716 ns/op 1980.06 MB/s 214520 B/op 20 allocs/op
|
||||
// PASS
|
||||
func BenchmarkCapture(b *testing.B) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
size int
|
||||
capture bool
|
||||
body bool
|
||||
}{
|
||||
{
|
||||
name: "2k",
|
||||
size: 2048,
|
||||
},
|
||||
{
|
||||
name: "20k",
|
||||
size: 20480,
|
||||
},
|
||||
{
|
||||
name: "100k",
|
||||
size: 102400,
|
||||
},
|
||||
{
|
||||
name: "2k captured",
|
||||
size: 2048,
|
||||
capture: true,
|
||||
},
|
||||
{
|
||||
name: "20k captured",
|
||||
size: 20480,
|
||||
capture: true,
|
||||
},
|
||||
{
|
||||
name: "100k captured",
|
||||
size: 102400,
|
||||
capture: true,
|
||||
},
|
||||
{
|
||||
name: "2k body",
|
||||
size: 2048,
|
||||
body: true,
|
||||
},
|
||||
{
|
||||
name: "20k body",
|
||||
size: 20480,
|
||||
body: true,
|
||||
},
|
||||
{
|
||||
name: "100k body",
|
||||
size: 102400,
|
||||
body: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range testCases {
|
||||
b.Run(test.name, func(b *testing.B) {
|
||||
baseBody := generateBytes(test.size)
|
||||
|
||||
next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
|
||||
n, err := rw.Write(baseBody)
|
||||
require.Equal(b, test.size, n)
|
||||
require.NoError(b, err)
|
||||
})
|
||||
|
||||
var body io.Reader
|
||||
if test.body {
|
||||
body = bytes.NewReader(baseBody)
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, "http://foo/", body)
|
||||
require.NoError(b, err)
|
||||
|
||||
chain := alice.New()
|
||||
if test.capture || test.body {
|
||||
captureWrapped := WrapHandler(&Handler{})
|
||||
chain = chain.Append(captureWrapped)
|
||||
}
|
||||
handlers, err := chain.Then(next)
|
||||
require.NoError(b, err)
|
||||
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(test.size))
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
runBenchmark(b, test.size, req, handlers)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func runBenchmark(b *testing.B, size int, req *http.Request, handler http.Handler) {
|
||||
b.Helper()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
handler.ServeHTTP(recorder, req)
|
||||
if code := recorder.Code; code != 200 {
|
||||
b.Fatalf("Expected 200 but got %d", code)
|
||||
}
|
||||
|
||||
assert.Equal(b, size, len(recorder.Body.String()))
|
||||
}
|
||||
|
||||
func generateBytes(length int) []byte {
|
||||
var value []byte
|
||||
for i := 0; i < length; i++ {
|
||||
value = append(value, 0x61+byte(i%26))
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func TestRequestReader(t *testing.T) {
|
||||
buff := bytes.NewBuffer([]byte("foo"))
|
||||
rr := readCounter{source: io.NopCloser(buff)}
|
||||
assert.Equal(t, int64(0), rr.size)
|
||||
|
||||
n, err := rr.Read([]byte("bar"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 3, n)
|
||||
|
||||
err = rr.Close()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(3), rr.size)
|
||||
}
|
||||
|
||||
type rwWithCloseNotify struct {
|
||||
*httptest.ResponseRecorder
|
||||
}
|
||||
|
||||
func (r *rwWithCloseNotify) CloseNotify() <-chan bool {
|
||||
panic("implement me")
|
||||
}
|
||||
|
||||
func TestCloseNotifier(t *testing.T) {
|
||||
testCases := []struct {
|
||||
rw http.ResponseWriter
|
||||
desc string
|
||||
implementsCloseNotifier bool
|
||||
}{
|
||||
{
|
||||
rw: httptest.NewRecorder(),
|
||||
desc: "does not implement CloseNotifier",
|
||||
implementsCloseNotifier: false,
|
||||
},
|
||||
{
|
||||
rw: &rwWithCloseNotify{httptest.NewRecorder()},
|
||||
desc: "implements CloseNotifier",
|
||||
implementsCloseNotifier: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range testCases {
|
||||
test := test
|
||||
t.Run(test.desc, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
_, ok := test.rw.(http.CloseNotifier)
|
||||
assert.Equal(t, test.implementsCloseNotifier, ok)
|
||||
|
||||
rw := newResponseWriter(test.rw)
|
||||
_, impl := rw.(http.CloseNotifier)
|
||||
assert.Equal(t, test.implementsCloseNotifier, impl)
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue