homelab/traefik
1
0
Fork 0
Code Activity

Add traffic size metrics

Browse source 
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 2022-09-12 17:10:09 +02:00 • committed by GitHub
parent 10528c973a
commit d578ed7327
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
26 changed files with 1125 additions and 339 deletions
Download patch file Download diff file Expand all files Collapse all files

20
pkg/middlewares/accesslog/capture_request_reader.go
View file

@ -1,20 +0,0 @@
package accesslog
import "io"
type captureRequestReader struct {
// source ReadCloser from where the request body is read.
source io.ReadCloser
// count Counts the number of bytes read (when captureRequestReader.Read is called).
count int64
}
func (r *captureRequestReader) Read(p []byte) (int, error) {
n, err := r.source.Read(p)
r.count += int64(n)
return n, err
}
func (r *captureRequestReader) Close() error {
return r.source.Close()
}

83
pkg/middlewares/accesslog/capture_response_writer.go
View file

@ -1,83 +0,0 @@
package accesslog
import (
"bufio"
"fmt"
"net"
"net/http"
"github.com/traefik/traefik/v2/pkg/middlewares"
)
var _ middlewares.Stateful = &captureResponseWriterWithCloseNotify{}
type capturer interface {
http.ResponseWriter
Size() int64
Status() int
}
func newCaptureResponseWriter(rw http.ResponseWriter) capturer {
capt := &captureResponseWriter{rw: rw}
if _, ok := rw.(http.CloseNotifier); !ok {
return capt
}
return &captureResponseWriterWithCloseNotify{capt}
}
// captureResponseWriter is a wrapper of type http.ResponseWriter
// that tracks request status and size.
type captureResponseWriter struct {
rw http.ResponseWriter
status int
size int64
}
type captureResponseWriterWithCloseNotify struct {
*captureResponseWriter
}
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone away.
func (r *captureResponseWriterWithCloseNotify) CloseNotify() <-chan bool {
return r.rw.(http.CloseNotifier).CloseNotify()
}
func (crw *captureResponseWriter) Header() http.Header {
return crw.rw.Header()
}
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)
}
func (crw *captureResponseWriter) Status() int {
return crw.status
}
func (crw *captureResponseWriter) Size() int64 {
return crw.size
}

50
pkg/middlewares/accesslog/capture_response_writer_test.go
View file

@ -1,50 +0,0 @@
package accesslog
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/assert"
)
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 := newCaptureResponseWriter(test.rw)
_, impl := rw.(http.CloseNotifier)
assert.Equal(t, test.implementsCloseNotifier, impl)
})
}
}

22
pkg/middlewares/accesslog/field_middleware.go
View file

