Rework servers load-balancer to use the WRR

Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>
This commit is contained in:
Julien Salleyron 2022-11-16 11:38:07 +01:00 committed by GitHub
parent 67d9c8da0b
commit fadee5e87b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
70 changed files with 2085 additions and 2211 deletions

View file

@ -41,14 +41,6 @@ func (f *FieldHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
}
}
// AddServiceFields add service fields.
func AddServiceFields(rw http.ResponseWriter, req *http.Request, next http.Handler, data *LogData) {
data.Core[ServiceURL] = req.URL // note that this is *not* the original incoming URL
data.Core[ServiceAddr] = req.URL.Host
next.ServeHTTP(rw, req)
}
// AddOriginFields add origin fields.
func AddOriginFields(rw http.ResponseWriter, req *http.Request, next http.Handler, data *LogData) {
start := time.Now().UTC()

View file

@ -62,13 +62,13 @@ func FromContext(ctx context.Context) (Capture, error) {
// Capture is the object populated by the capture middleware,
// holding probes that allow to gather information about the request and response.
type Capture struct {
rr *readCounter
rw responseWriter
rr *readCounter
crw *captureResponseWriter
}
// NeedsReset returns whether the given http.ResponseWriter is the capture's probe.
func (c *Capture) NeedsReset(rw http.ResponseWriter) bool {
return c.rw != rw
return c.crw.rw != rw
}
// Reset returns a new handler that renews the Capture's probes, and inserts
@ -83,18 +83,18 @@ func (c *Capture) Reset(next http.Handler) http.Handler {
c.rr = readCounter
newReq.Body = readCounter
}
c.rw = newResponseWriter(rw)
c.crw = &captureResponseWriter{rw: rw}
next.ServeHTTP(c.rw, newReq)
next.ServeHTTP(c.crw, newReq)
})
}
func (c *Capture) ResponseSize() int64 {
return c.rw.Size()
return c.crw.Size()
}
func (c *Capture) StatusCode() int {
return c.rw.Status()
return c.crw.Status()
}
// RequestSize returns the size of the request's body if it applies,
@ -123,22 +123,7 @@ 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}
}
var _ middlewares.Stateful = &captureResponseWriter{}
// captureResponseWriter is a wrapper of type http.ResponseWriter
// that tracks response status and size.
@ -189,13 +174,3 @@ func (crw *captureResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error)
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()
}

View file

@ -189,44 +189,3 @@ func TestRequestReader(t *testing.T) {
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)
})
}
}

View file

@ -21,9 +21,8 @@ import (
// Compile time validation that the response recorder implements http interfaces correctly.
var (
// TODO: maybe remove at least for codeModifierWithCloseNotify.
_ middlewares.Stateful = &codeModifierWithCloseNotify{}
_ middlewares.Stateful = &codeCatcherWithCloseNotify{}
_ middlewares.Stateful = &codeModifier{}
_ middlewares.Stateful = &codeCatcher{}
)
const typeName = "customError"
@ -124,13 +123,6 @@ func newRequest(baseURL string) (*http.Request, error) {
return req, nil
}
type responseInterceptor interface {
http.ResponseWriter
http.Flusher
getCode() int
isFilteredCode() bool
}
// codeCatcher is a response writer that detects as soon as possible whether the
// response is a code within the ranges of codes it watches for. If it is, it
// simply drops the data from the response. Otherwise, it forwards it directly to
@ -144,27 +136,13 @@ type codeCatcher struct {
headersSent bool
}
type codeCatcherWithCloseNotify struct {
*codeCatcher
}
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone away.
func (cc *codeCatcherWithCloseNotify) CloseNotify() <-chan bool {
return cc.responseWriter.(http.CloseNotifier).CloseNotify()
}
func newCodeCatcher(rw http.ResponseWriter, httpCodeRanges types.HTTPCodeRanges) responseInterceptor {
catcher := &codeCatcher{
func newCodeCatcher(rw http.ResponseWriter, httpCodeRanges types.HTTPCodeRanges) *codeCatcher {
return &codeCatcher{
headerMap: make(http.Header),
code: http.StatusOK, // If backend does not call WriteHeader on us, we consider it's a 200.
responseWriter: rw,
httpCodeRanges: httpCodeRanges,
}
if _, ok := rw.(http.CloseNotifier); ok {
return &codeCatcherWithCloseNotify{catcher}
}
return catcher
}
func (cc *codeCatcher) Header() http.Header {
@ -240,24 +218,7 @@ func (cc *codeCatcher) Flush() {
// codeModifier forwards a response back to the client,
// while enforcing a given response code.
type codeModifier interface {
http.ResponseWriter
}
// newCodeModifier returns a codeModifier that enforces the given code.
func newCodeModifier(rw http.ResponseWriter, code int) codeModifier {
codeMod := &codeModifierWithoutCloseNotify{
headerMap: make(http.Header),
code: code,
responseWriter: rw,
}
if _, ok := rw.(http.CloseNotifier); ok {
return &codeModifierWithCloseNotify{codeMod}
}
return codeMod
}
type codeModifierWithoutCloseNotify struct {
type codeModifier struct {
code int // the code enforced in the response.
// headerSent is whether the headers have already been sent,
@ -268,18 +229,17 @@ type codeModifierWithoutCloseNotify struct {
responseWriter http.ResponseWriter
}
type codeModifierWithCloseNotify struct {
*codeModifierWithoutCloseNotify
}
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone away.
func (r *codeModifierWithCloseNotify) CloseNotify() <-chan bool {
return r.responseWriter.(http.CloseNotifier).CloseNotify()
// newCodeModifier returns a codeModifier that enforces the given code.
func newCodeModifier(rw http.ResponseWriter, code int) *codeModifier {
return &codeModifier{
headerMap: make(http.Header),
code: code,
responseWriter: rw,
}
}
// Header returns the response headers.
func (r *codeModifierWithoutCloseNotify) Header() http.Header {
func (r *codeModifier) Header() http.Header {
if r.headerMap == nil {
r.headerMap = make(http.Header)
}
@ -289,14 +249,14 @@ func (r *codeModifierWithoutCloseNotify) Header() http.Header {
// Write calls WriteHeader to send the enforced code,
// then writes the data directly to r.responseWriter.
func (r *codeModifierWithoutCloseNotify) Write(buf []byte) (int, error) {
func (r *codeModifier) Write(buf []byte) (int, error) {
r.WriteHeader(r.code)
return r.responseWriter.Write(buf)
}
// WriteHeader sends the headers, with the enforced code (the code in argument
// is always ignored), if it hasn't already been done.
func (r *codeModifierWithoutCloseNotify) WriteHeader(_ int) {
func (r *codeModifier) WriteHeader(_ int) {
if r.headerSent {
return
}
@ -307,7 +267,7 @@ func (r *codeModifierWithoutCloseNotify) WriteHeader(_ int) {
}
// Hijack hijacks the connection.
func (r *codeModifierWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
func (r *codeModifier) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := r.responseWriter.(http.Hijacker)
if !ok {
return nil, nil, fmt.Errorf("%T is not a http.Hijacker", r.responseWriter)
@ -316,7 +276,7 @@ func (r *codeModifierWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter,
}
// Flush sends any buffered data to the client.
func (r *codeModifierWithoutCloseNotify) Flush() {
func (r *codeModifier) Flush() {
r.WriteHeader(r.code)
if flusher, ok := r.responseWriter.(http.Flusher); ok {

View file

@ -188,50 +188,3 @@ type mockServiceBuilder struct {
func (m *mockServiceBuilder) BuildHTTP(_ context.Context, _ string) (http.Handler, error) {
return m.handler, nil
}
func TestNewResponseRecorder(t *testing.T) {
testCases := []struct {
desc string
rw http.ResponseWriter
expected http.ResponseWriter
}{
{
desc: "Without Close Notify",
rw: httptest.NewRecorder(),
expected: &codeModifierWithoutCloseNotify{},
},
{
desc: "With Close Notify",
rw: &mockRWCloseNotify{},
expected: &codeModifierWithCloseNotify{},
},
}
for _, test := range testCases {
test := test
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
rec := newCodeModifier(test.rw, 0)
assert.IsType(t, rec, test.expected)
})
}
}
type mockRWCloseNotify struct{}
func (m *mockRWCloseNotify) CloseNotify() <-chan bool {
panic("implement me")
}
func (m *mockRWCloseNotify) Header() http.Header {
panic("implement me")
}
func (m *mockRWCloseNotify) Write([]byte) (int, error) {
panic("implement me")
}
func (m *mockRWCloseNotify) WriteHeader(int) {
panic("implement me")
}

View file

@ -1,34 +0,0 @@
package emptybackendhandler
import (
"net/http"
"github.com/traefik/traefik/v2/pkg/healthcheck"
)
// EmptyBackend is a middleware that checks whether the current Backend
// has at least one active Server in respect to the healthchecks and if this
// is not the case, it will stop the middleware chain and respond with 503.
type emptyBackend struct {
healthcheck.BalancerStatusHandler
}
// New creates a new EmptyBackend middleware.
func New(lb healthcheck.BalancerStatusHandler) http.Handler {
return &emptyBackend{BalancerStatusHandler: lb}
}
// ServeHTTP responds with 503 when there is no active Server and otherwise
// invokes the next handler in the middleware chain.
func (e *emptyBackend) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
if len(e.BalancerStatusHandler.Servers()) != 0 {
e.BalancerStatusHandler.ServeHTTP(rw, req)
return
}
rw.WriteHeader(http.StatusServiceUnavailable)
if _, err := rw.Write([]byte(http.StatusText(http.StatusServiceUnavailable))); err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
}

View file

@ -1,85 +0,0 @@
package emptybackendhandler
import (
"fmt"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"github.com/stretchr/testify/assert"
"github.com/traefik/traefik/v2/pkg/testhelpers"
"github.com/vulcand/oxy/roundrobin"
)
func TestEmptyBackendHandler(t *testing.T) {
testCases := []struct {
amountServer int
expectedStatusCode int
}{
{
amountServer: 0,
expectedStatusCode: http.StatusServiceUnavailable,
},
{
amountServer: 1,
expectedStatusCode: http.StatusOK,
},
}
for _, test := range testCases {
test := test
t.Run(fmt.Sprintf("amount servers %d", test.amountServer), func(t *testing.T) {
t.Parallel()
handler := New(&healthCheckLoadBalancer{amountServer: test.amountServer})
recorder := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
handler.ServeHTTP(recorder, req)
assert.Equal(t, test.expectedStatusCode, recorder.Result().StatusCode)
})
}
}
type healthCheckLoadBalancer struct {
amountServer int
}
func (lb *healthCheckLoadBalancer) RegisterStatusUpdater(fn func(up bool)) error {
return nil
}
func (lb *healthCheckLoadBalancer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}
func (lb *healthCheckLoadBalancer) Servers() []*url.URL {
servers := make([]*url.URL, lb.amountServer)
for i := 0; i < lb.amountServer; i++ {
servers = append(servers, testhelpers.MustParseURL("http://localhost"))
}
return servers
}
func (lb *healthCheckLoadBalancer) RemoveServer(u *url.URL) error {
return nil
}
func (lb *healthCheckLoadBalancer) UpsertServer(u *url.URL, options ...roundrobin.ServerOption) error {
return nil
}
func (lb *healthCheckLoadBalancer) ServerWeight(u *url.URL) (int, bool) {
return 0, false
}
func (lb *healthCheckLoadBalancer) NextServer() (*url.URL, error) {
return nil, nil
}
func (lb *healthCheckLoadBalancer) Next() http.Handler {
return nil
}

View file

@ -23,17 +23,12 @@ type responseModifier struct {
// modifier can be nil.
func newResponseModifier(w http.ResponseWriter, r *http.Request, modifier func(*http.Response) error) http.ResponseWriter {
rm := &responseModifier{
return &responseModifier{
req: r,
rw: w,
modifier: modifier,
code: http.StatusOK,
}
if _, ok := w.(http.CloseNotifier); ok {
return responseModifierWithCloseNotify{responseModifier: rm}
}
return rm
}
func (r *responseModifier) WriteHeader(code int) {
@ -97,12 +92,3 @@ func (r *responseModifier) Flush() {
flusher.Flush()
}
}
type responseModifierWithCloseNotify struct {
*responseModifier
}
// CloseNotify implements http.CloseNotifier.
func (r *responseModifierWithCloseNotify) CloseNotify() <-chan bool {
return r.responseModifier.rw.(http.CloseNotifier).CloseNotify()
}

View file

@ -104,13 +104,6 @@ func WrapRouterHandler(ctx context.Context, registry metrics.Registry, routerNam
}
}
// WrapServiceHandler Wraps metrics service to alice.Constructor.
func WrapServiceHandler(ctx context.Context, registry metrics.Registry, serviceName string) alice.Constructor {
return func(next http.Handler) (http.Handler, error) {
return NewServiceMiddleware(ctx, next, registry, serviceName), nil
}
}
func (m *metricsMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
proto := getRequestProtocol(req)

View file

@ -60,47 +60,6 @@ func (m *collectingRetryMetrics) ServiceRetriesCounter() metrics.Counter {
return m.retriesCounter
}
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 := newResponseRecorder(test.rw)
_, impl := rw.(http.CloseNotifier)
assert.Equal(t, test.implementsCloseNotifier, impl)
})
}
}
func Test_getMethod(t *testing.T) {
testCases := []struct {
method string

View file

@ -1,63 +0,0 @@
package metrics
import (
"bufio"
"net"
"net/http"
)
type recorder interface {
http.ResponseWriter
http.Flusher
getCode() int
}
func newResponseRecorder(rw http.ResponseWriter) recorder {
rec := &responseRecorder{
ResponseWriter: rw,
statusCode: http.StatusOK,
}
if _, ok := rw.(http.CloseNotifier); !ok {
return rec
}
return &responseRecorderWithCloseNotify{rec}
}
// responseRecorder captures information from the response and preserves it for
// later analysis.
type responseRecorder struct {
http.ResponseWriter
statusCode int
}
type responseRecorderWithCloseNotify struct {
*responseRecorder
}
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone away.
func (r *responseRecorderWithCloseNotify) CloseNotify() <-chan bool {
return r.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
func (r *responseRecorder) getCode() int {
return r.statusCode
}
// WriteHeader captures the status code for later retrieval.
func (r *responseRecorder) WriteHeader(status int) {
r.ResponseWriter.WriteHeader(status)
r.statusCode = status
}
// Hijack hijacks the connection.
func (r *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return r.ResponseWriter.(http.Hijacker).Hijack()
}
// Flush sends any buffered data to the client.
func (r *responseRecorder) Flush() {
if f, ok := r.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}

View file

@ -1,70 +0,0 @@
package pipelining
import (
"bufio"
"context"
"net"
"net/http"
"github.com/traefik/traefik/v2/pkg/log"
"github.com/traefik/traefik/v2/pkg/middlewares"
)
const (
typeName = "Pipelining"
)
// pipelining returns a middleware.
type pipelining struct {
next http.Handler
}
// New returns a new pipelining instance.
func New(ctx context.Context, next http.Handler, name string) http.Handler {
log.FromContext(middlewares.GetLoggerCtx(ctx, name, typeName)).Debug("Creating middleware")
return &pipelining{
next: next,
}
}
func (p *pipelining) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
// https://github.com/golang/go/blob/3d59583836630cf13ec4bfbed977d27b1b7adbdc/src/net/http/server.go#L201-L218
if r.Method == http.MethodPut || r.Method == http.MethodPost {
p.next.ServeHTTP(rw, r)
} else {
p.next.ServeHTTP(&writerWithoutCloseNotify{rw}, r)
}
}
// writerWithoutCloseNotify helps to disable closeNotify.
type writerWithoutCloseNotify struct {
W http.ResponseWriter
}
// Header returns the response headers.
func (w *writerWithoutCloseNotify) Header() http.Header {
return w.W.Header()
}
// Write writes the data to the connection as part of an HTTP reply.
func (w *writerWithoutCloseNotify) Write(buf []byte) (int, error) {
return w.W.Write(buf)
}
// WriteHeader sends an HTTP response header with the provided status code.
func (w *writerWithoutCloseNotify) WriteHeader(code int) {
w.W.WriteHeader(code)
}
// Flush sends any buffered data to the client.
func (w *writerWithoutCloseNotify) Flush() {
if f, ok := w.W.(http.Flusher); ok {
f.Flush()
}
}
// Hijack hijacks the connection.
func (w *writerWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return w.W.(http.Hijacker).Hijack()
}

View file

@ -1,70 +0,0 @@
package pipelining
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/assert"
)
type recorderWithCloseNotify struct {
*httptest.ResponseRecorder
}
func (r *recorderWithCloseNotify) CloseNotify() <-chan bool {
panic("implement me")
}
func TestNew(t *testing.T) {
testCases := []struct {
desc string
HTTPMethod string
implementCloseNotifier bool
}{
{
desc: "should not implement CloseNotifier with GET method",
HTTPMethod: http.MethodGet,
implementCloseNotifier: false,
},
{
desc: "should implement CloseNotifier with PUT method",
HTTPMethod: http.MethodPut,
implementCloseNotifier: true,
},
{
desc: "should implement CloseNotifier with POST method",
HTTPMethod: http.MethodPost,
implementCloseNotifier: true,
},
{
desc: "should not implement CloseNotifier with GET method",
HTTPMethod: http.MethodHead,
implementCloseNotifier: false,
},
{
desc: "should not implement CloseNotifier with PROPFIND method",
HTTPMethod: "PROPFIND",
implementCloseNotifier: false,
},
}
for _, test := range testCases {
test := test
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, ok := w.(http.CloseNotifier)
assert.Equal(t, test.implementCloseNotifier, ok)
w.WriteHeader(http.StatusOK)
})
handler := New(context.Background(), nextHandler, "pipe")
req := httptest.NewRequest(test.HTTPMethod, "http://localhost", nil)
handler.ServeHTTP(&recorderWithCloseNotify{httptest.NewRecorder()}, req)
})
}
}

View file

@ -20,11 +20,9 @@ import (
)
// Compile time validation that the response writer implements http interfaces correctly.
var _ middlewares.Stateful = &responseWriterWithCloseNotify{}
var _ middlewares.Stateful = &responseWriter{}
const (
typeName = "Retry"
)
const typeName = "Retry"
// Listener is used to inform about retry attempts.
type Listener interface {
@ -149,57 +147,44 @@ func (l Listeners) Retried(req *http.Request, attempt int) {
}
}
type responseWriter interface {
http.ResponseWriter
http.Flusher
ShouldRetry() bool
DisableRetries()
}
func newResponseWriter(rw http.ResponseWriter, shouldRetry bool) responseWriter {
responseWriter := &responseWriterWithoutCloseNotify{
func newResponseWriter(rw http.ResponseWriter, shouldRetry bool) *responseWriter {
return &responseWriter{
responseWriter: rw,
headers: make(http.Header),
shouldRetry: shouldRetry,
}
if _, ok := rw.(http.CloseNotifier); ok {
return &responseWriterWithCloseNotify{
responseWriterWithoutCloseNotify: responseWriter,
}
}
return responseWriter
}
type responseWriterWithoutCloseNotify struct {
type responseWriter struct {
responseWriter http.ResponseWriter
headers http.Header
shouldRetry bool
written bool
}
func (r *responseWriterWithoutCloseNotify) ShouldRetry() bool {
func (r *responseWriter) ShouldRetry() bool {
return r.shouldRetry
}
func (r *responseWriterWithoutCloseNotify) DisableRetries() {
func (r *responseWriter) DisableRetries() {
r.shouldRetry = false
}
func (r *responseWriterWithoutCloseNotify) Header() http.Header {
func (r *responseWriter) Header() http.Header {
if r.written {
return r.responseWriter.Header()
}
return r.headers
}
func (r *responseWriterWithoutCloseNotify) Write(buf []byte) (int, error) {
func (r *responseWriter) Write(buf []byte) (int, error) {
if r.ShouldRetry() {
return len(buf), nil
}
return r.responseWriter.Write(buf)
}
func (r *responseWriterWithoutCloseNotify) WriteHeader(code int) {
func (r *responseWriter) WriteHeader(code int) {
if r.ShouldRetry() && code == http.StatusServiceUnavailable {
// We get a 503 HTTP Status Code when there is no backend server in the pool
// to which the request could be sent. Also, note that r.ShouldRetry()
@ -226,7 +211,7 @@ func (r *responseWriterWithoutCloseNotify) WriteHeader(code int) {
r.written = true
}
func (r *responseWriterWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
func (r *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := r.responseWriter.(http.Hijacker)
if !ok {
return nil, nil, fmt.Errorf("%T is not a http.Hijacker", r.responseWriter)
@ -234,16 +219,8 @@ func (r *responseWriterWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter
return hijacker.Hijack()
}
func (r *responseWriterWithoutCloseNotify) Flush() {
func (r *responseWriter) Flush() {
if flusher, ok := r.responseWriter.(http.Flusher); ok {
flusher.Flush()
}
}
type responseWriterWithCloseNotify struct {
*responseWriterWithoutCloseNotify
}
func (r *responseWriterWithCloseNotify) CloseNotify() <-chan bool {
return r.responseWriter.(http.CloseNotifier).CloseNotify()
}

View file

@ -8,5 +8,4 @@ type Stateful interface {
http.ResponseWriter
http.Hijacker
http.Flusher
http.CloseNotifier
}

View file

@ -48,7 +48,7 @@ func (e *entryPointMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Reque
req = req.WithContext(tracing.WithTracing(req.Context(), e.Tracing))
recorder := newStatusCodeRecoder(rw, http.StatusOK)
recorder := newStatusCodeRecorder(rw, http.StatusOK)
e.next.ServeHTTP(recorder, req)
tracing.LogResponseCode(span, recorder.Status())

View file

@ -51,7 +51,7 @@ func (f *forwarderMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Reques
tracing.InjectRequestHeaders(req)
recorder := newStatusCodeRecoder(rw, 200)
recorder := newStatusCodeRecorder(rw, 200)
f.next.ServeHTTP(recorder, req)

View file

@ -6,52 +6,35 @@ import (
"net/http"
)
type statusCodeRecoder interface {
http.ResponseWriter
Status() int
// newStatusCodeRecorder returns an initialized statusCodeRecoder.
func newStatusCodeRecorder(rw http.ResponseWriter, status int) *statusCodeRecorder {
return &statusCodeRecorder{rw, status}
}
// newStatusCodeRecoder returns an initialized statusCodeRecoder.
func newStatusCodeRecoder(rw http.ResponseWriter, status int) statusCodeRecoder {
recorder := &statusCodeWithoutCloseNotify{rw, status}
if _, ok := rw.(http.CloseNotifier); ok {
return &statusCodeWithCloseNotify{recorder}
}
return recorder
}
type statusCodeWithoutCloseNotify struct {
type statusCodeRecorder struct {
http.ResponseWriter
status int
}
// WriteHeader captures the status code for later retrieval.
func (s *statusCodeWithoutCloseNotify) WriteHeader(status int) {
func (s *statusCodeRecorder) WriteHeader(status int) {
s.status = status
s.ResponseWriter.WriteHeader(status)
}
// Status get response status.
func (s *statusCodeWithoutCloseNotify) Status() int {
func (s *statusCodeRecorder) Status() int {
return s.status
}
// Hijack hijacks the connection.
func (s *statusCodeWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
func (s *statusCodeRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return s.ResponseWriter.(http.Hijacker).Hijack()
}
// Flush sends any buffered data to the client.
func (s *statusCodeWithoutCloseNotify) Flush() {
func (s *statusCodeRecorder) Flush() {
if flusher, ok := s.ResponseWriter.(http.Flusher); ok {
flusher.Flush()
}
}
type statusCodeWithCloseNotify struct {
*statusCodeWithoutCloseNotify
}
func (s *statusCodeWithCloseNotify) CloseNotify() <-chan bool {
return s.ResponseWriter.(http.CloseNotifier).CloseNotify()
}