From 492a62e945555bbf94a6f9dd6d430f712738c5e0 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Mon, 27 Apr 2015 18:10:20 -0700
Subject: [PATCH] net/http/httputil: add hook for managing io.Copy buffers per
 request

Adds ReverseProxy.BufferPool for users with sensitive allocation
requirements. Permits avoiding 32 KB of io.Copy garbage per request.

Fixes #10277

Change-Id: I5dfd58fa70a363ead4be56405e507df90d871719
Reviewed-on: https://go-review.googlesource.com/9399
Reviewed-by: Keith Randall <khr@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
---
 src/net/http/httputil/reverseproxy.go      | 21 ++++++-
 src/net/http/httputil/reverseproxy_test.go | 68 ++++++++++++++++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/src/net/http/httputil/reverseproxy.go b/src/net/http/httputil/reverseproxy.go
index 95a99ddb9dd0b..4dba352a4fa51 100644
--- a/src/net/http/httputil/reverseproxy.go
+++ b/src/net/http/httputil/reverseproxy.go
@@ -46,6 +46,18 @@ type ReverseProxy struct {
 	// If nil, logging goes to os.Stderr via the log package's
 	// standard logger.
 	ErrorLog *log.Logger
+
+	// BufferPool optionally specifies a buffer pool to
+	// get byte slices for use by io.CopyBuffer when
+	// copying HTTP response bodies.
+	BufferPool BufferPool
+}
+
+// A BufferPool is an interface for getting and returning temporary
+// byte slices for use by io.CopyBuffer.
+type BufferPool interface {
+	Get() []byte
+	Put([]byte)
 }
 
 func singleJoiningSlash(a, b string) string {
@@ -245,7 +257,14 @@ func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
 		}
 	}
 
-	io.Copy(dst, src)
+	var buf []byte
+	if p.BufferPool != nil {
+		buf = p.BufferPool.Get()
+	}
+	io.CopyBuffer(dst, src, buf)
+	if p.BufferPool != nil {
+		p.BufferPool.Put(buf)
+	}
 }
 
 func (p *ReverseProxy) logf(format string, args ...interface{}) {
diff --git a/src/net/http/httputil/reverseproxy_test.go b/src/net/http/httputil/reverseproxy_test.go
index 1d309614e2b64..5f6fc56e07a5b 100644
--- a/src/net/http/httputil/reverseproxy_test.go
+++ b/src/net/http/httputil/reverseproxy_test.go
@@ -8,13 +8,16 @@ package httputil
 
 import (
 	"bufio"
+	"io"
 	"io/ioutil"
 	"log"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"reflect"
+	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 )
@@ -316,3 +319,68 @@ func TestNilBody(t *testing.T) {
 		t.Errorf("Got %q; want %q", slurp, "hi")
 	}
 }
+
+type bufferPool struct {
+	get func() []byte
+	put func([]byte)
+}
+
+func (bp bufferPool) Get() []byte  { return bp.get() }
+func (bp bufferPool) Put(v []byte) { bp.put(v) }
+
+func TestReverseProxyGetPutBuffer(t *testing.T) {
+	const msg = "hi"
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		io.WriteString(w, msg)
+	}))
+	defer backend.Close()
+
+	backendURL, err := url.Parse(backend.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var (
+		mu  sync.Mutex
+		log []string
+	)
+	addLog := func(event string) {
+		mu.Lock()
+		defer mu.Unlock()
+		log = append(log, event)
+	}
+	rp := NewSingleHostReverseProxy(backendURL)
+	const size = 1234
+	rp.BufferPool = bufferPool{
+		get: func() []byte {
+			addLog("getBuf")
+			return make([]byte, size)
+		},
+		put: func(p []byte) {
+			addLog("putBuf-" + strconv.Itoa(len(p)))
+		},
+	}
+	frontend := httptest.NewServer(rp)
+	defer frontend.Close()
+
+	req, _ := http.NewRequest("GET", frontend.URL, nil)
+	req.Close = true
+	res, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("Get: %v", err)
+	}
+	slurp, err := ioutil.ReadAll(res.Body)
+	res.Body.Close()
+	if err != nil {
+		t.Fatalf("reading body: %v", err)
+	}
+	if string(slurp) != msg {
+		t.Errorf("msg = %q; want %q", slurp, msg)
+	}
+	wantLog := []string{"getBuf", "putBuf-" + strconv.Itoa(size)}
+	mu.Lock()
+	defer mu.Unlock()
+	if !reflect.DeepEqual(log, wantLog) {
+		t.Errorf("Log events = %q; want %q", log, wantLog)
+	}
+}