@ -4,6 +4,8 @@ import (
"net/http"
"time"
"github.com/traefik/traefik/v2/pkg/log"
"github.com/traefik/traefik/v2/pkg/middlewares/capture"
"github.com/vulcand/oxy/utils"
)
@ -49,16 +51,24 @@ func AddServiceFields(rw http.ResponseWriter, req *http.Request, next http.Handl
// AddOriginFields add origin fields.
func AddOriginFields(rw http.ResponseWriter, req *http.Request, next http.Handler, data *LogData) {
crw := newCaptureResponseWriter(rw)
start := time.Now().UTC()
next.ServeHTTP(crw, req)
next.ServeHTTP(rw, req)
// use UTC to handle switchover of daylight saving correctly
data.Core[OriginDuration] = time.Now().UTC().Sub(start)
data.Core[OriginStatus] = crw.Status()
// make copy of headers so we can ensure there is no subsequent mutation during response processing
// make copy of headers, so we can ensure there is no subsequent mutation
// during response processing
data.OriginResponse = make(http.Header)
utils.CopyHeaders(data.OriginResponse, crw.Header())
data.Core[OriginContentSize] = crw.Size()
utils.CopyHeaders(data.OriginResponse, rw.Header())
ctx := req.Context()
capt, err := capture.FromContext(ctx)
if err != nil {
log.FromContext(log.With(ctx, log.Str(log.MiddlewareType, "AccessLogs"))).Errorf("Could not get Capture: %v", err)
return
}
data.Core[OriginStatus] = capt.StatusCode()
data.Core[OriginContentSize] = capt.ResponseSize()
}

43
pkg/middlewares/accesslog/logger.go
View file

@ -19,6 +19,7 @@ import (
"github.com/sirupsen/logrus"
ptypes "github.com/traefik/paerser/types"
"github.com/traefik/traefik/v2/pkg/log"
"github.com/traefik/traefik/v2/pkg/middlewares/capture"
traefiktls "github.com/traefik/traefik/v2/pkg/tls"
"github.com/traefik/traefik/v2/pkg/types"
)
@ -182,13 +183,17 @@ func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http
},
}
reqWithDataTable := req.WithContext(context.WithValue(req.Context(), DataTableKey, logDataTable))
defer func() {
if h.config.BufferingSize > 0 {
h.logHandlerChan <- handlerParams{
logDataTable: logDataTable,
}
return
}
h.logTheRoundTrip(logDataTable)
}()
var crr *captureRequestReader
if req.Body != nil {
crr = &captureRequestReader{source: req.Body, count: 0}
reqWithDataTable.Body = crr
}
reqWithDataTable := req.WithContext(context.WithValue(req.Context(), DataTableKey, logDataTable))
core[RequestCount] = nextRequestCount()
if req.Host != "" {
@ -222,30 +227,26 @@ func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http
core[ClientHost] = forwardedFor
}
crw := newCaptureResponseWriter(rw)
next.ServeHTTP(crw, reqWithDataTable)
next.ServeHTTP(rw, reqWithDataTable)
if _, ok := core[ClientUsername]; !ok {
core[ClientUsername] = usernameIfPresent(reqWithDataTable.URL)
}
logDataTable.DownstreamResponse = downstreamResponse{
headers: crw.Header().Clone(),
status: crw.Status(),
size: crw.Size(),
}
if crr != nil {
logDataTable.Request.size = crr.count
headers: rw.Header().Clone(),
}
if h.config.BufferingSize > 0 {
h.logHandlerChan <- handlerParams{
logDataTable: logDataTable,
}
} else {
h.logTheRoundTrip(logDataTable)
ctx := req.Context()
capt, err := capture.FromContext(ctx)
if err != nil {
log.FromContext(log.With(ctx, log.Str(log.MiddlewareType, "AccessLogs"))).Errorf("Could not get Capture: %v", err)
return
}
logDataTable.DownstreamResponse.status = capt.StatusCode()
logDataTable.DownstreamResponse.size = capt.ResponseSize()
logDataTable.Request.size = capt.RequestSize()
}
// Close closes the Logger (i.e. the file, drain logHandlerChan, etc).

66
pkg/middlewares/accesslog/logger_test.go
View file

