From d0b684a26643404a9efeb0ae90069d6facc8a467 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Wed, 17 Jan 2024 11:53:22 +0100
Subject: [PATCH] Fixes and tests

---
 pkg/proxy/middleware.go               | 138 +++++++++++++++++---------
 pkg/proxy/middleware_test.go          | 115 ++++++++++++++++++++-
 tests/integration/integration_test.go |   2 +-
 3 files changed, 204 insertions(+), 51 deletions(-)

diff --git a/pkg/proxy/middleware.go b/pkg/proxy/middleware.go
index e8a5820f..6964a84a 100644
--- a/pkg/proxy/middleware.go
+++ b/pkg/proxy/middleware.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"errors"
 	"io"
-	"log"
 	"math/rand"
 	"net/http"
 	"time"
@@ -44,35 +43,35 @@ func NewRetryMiddleware(maxRetries int, other http.Handler, optRetryStatusCodes
 	}
 }
 
+type xResponseWriter interface {
+	http.ResponseWriter
+	discardedResponse() bool
+	capturedStatusCode() int
+}
+type xBodyCapturer interface {
+	io.ReadCloser
+	Capture()
+}
+
 func (r RetryMiddleware) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
-	lazyBody := &lazyBodyCapturer{
-		reader: request.Body,
-		buf:    bytes.NewBuffer([]byte{}),
-	}
+	lazyBody := newLazyBodyCapturer(request.Body)
 	request.Body = lazyBody
 	for i := 0; ; i++ {
-		capturedResp := &responseWriterDelegator{
-			isRetryableStatusCode: r.isRetryableStatusCode,
-			ResponseWriter:        writer,
-			headerBuf:             make(http.Header),
-			discardErrResp: i < r.maxRetries &&
-				request.Context().Err() == nil, // abort early on timeout, context cancel
-		}
+		discardErrResp := i < r.maxRetries && request.Context().Err() == nil
+		capturedResp := newResponseWriterDelegator(writer, r.isRetryableStatusCode, discardErrResp)
 		// call next handler in chain
-		req, err := http.NewRequestWithContext(request.Context(), request.Method, request.URL.String(), lazyBody)
-		if err != nil {
-			log.Printf("clone request: %v", err)
-			writer.WriteHeader(http.StatusInternalServerError)
-			return
-		}
-		r.nextHandler.ServeHTTP(capturedResp, req)
-		lazyBody.Capture()
-		if !capturedResp.discardErrResp || // max retries reached
-			!r.isRetryableStatusCode(capturedResp.statusCode) {
+		reqClone := request.Clone(request.Context()) // also copies the reference to the lazy body capturer
+		r.nextHandler.ServeHTTP(capturedResp, reqClone)
+
+		if !capturedResp.discardedResponse() || // max retries reached or context error
+			!r.isRetryableStatusCode(capturedResp.capturedStatusCode()) {
 			break
 		}
+		// setup for retry
+		lazyBody.Capture()
+
 		totalRetries.Inc()
-		// Exponential backoff
+		// exponential backoff
 		jitter := time.Duration(r.rSource.Intn(50))
 		time.Sleep(time.Millisecond*time.Duration(1<<uint(i)) + jitter)
 	}
@@ -85,12 +84,12 @@ func (r RetryMiddleware) isRetryableStatusCode(status int) bool {
 
 var (
 	_ http.Flusher  = &responseWriterDelegator{}
-	_ io.ReaderFrom = &responseWriterDelegator{}
+	_ io.ReaderFrom = &xResponseWriterDelegator{}
 )
 
 // responseWriterDelegator represents a wrapper around http.ResponseWriter that provides additional
 // functionalities for handling response writing. Depending on the status code and discard settings,
-// the heeader + content on write is skipped so that it can be re-used on retry.
+// the header + content on write is skipped so that it can be re-used on retry.
 type responseWriterDelegator struct {
 	http.ResponseWriter
 	headerBuf   http.Header
@@ -101,6 +100,28 @@ type responseWriterDelegator struct {
 	isRetryableStatusCode func(status int) bool
 }
 
+// newResponseWriterDelegator constructor
+func newResponseWriterDelegator(writer http.ResponseWriter, isRetryableStatusCode func(status int) bool, discardErrResp bool) xResponseWriter {
+	d := &responseWriterDelegator{
+		isRetryableStatusCode: isRetryableStatusCode,
+		ResponseWriter:        writer,
+		headerBuf:             make(http.Header),
+		discardErrResp:        discardErrResp, // abort early on timeout, context cancel
+	}
+	if _, ok := writer.(io.ReaderFrom); ok {
+		return &xResponseWriterDelegator{responseWriterDelegator: d}
+	}
+	return d
+}
+
+func (r *responseWriterDelegator) discardedResponse() bool {
+	return r.discardErrResp
+}
+
+func (r *responseWriterDelegator) capturedStatusCode() int {
+	return r.statusCode
+}
+
 func (r *responseWriterDelegator) Header() http.Header {
 	return r.headerBuf
 }
@@ -141,7 +162,22 @@ func (r *responseWriterDelegator) Write(data []byte) (int, error) {
 	}
 }
 
-func (r *responseWriterDelegator) ReadFrom(re io.Reader) (int64, error) {
+func (r *responseWriterDelegator) Flush() {
+	if f, ok := r.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+// xResponseWriterDelegator provides the same functionalities as responseWriterDelegator but also implements the
+// io.ReaderFrom interface.
+// The ReadFrom method ensures that the header is set before reading from the reader.
+// In case discardErrResp is true and the response status code is retryable, the content is discarded.
+// Otherwise, it calls the ReadFrom method of the underlying response writer and returns the result.
+type xResponseWriterDelegator struct {
+	*responseWriterDelegator
+}
+
+func (r *xResponseWriterDelegator) ReadFrom(re io.Reader) (int64, error) {
 	// ensure header is set. default is 200 in Go stdlib
 	if !r.wroteHeader {
 		r.WriteHeader(http.StatusOK)
@@ -153,17 +189,14 @@ func (r *responseWriterDelegator) ReadFrom(re io.Reader) (int64, error) {
 	}
 }
 
-func (r *responseWriterDelegator) Flush() {
-	if f, ok := r.ResponseWriter.(http.Flusher); ok {
-		f.Flush()
-	}
-}
-
 var (
 	_ io.ReadCloser = &lazyBodyCapturer{}
-	_ io.WriterTo   = &lazyBodyCapturer{}
+	_ io.WriterTo   = &lazyBodyCapturerWriteTo{}
 )
 
+// lazyBodyCapturer represents a type that captures the request body lazily.
+// It wraps an io.ReadCloser and provides methods for reading, closing,
+// writing to an io.Writer, and capturing the body content.
 type lazyBodyCapturer struct {
 	reader       io.ReadCloser
 	capturedBody []byte
@@ -171,6 +204,18 @@ type lazyBodyCapturer struct {
 	allRead      bool
 }
 
+// newLazyBodyCapturer constructor
+func newLazyBodyCapturer(body io.ReadCloser) xBodyCapturer {
+	c := &lazyBodyCapturer{
+		reader: body,
+		buf:    bytes.NewBuffer([]byte{}),
+	}
+	if _, ok := c.reader.(io.WriterTo); ok {
+		return &lazyBodyCapturerWriteTo{c}
+	}
+	return c
+}
+
 func (m *lazyBodyCapturer) Read(p []byte) (int, error) {
 	if m.allRead {
 		return m.reader.Read(p)
@@ -186,23 +231,24 @@ func (m *lazyBodyCapturer) Close() error {
 	return m.reader.Close()
 }
 
-func (m *lazyBodyCapturer) WriteTo(w io.Writer) (int64, error) {
-	if m.allRead {
-		return m.reader.(io.WriterTo).WriteTo(w)
-	}
-	n, err := m.reader.(io.WriterTo).WriteTo(io.MultiWriter(w, m.buf))
-	if errors.Is(err, io.EOF) {
-		m.allRead = true
-	}
-	return n, err
-}
-
 func (m *lazyBodyCapturer) Capture() {
 	m.allRead = true
 	if m.buf != nil {
 		m.capturedBody = m.buf.Bytes()
 		m.buf = nil
-	} else {
-		m.reader = io.NopCloser(bytes.NewReader(m.capturedBody))
 	}
+	m.reader = io.NopCloser(bytes.NewReader(m.capturedBody))
+}
+
+type lazyBodyCapturerWriteTo struct {
+	*lazyBodyCapturer
+}
+
+func (m *lazyBodyCapturerWriteTo) WriteTo(w io.Writer) (int64, error) {
+	if m.allRead {
+		return m.reader.(io.WriterTo).WriteTo(w)
+	}
+	n, err := m.reader.(io.WriterTo).WriteTo(io.MultiWriter(w, m.buf))
+	m.allRead = true
+	return n, err
 }
diff --git a/pkg/proxy/middleware_test.go b/pkg/proxy/middleware_test.go
index d11f1a11..452fee12 100644
--- a/pkg/proxy/middleware_test.go
+++ b/pkg/proxy/middleware_test.go
@@ -1,9 +1,13 @@
 package proxy
 
 import (
+	"bytes"
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -12,38 +16,45 @@ import (
 )
 
 func TestServeHTTP(t *testing.T) {
+	myHeader := map[string][]string{"Foo": {"bar1", "bar2"}}
 	specs := map[string]struct {
 		context    func() context.Context
 		maxRetries int
+		headers    http.Header
 		respStatus int
 		expRetries int
 	}{
 		"no retry on 200": {
 			context:    func() context.Context { return context.TODO() },
+			headers:    myHeader,
 			maxRetries: 3,
 			respStatus: http.StatusOK,
 			expRetries: 0,
 		},
 		"no retry on 500": {
 			context:    func() context.Context { return context.TODO() },
+			headers:    myHeader,
 			maxRetries: 3,
 			respStatus: http.StatusInternalServerError,
 			expRetries: 0,
 		},
 		"max retries on 503": {
 			context:    func() context.Context { return context.TODO() },
+			headers:    myHeader,
 			maxRetries: 3,
 			respStatus: http.StatusServiceUnavailable,
 			expRetries: 3,
 		},
 		"max retries on 502": {
 			context:    func() context.Context { return context.TODO() },
+			headers:    myHeader,
 			maxRetries: 3,
 			respStatus: http.StatusBadGateway,
 			expRetries: 3,
 		},
 		"not buffered on 100": {
 			context:    func() context.Context { return context.TODO() },
+			headers:    myHeader,
 			maxRetries: 3,
 			respStatus: http.StatusContinue,
 			expRetries: 0,
@@ -54,6 +65,7 @@ func TestServeHTTP(t *testing.T) {
 				cancel()
 				return ctx
 			},
+			headers:    myHeader,
 			maxRetries: 3,
 			respStatus: http.StatusBadGateway,
 			expRetries: 0,
@@ -62,31 +74,106 @@ func TestServeHTTP(t *testing.T) {
 	for name, spec := range specs {
 		t.Run(name, func(t *testing.T) {
 			counterBefore := counterValue(t, totalRetries)
-			req, err := http.NewRequestWithContext(spec.context(), "GET", "/test", nil)
+			const myBody = "my-request-body"
+			req, err := http.NewRequestWithContext(spec.context(), "GET", "/test", strings.NewReader(myBody))
 			require.NoError(t, err)
+			req.Header = spec.headers.Clone()
 
 			respRecorder := httptest.NewRecorder()
 
 			var counter int
-			testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			testBackend := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				counter++
+				// return all headers
+				for k, vals := range r.Header {
+					for _, v := range vals {
+						w.Header().Add(k, v)
+					}
+				}
 				w.WriteHeader(spec.respStatus)
+				reqBody, err := io.ReadAll(req.Body)
+				require.NoError(t, err)
+				_, err = w.Write(append(reqBody, []byte(strconv.Itoa(spec.respStatus))...))
+				require.NoError(t, err)
 			})
 
 			// when
-			middleware := NewRetryMiddleware(spec.maxRetries, testHandler)
+			middleware := NewRetryMiddleware(spec.maxRetries, testBackend)
 			middleware.ServeHTTP(respRecorder, req)
 
 			// then
 			resp := respRecorder.Result()
 			require.Equal(t, spec.respStatus, resp.StatusCode)
 			assert.Equal(t, spec.expRetries, counter-1)
-
+			// and headers matches
+			assert.Equal(t, spec.headers, resp.Header)
+			// and body matches
+			bodyRead, err := io.ReadAll(resp.Body)
+			require.NoError(t, err)
+			require.NoError(t, resp.Body.Close())
+			assert.Equal(t, myBody+strconv.Itoa(spec.respStatus), string(bodyRead))
+			// and prometheus metric updated
 			assert.Equal(t, spec.expRetries, int(counterValue(t, totalRetries)-counterBefore))
 		})
 	}
 }
 
+func TestWriteDelegatorReadFrom(t *testing.T) {
+	const myTestContent = `my body content`
+
+	rec := &testResponseWriter{ResponseRecorder: httptest.NewRecorder()}
+	// scenario: discard on error disabled
+	d := newResponseWriterDelegator(rec, func(int) bool { return true }, false)
+	// when
+	n, err := d.(io.ReaderFrom).ReadFrom(strings.NewReader(myTestContent))
+	// then the content is written
+	require.NoError(t, err)
+	assert.Equal(t, len(myTestContent), int(n))
+	assert.Equal(t, myTestContent, rec.Body.String())
+	assert.Equal(t, http.StatusOK, d.capturedStatusCode())
+
+	// scenario: discard on error enabled
+	rec = &testResponseWriter{ResponseRecorder: httptest.NewRecorder()}
+	// with discard on error disabled
+	d = newResponseWriterDelegator(rec, func(int) bool { return true }, true)
+	// when
+	n, err = d.(io.ReaderFrom).ReadFrom(strings.NewReader(myTestContent))
+	// then the content is not written
+	require.NoError(t, err)
+	assert.Equal(t, len(myTestContent), int(n))
+	assert.Equal(t, "", rec.Body.String())
+	assert.Equal(t, http.StatusOK, d.capturedStatusCode())
+
+	// scenario: not implementing io.ReaderFrom
+	d = newResponseWriterDelegator(httptest.NewRecorder(), func(int) bool { return true }, false)
+	_, ok := d.(io.ReaderFrom)
+	require.False(t, ok)
+}
+
+func TestLazyBodyCapturer(t *testing.T) {
+	const myTestContent = "my-test-content"
+	c := newLazyBodyCapturer(io.NopCloser(strings.NewReader(myTestContent)))
+	var buf bytes.Buffer
+	n, err := c.(io.WriterTo).WriteTo(&buf)
+	require.NoError(t, err)
+	assert.Len(t, myTestContent, int(n))
+	assert.Equal(t, myTestContent, buf.String())
+	// when captured
+	c.Capture()
+	// then data is buffered for second read
+	buf.Reset()
+	n, err = c.(io.WriterTo).WriteTo(&buf)
+	require.NoError(t, err)
+	assert.Equal(t, len(myTestContent), 15)
+	assert.Equal(t, myTestContent, buf.String())
+
+	// scenario: source reader does not implement WriteTo
+	c = newLazyBodyCapturer(testReader{strings.NewReader(myTestContent)})
+	// then instance also does not implement it
+	_, ok := c.(io.WriterTo)
+	require.False(t, ok)
+}
+
 func counterValue(t *testing.T, counter prometheus.Counter) float64 {
 	registry := prometheus.NewPedanticRegistry()
 	registry.MustRegister(counter)
@@ -96,3 +183,23 @@ func counterValue(t *testing.T, counter prometheus.Counter) float64 {
 	require.Len(t, gathered[0].Metric, 1)
 	return gathered[0].Metric[0].GetCounter().GetValue()
 }
+
+type testResponseWriter struct {
+	*httptest.ResponseRecorder
+}
+
+func (r *testResponseWriter) ReadFrom(re io.Reader) (int64, error) {
+	return r.ResponseRecorder.Body.ReadFrom(re)
+}
+
+type testReader struct {
+	r io.Reader
+}
+
+func (t testReader) Read(p []byte) (n int, err error) {
+	return t.r.Read(p)
+}
+
+func (t testReader) Close() error {
+	return nil
+}
diff --git a/tests/integration/integration_test.go b/tests/integration/integration_test.go
index 3ba2ce18..003f976b 100644
--- a/tests/integration/integration_test.go
+++ b/tests/integration/integration_test.go
@@ -229,6 +229,7 @@ func sendRequest(t *testing.T, wg *sync.WaitGroup, modelName string, expCode int
 	bodyRespChan := make(chan string, 1)
 	go func() {
 		defer wg.Done()
+		defer close(bodyRespChan)
 
 		body := []byte(fmt.Sprintf(`{"model": %q}`, modelName))
 		req, err := http.NewRequest(http.MethodPost, testServer.URL, bytes.NewReader(body))
@@ -241,7 +242,6 @@ func sendRequest(t *testing.T, wg *sync.WaitGroup, modelName string, expCode int
 		_ = res.Body.Close()
 		require.NoError(t, err)
 		bodyRespChan <- string(got)
-		close(bodyRespChan)
 	}()
 	return bodyRespChan
 }