Rework servers load-balancer to use the WRR
Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>
This commit is contained in:
parent
67d9c8da0b
commit
fadee5e87b
70 changed files with 2085 additions and 2211 deletions
|
@ -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()
|
||||
|
|
|
@ -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()
|
||||
}
|
||||
|
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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 {
|
||||
|
|
|
@ -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")
|
||||
}
|
||||
|
|
|
@ -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
|
||||
}
|
||||
}
|
|
@ -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
|
||||
}
|
|
@ -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()
|
||||
}
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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()
|
||||
}
|
||||
}
|
|
@ -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()
|
||||
}
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
}
|
|
@ -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()
|
||||
}
|
||||
|
|
|
@ -8,5 +8,4 @@ type Stateful interface {
|
|||
http.ResponseWriter
|
||||
http.Hijacker
|
||||
http.Flusher
|
||||
http.CloseNotifier
|
||||
}
|
||||
|
|
|
@ -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())
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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()
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue