homelab/traefik
1
0
Fork 0
Code Activity

Enable the retry middleware in the proxy

Browse source 
Co-authored-by: Romain <rtribotte@users.noreply.github.com>
This commit is contained in:
Kevin Pollet 2025-02-21 11:36:05 +01:00 • committed by GitHub
parent c2a294c872
commit f196de90e1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 99 additions and 64 deletions
Download patch file Download diff file Expand all files Collapse all files

63
pkg/middlewares/retry/retry_test.go
View file

@ -105,12 +105,21 @@ func TestRetry(t *testing.T) {
t.Parallel()
retryAttempts := 0
next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
// This signals that a connection will be established with the backend
// to enable the Retry middleware mechanism.
shouldRetry := ContextShouldRetry(req.Context())
if shouldRetry != nil {
shouldRetry(true)
}
retryAttempts++
if retryAttempts > test.amountFaultyEndpoints {
// calls WroteHeaders on httptrace.
_ = r.Write(io.Discard)
// This signals that request headers have been sent to the backend.
if shouldRetry != nil {
shouldRetry(false)
}
rw.WriteHeader(http.StatusOK)
return
@ -152,26 +161,16 @@ func TestRetryEmptyServerList(t *testing.T) {
assert.Equal(t, 0, retryListener.timesCalled)
}
func TestRetryListeners(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/", nil)
retryListeners := Listeners{&countingRetryListener{}, &countingRetryListener{}}
retryListeners.Retried(req, 1)
retryListeners.Retried(req, 1)
for _, retryListener := range retryListeners {
listener := retryListener.(*countingRetryListener)
if listener.timesCalled != 2 {
t.Errorf("retry listener was called %d time(s), want %d time(s)", listener.timesCalled, 2)
}
}
}
func TestMultipleRetriesShouldNotLooseHeaders(t *testing.T) {
attempt := 0
expectedHeaderValue := "bar"
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
shouldRetry := ContextShouldRetry(req.Context())
if shouldRetry != nil {
shouldRetry(true)
}
headerName := fmt.Sprintf("X-Foo-Test-%d", attempt)
rw.Header().Add(headerName, expectedHeaderValue)
if attempt < 2 {
@ -179,9 +178,8 @@ func TestMultipleRetriesShouldNotLooseHeaders(t *testing.T) {
return
}
// Request has been successfully written to backend.
trace := httptrace.ContextClientTrace(req.Context())
trace.WroteHeaders()
// Request has been successfully written to backend
shouldRetry(false)
// And we decide to answer to client.
rw.WriteHeader(http.StatusNoContent)
@ -212,9 +210,10 @@ func TestRetryShouldNotLooseHeadersOnWrite(t *testing.T) {
rw.Header().Add("X-Foo-Test", "bar")
// Request has been successfully written to backend.
trace := httptrace.ContextClientTrace(req.Context())
trace.WroteHeaders()
shouldRetry := ContextShouldRetry(req.Context())
if shouldRetry != nil {
shouldRetry(false)
}
// And we decide to answer to client without calling WriteHeader.
_, err := rw.Write([]byte("bar"))
require.NoError(t, err)
@ -285,12 +284,24 @@ func TestRetryWebsocket(t *testing.T) {
t.Parallel()
retryAttempts := 0
next := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
// This signals that a connection will be established with the backend
// to enable the Retry middleware mechanism.
shouldRetry := ContextShouldRetry(req.Context())
if shouldRetry != nil {
shouldRetry(true)
}
retryAttempts++
if retryAttempts > test.amountFaultyEndpoints {
// This signals that request headers have been sent to the backend.
if shouldRetry != nil {
shouldRetry(false)
}
upgrader := websocket.Upgrader{}
_, err := upgrader.Upgrade(rw, r, nil)
_, err := upgrader.Upgrade(rw, req, nil)
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
}