fix(buffer): incorrect short-circuiting in buffer logic (#253)

Author: ldez

buffer/buffer.go

@@ -250,7 +250,7 @@ func (b *Buffer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 			reader = rdr
 		}
 
-		if body == nil || (b.retryPredicate == nil || attempt > DefaultMaxRetryAttempts) ||
+		if b.retryPredicate == nil || attempt > DefaultMaxRetryAttempts ||
 			!b.retryPredicate(&context{r: req, attempt: attempt, responseCode: bw.code}) {
 			utils.CopyHeaders(w.Header(), bw.Header())
 			w.WriteHeader(bw.code)

roundrobin/rr_test.go

@@ -1,12 +1,15 @@
 package roundrobin
 
 import (
+	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/vulcand/oxy/v2/buffer"
 	"github.com/vulcand/oxy/v2/forward"
 	"github.com/vulcand/oxy/v2/testutils"
 	"github.com/vulcand/oxy/v2/utils"
@@ -201,13 +204,55 @@ func TestRoundRobinRequestRewriteListener(t *testing.T) {
 	assert.NotNil(t, lb.requestRewriteListener)
 }
 
-func seq(t *testing.T, url string, repeat int) []string {
+func TestRoundRobinFailWell(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+
+	rr := httptest.NewRecorder()
+
+	var oneCount atomic.Int32
+
+	oneServer := testutils.NewResponderWithCount(&oneCount)
+
+	t.Cleanup(func() { oneServer.Close() })
+
+	var twoCount atomic.Int32
+
+	twoServer := testutils.NewResponderWithCount(&twoCount)
+
+	// explicit close now
+	twoServer.Close()
+
+	lb, err := New(forward.New(false))
+	require.NoError(t, err)
+
+	require.NoError(t, lb.UpsertServer(testutils.MustParseRequestURI(oneServer.URL)))
+	require.NoError(t, lb.UpsertServer(testutils.MustParseRequestURI(twoServer.URL)))
+
+	buff, err := buffer.New(lb, buffer.Retry(fmt.Sprintf("IsNetworkError() && Attempts() < %d", 2)))
+	require.NoError(t, err)
+
+	okCount := 0
+
+	for range 10 {
+		buff.ServeHTTP(rr, req)
+
+		assert.Equal(t, http.StatusOK, rr.Code)
+
+		okCount++
+	}
+
+	assert.Equal(t, 10, okCount)
+	assert.Equal(t, int32(10), oneCount.Load())
+	assert.Zero(t, twoCount.Load())
+}
+
+func seq(t *testing.T, uri string, repeat int) []string {
 	t.Helper()
 
 	var out []string
 
 	for range repeat {
-		_, body, err := testutils.Get(url)
+		_, body, err := testutils.Get(uri)
 		require.NoError(t, err)
 
 		out = append(out, string(body))

testutils/utils.go

@@ -8,6 +8,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/vulcand/oxy/v2/internal/holsterv4/clock"
@@ -19,7 +20,7 @@ func NewHandler(handler http.HandlerFunc) *httptest.Server {
 	return httptest.NewServer(handler)
 }
 
-// NewResponder creates a new Server with response.
+// NewResponder creates a new Server with a response.
 func NewResponder(t *testing.T, response string) *httptest.Server {
 	t.Helper()
 
@@ -32,6 +33,15 @@ func NewResponder(t *testing.T, response string) *httptest.Server {
 	return server
 }
 
+// NewResponderWithCount creates a new Server with a call counter.
+func NewResponderWithCount(count *atomic.Int32) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		count.Add(1)
+
+		_, _ = w.Write([]byte("Ok"))
+	}))
+}
+
 // MustParseRequestURI is the version of url.ParseRequestURI that panics if incorrect, helpful to shorten the tests.
 func MustParseRequestURI(uri string) *url.URL {
 	out, err := url.ParseRequestURI(uri)