perf: reuse reverse proxy copy buffers

context_benchmark_test.go

@@ -23,7 +23,7 @@ func TestContextResetKeepsKeysNilUntilSet(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to build request: %v", err)
 	}
-	c.reset(c.Writer, req)
+	c.reset(UnwrapResponseWriter(c.Writer), req)
 
 	if c.Keys != nil {
 		t.Fatalf("expected reset to clear Keys without allocating a new map")
@@ -47,6 +47,7 @@ func TestContextResetKeepsKeysNilUntilSet(t *testing.T) {
 func BenchmarkContextReset(b *testing.B) {
 	b.Run("NoKeysUse", func(b *testing.B) {
 		c, _ := CreateTestContext(nil)
+		rawWriter := UnwrapResponseWriter(c.Writer)
 		req, err := http.NewRequest(http.MethodGet, "/", nil)
 		if err != nil {
 			b.Fatalf("failed to build request: %v", err)
@@ -56,12 +57,13 @@ func BenchmarkContextReset(b *testing.B) {
 		b.ResetTimer()
 
 		for i := 0; i < b.N; i++ {
-			c.reset(c.Writer, req)
+			c.reset(rawWriter, req)
 		}
 	})
 
 	b.Run("WithKeysUse", func(b *testing.B) {
 		c, _ := CreateTestContext(nil)
+		rawWriter := UnwrapResponseWriter(c.Writer)
 		req, err := http.NewRequest(http.MethodGet, "/", nil)
 		if err != nil {
 			b.Fatalf("failed to build request: %v", err)
@@ -71,8 +73,9 @@ func BenchmarkContextReset(b *testing.B) {
 		b.ResetTimer()
 
 		for i := 0; i < b.N; i++ {
-			c.reset(c.Writer, req)
+			c.reset(rawWriter, req)
 			c.Set("request-id", i)
 		}
 	})
+
 }

ecw_benchmark_test.go

@@ -0,0 +1,59 @@
+package touka
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestErrorCapturingResponseWriterResetClearsHeaderSnapshot(t *testing.T) {
+	c, _ := CreateTestContext(nil)
+	ecw := AcquireErrorCapturingResponseWriter(c)
+	defer ReleaseErrorCapturingResponseWriter(ecw)
+
+	ecw.capturedErrorSignal = true
+	ecw.Header().Set("Content-Type", "text/plain")
+	ecw.Header().Add("X-Test", "one")
+
+	req, err := http.NewRequest(http.MethodGet, "/", nil)
+	if err != nil {
+		t.Fatalf("failed to build request: %v", err)
+	}
+
+	ecw.reset(httptest.NewRecorder(), req, c, c.engine.errorHandle.handler)
+
+	if len(ecw.headerSnapshot) != 0 {
+		t.Fatalf("expected header snapshot to be empty after reset, got %#v", ecw.headerSnapshot)
+	}
+}
+
+func BenchmarkErrorCapturingResponseWriterReset(b *testing.B) {
+	c, _ := CreateTestContext(nil)
+	ecw := AcquireErrorCapturingResponseWriter(c)
+	defer ReleaseErrorCapturingResponseWriter(ecw)
+
+	rawWriter := httptest.NewRecorder()
+	req, err := http.NewRequest(http.MethodGet, "/", nil)
+	if err != nil {
+		b.Fatalf("failed to build request: %v", err)
+	}
+
+	keys := make([]string, 16)
+	for i := range keys {
+		keys[i] = http.CanonicalHeaderKey("X-Test-" + string(rune('A'+i)))
+	}
+	values := []string{"one", "two", "three"}
+	for _, key := range keys {
+		ecw.headerSnapshot[key] = values
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		ecw.reset(rawWriter, req, c, c.engine.errorHandle.handler)
+		for _, key := range keys {
+			ecw.headerSnapshot[key] = values
+		}
+	}
+}

reverseproxy.go

@@ -84,6 +84,13 @@ type reverseProxyHandler struct {
 	roundRobin  atomic.Uint64
 }
 
+var reverseProxyCopyBufferPool = sync.Pool{
+	New: func() any {
+		buf := make([]byte, 32*1024)
+		return &buf
+	},
+}
+
 type reverseProxyStatusError struct {
 	status int
 	err    error
@@ -1153,6 +1160,10 @@ func (p *reverseProxyHandler) copyResponse(dst ResponseWriter, src io.Reader, fl
 	if p.config.BufferPool != nil {
 		buf = p.config.BufferPool.Get()
 		defer p.config.BufferPool.Put(buf)
+	} else {
+		bufp := reverseProxyCopyBufferPool.Get().(*[]byte)
+		buf = *bufp
+		defer reverseProxyCopyBufferPool.Put(bufp)
 	}
 	_, err := p.copyBuffer(writer, src, buf)
 	return err

reverseproxy_benchmark_test.go

@@ -0,0 +1,206 @@
+package touka
+
+import (
+	"bufio"
+	"bytes"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+)
+
+type benchmarkReadSeeker struct {
+	data []byte
+	off  int
+}
+
+func (r *benchmarkReadSeeker) Read(p []byte) (int, error) {
+	if r.off >= len(r.data) {
+		return 0, io.EOF
+	}
+	n := copy(p, r.data[r.off:])
+	r.off += n
+	return n, nil
+}
+
+func (r *benchmarkReadSeeker) Reset() {
+	r.off = 0
+}
+
+type benchmarkResponseWriter struct {
+	header http.Header
+	status int
+	size   int
+}
+
+func newBenchmarkResponseWriter() *benchmarkResponseWriter {
+	return &benchmarkResponseWriter{header: make(http.Header)}
+}
+
+func (w *benchmarkResponseWriter) Header() http.Header {
+	return w.header
+}
+
+func (w *benchmarkResponseWriter) WriteHeader(statusCode int) {
+	if w.status == 0 {
+		w.status = statusCode
+	}
+}
+
+func (w *benchmarkResponseWriter) Write(p []byte) (int, error) {
+	if w.status == 0 {
+		w.status = http.StatusOK
+	}
+	w.size += len(p)
+	return len(p), nil
+}
+
+func (w *benchmarkResponseWriter) Flush() {}
+
+func (w *benchmarkResponseWriter) Status() int {
+	return w.status
+}
+
+func (w *benchmarkResponseWriter) Size() int {
+	return w.size
+}
+
+func (w *benchmarkResponseWriter) Written() bool {
+	return w.status != 0
+}
+
+func (w *benchmarkResponseWriter) IsHijacked() bool {
+	return false
+}
+
+func (w *benchmarkResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, http.ErrNotSupported
+}
+
+func (w *benchmarkResponseWriter) reset() {
+	clear(w.header)
+	w.status = 0
+	w.size = 0
+}
+
+var benchmarkReverseProxySink int
+
+func BenchmarkReverseProxyCopyResponse(b *testing.B) {
+	body := bytes.Repeat([]byte("0123456789abcdef"), 4096)
+	proxy := newReverseProxyHandler(ReverseProxyConfig{})
+	dst := newBenchmarkResponseWriter()
+	src := &benchmarkReadSeeker{data: body}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		dst.reset()
+		src.Reset()
+		if err := proxy.copyResponse(dst, src, 0); err != nil {
+			b.Fatalf("copyResponse failed: %v", err)
+		}
+	}
+
+	benchmarkReverseProxySink = dst.Size()
+}
+
+func BenchmarkReverseProxyAvailableUpstreams(b *testing.B) {
+	proxy := &reverseProxyHandler{
+		upstreams: []*reverseProxyUpstream{
+			{key: "a"},
+			{key: "b"},
+			{key: "c"},
+			{key: "d"},
+		},
+		config: ReverseProxyConfig{
+			PassiveHealth: ReverseProxyPassiveHealthConfig{
+				FailDuration: time.Minute,
+				MaxFails:     3,
+			},
+		},
+	}
+
+	now := time.Now()
+	proxy.upstreams[0].failures = []time.Time{now.Add(-30 * time.Second)}
+	proxy.upstreams[1].failures = []time.Time{now.Add(-20 * time.Second), now.Add(-10 * time.Second)}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		benchmarkReverseProxySink = len(proxy.availableUpstreams(now, nil))
+	}
+}
+
+func TestReverseProxyCopyResponseWithoutBufferPool(t *testing.T) {
+	proxy := newReverseProxyHandler(ReverseProxyConfig{})
+	dst := newBenchmarkResponseWriter()
+	src := bytes.NewBufferString("hello, reverse proxy")
+
+	if err := proxy.copyResponse(dst, src, 0); err != nil {
+		t.Fatalf("copyResponse failed: %v", err)
+	}
+
+	if got, want := dst.Size(), len("hello, reverse proxy"); got != want {
+		t.Fatalf("expected %d bytes copied, got %d", want, got)
+	}
+}
+
+type fixedLenBufferPool struct {
+	buf []byte
+}
+
+func (p *fixedLenBufferPool) Get() []byte {
+	return p.buf
+}
+
+func (p *fixedLenBufferPool) Put(buf []byte) {
+	p.buf = buf
+}
+
+type recordingReader struct {
+	chunk int
+	reads []int
+	left  int
+}
+
+func (r *recordingReader) Read(p []byte) (int, error) {
+	if r.left == 0 {
+		return 0, io.EOF
+	}
+	n := min(r.chunk, len(p), r.left)
+	if n == 0 {
+		return 0, errors.New("reader received zero-length buffer")
+	}
+	for i := 0; i < n; i++ {
+		p[i] = 'x'
+	}
+	r.left -= n
+	r.reads = append(r.reads, len(p))
+	return n, nil
+}
+
+func TestReverseProxyCopyResponseRespectsCustomBufferLength(t *testing.T) {
+	pool := &fixedLenBufferPool{buf: make([]byte, 8, 32*1024)}
+	proxy := newReverseProxyHandler(ReverseProxyConfig{BufferPool: pool})
+	dst := newBenchmarkResponseWriter()
+	src := &recordingReader{chunk: 8, left: 24}
+
+	if err := proxy.copyResponse(dst, src, 0); err != nil {
+		t.Fatalf("copyResponse failed: %v", err)
+	}
+
+	if len(src.reads) == 0 {
+		t.Fatal("expected reader to be used")
+	}
+	for _, size := range src.reads {
+		if size != 8 {
+			t.Fatalf("expected custom buffer length 8 to be preserved, got read size %d", size)
+		}
+	}
+}
+
+var _ io.Writer = (*benchmarkResponseWriter)(nil)