@ -1,9 +1,11 @@
package accesslog
import (
"bytes"
"crypto/tls"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"net/url"
@ -14,9 +16,11 @@ import (
"testing"
"time"
"github.com/containous/alice"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
ptypes "github.com/traefik/paerser/types"
"github.com/traefik/traefik/v2/pkg/middlewares/capture"
"github.com/traefik/traefik/v2/pkg/types"
)
@ -46,23 +50,29 @@ func TestLogRotation(t *testing.T) {
config := &types.AccessLog{FilePath: fileName, Format: CommonFormat}
logHandler, err := NewHandler(config)
if err != nil {
t.Fatalf("Error creating new log handler: %s", err)
}
defer logHandler.Close()
require.NoError(t, err)
t.Cleanup(func() {
err := logHandler.Close()
require.NoError(t, err)
})
chain := alice.New()
chain = chain.Append(capture.WrapHandler(&capture.Handler{}))
chain = chain.Append(WrapHandler(logHandler))
handler, err := chain.Then(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusOK)
}))
require.NoError(t, err)
recorder := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
next := func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusOK)
}
iterations := 20
halfDone := make(chan bool)
writeDone := make(chan bool)
go func() {
for i := 0; i < iterations; i++ {
logHandler.ServeHTTP(recorder, req, http.HandlerFunc(next))
handler.ServeHTTP(recorder, req)
if i == iterations/2 {
halfDone <- true
}
@ -178,7 +188,10 @@ func TestLoggerHeaderFields(t *testing.T) {
logger, err := NewHandler(config)
require.NoError(t, err)
defer logger.Close()
t.Cleanup(func() {
err := logger.Close()
require.NoError(t, err)
})
if config.FilePath != "" {
_, err = os.Stat(config.FilePath)
@ -196,9 +209,14 @@ func TestLoggerHeaderFields(t *testing.T) {
req.Header.Add(test.header, s)
}
logger.ServeHTTP(httptest.NewRecorder(), req, http.HandlerFunc(func(writer http.ResponseWriter, r *http.Request) {
writer.WriteHeader(http.StatusOK)
chain := alice.New()
chain = chain.Append(capture.WrapHandler(&capture.Handler{}))
chain = chain.Append(WrapHandler(logger))
handler, err := chain.Then(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusOK)
}))
require.NoError(t, err)
handler.ServeHTTP(httptest.NewRecorder(), req)
logData, err := os.ReadFile(logFile.Name())
require.NoError(t, err)
@ -224,11 +242,14 @@ func TestLoggerCLF(t *testing.T) {
assertValidLogData(t, expectedLog, logData)
}
func TestAsyncLoggerCLF(t *testing.T) {
func TestLoggerCLFWithBufferingSize(t *testing.T) {
logFilePath := filepath.Join(t.TempDir(), logFileNameSuffix)
config := &types.AccessLog{FilePath: logFilePath, Format: CommonFormat, BufferingSize: 1024}
doLogging(t, config)
// wait a bit for the buffer to be written in the file.
time.Sleep(50 * time.Millisecond)
logData, err := os.ReadFile(logFilePath)
require.NoError(t, err)
@ -691,9 +712,7 @@ func assertValidLogData(t *testing.T, expected string, logData []byte) {
resultExpected, err := ParseAccessLog(expected)
require.NoError(t, err)
formatErrMessage := fmt.Sprintf(`
Expected: %s
Actual: %s`, expected, string(logData))
formatErrMessage := fmt.Sprintf("Expected:\t%q\nActual:\t%q", expected, string(logData))
require.Equal(t, len(resultExpected), len(result), formatErrMessage)
assert.Equal(t, resultExpected[ClientHost], result[ClientHost], formatErrMessage)
@ -722,7 +741,7 @@ func captureStdout(t *testing.T) (out *os.File, restoreStdout func()) {
restoreStdout = func() {
os.Stdout = original
os.RemoveAll(file.Name())
_ = os.RemoveAll(file.Name())
}
return file, restoreStdout
@ -730,10 +749,12 @@ func captureStdout(t *testing.T) (out *os.File, restoreStdout func()) {
func doLoggingTLSOpt(t *testing.T, config *types.AccessLog, enableTLS bool) {
t.Helper()
logger, err := NewHandler(config)
require.NoError(t, err)
defer logger.Close()
t.Cleanup(func() {
err := logger.Close()
require.NoError(t, err)
})
if config.FilePath != "" {
_, err = os.Stat(config.FilePath)
@ -753,6 +774,7 @@ func doLoggingTLSOpt(t *testing.T, config *types.AccessLog, enableTLS bool) {
User: url.UserPassword(testUsername, ""),
Path: testPath,
},
Body: io.NopCloser(bytes.NewReader([]byte("bar"))),
}
if enableTLS {
req.TLS = &tls.ConnectionState{
@ -761,7 +783,13 @@ func doLoggingTLSOpt(t *testing.T, config *types.AccessLog, enableTLS bool) {
}
}
logger.ServeHTTP(httptest.NewRecorder(), req, http.HandlerFunc(logWriterTestHandlerFunc))
chain := alice.New()
chain = chain.Append(capture.WrapHandler(&capture.Handler{}))
chain = chain.Append(WrapHandler(logger))
handler, err := chain.Then(http.HandlerFunc(logWriterTestHandlerFunc))
require.NoError(t, err)
handler.ServeHTTP(httptest.NewRecorder(), req)
}
func doLoggingTLS(t *testing.T, config *types.AccessLog) {

198
pkg/middlewares/capture/capture.go Normal file
View 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
View 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)
})
}
}

25
pkg/middlewares/metrics/metrics.go
View file

@ -13,6 +13,7 @@ import (
"github.com/traefik/traefik/v2/pkg/log"
"github.com/traefik/traefik/v2/pkg/metrics"
"github.com/traefik/traefik/v2/pkg/middlewares"
"github.com/traefik/traefik/v2/pkg/middlewares/capture"
"github.com/traefik/traefik/v2/pkg/middlewares/retry"
traefiktls "github.com/traefik/traefik/v2/pkg/tls"
)
@ -32,6 +33,8 @@ type metricsMiddleware struct {
reqsTLSCounter gokitmetrics.Counter
reqDurationHistogram metrics.ScalableHistogram
openConnsGauge gokitmetrics.Gauge
reqsBytesCounter gokitmetrics.Counter
respsBytesCounter gokitmetrics.Counter
baseLabels []string
}
@ -45,6 +48,8 @@ func NewEntryPointMiddleware(ctx context.Context, next http.Handler, registry me
reqsTLSCounter: registry.EntryPointReqsTLSCounter(),
reqDurationHistogram: registry.EntryPointReqDurationHistogram(),
openConnsGauge: registry.EntryPointOpenConnsGauge(),
reqsBytesCounter: registry.EntryPointReqsBytesCounter(),
respsBytesCounter: registry.EntryPointRespsBytesCounter(),
baseLabels: []string{"entrypoint", entryPointName},
}
}
@ -59,6 +64,8 @@ func NewRouterMiddleware(ctx context.Context, next http.Handler, registry metric
reqsTLSCounter: registry.RouterReqsTLSCounter(),
reqDurationHistogram: registry.RouterReqDurationHistogram(),
openConnsGauge: registry.RouterOpenConnsGauge(),
reqsBytesCounter: registry.RouterReqsBytesCounter(),
respsBytesCounter: registry.RouterRespsBytesCounter(),
baseLabels: []string{"router", routerName, "service", serviceName},
}
}
@ -73,6 +80,8 @@ func NewServiceMiddleware(ctx context.Context, next http.Handler, registry metri
reqsTLSCounter: registry.ServiceReqsTLSCounter(),
reqDurationHistogram: registry.ServiceReqDurationHistogram(),
openConnsGauge: registry.ServiceOpenConnsGauge(),
reqsBytesCounter: registry.ServiceReqsBytesCounter(),
respsBytesCounter: registry.ServiceRespsBytesCounter(),
baseLabels: []string{"service", serviceName},
}
}
@ -116,16 +125,22 @@ func (m *metricsMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Request)
m.reqsTLSCounter.With(tlsLabels...).Add(1)
}
recorder := newResponseRecorder(rw)
start := time.Now()
m.next.ServeHTTP(recorder, req)
m.next.ServeHTTP(rw, req)
labels = append(labels, "code", strconv.Itoa(recorder.getCode()))
ctx := req.Context()
capt, err := capture.FromContext(ctx)
if err != nil {
log.FromContext(middlewares.GetLoggerCtx(ctx, nameEntrypoint, typeName)).Errorf("Could not get Capture: %w", err)
return
}
labels = append(labels, "code", strconv.Itoa(capt.StatusCode()))
m.reqDurationHistogram.With(labels...).ObserveFromStart(start)
m.reqsCounter.With(labels...).Add(1)
m.respsBytesCounter.With(labels...).Add(float64(capt.ResponseSize()))
m.reqsBytesCounter.With(labels...).Add(float64(capt.RequestSize()))
}
func getRequestProtocol(req *http.Request) string {
@ -201,6 +216,6 @@ type RetryListener struct {
}
// Retried tracks the retry in the RequestMetrics implementation.
func (m *RetryListener) Retried(req *http.Request, attempt int) {
func (m *RetryListener) Retried(_ *http.Request, _ int) {
m.retryMetrics.ServiceRetriesCounter().With("service", m.serviceName).Add(1)
}