http: add Transport.MaxIdleConnsPerHost

R=rsc
CC=golang-dev
https://golang.org/cl/4280079

2 files changed, 73 insertions, 3 deletions

diff --git a/src/pkg/http/transport.go b/src/pkg/http/transport.go
index ed7843bc71..7f85c8c281 100644
--- a/src/pkg/http/transport.go
+++ b/src/pkg/http/transport.go
@@ -24,6 +24,10 @@ import (
 // environment variables.
 var DefaultTransport RoundTripper = &Transport{}
 
+// DefaultMaxIdleConnsPerHost is the default value of Transport's
+// MaxIdleConnsPerHost.
+const DefaultMaxIdleConnsPerHost = 2
+
 // Transport is an implementation of RoundTripper that supports http,
 // https, and http proxies (for either http or https with CONNECT).
 // Transport can also cache connections for future re-use.
@@ -31,11 +35,17 @@ type Transport struct {
 	lk       sync.Mutex
 	idleConn map[string][]*persistConn
 
-	// TODO: tunables on max cached connections (total, per-server), duration
+	// TODO: tunable on global max cached connections
+	// TODO: tunable on timeout on cached connections
 	// TODO: optional pipelining
 
 	IgnoreEnvironment bool // don't look at environment variables for proxy configuration
 	DisableKeepAlives bool
+
+	// MaxIdleConnsPerHost, if non-zero, controls the maximum idle
+	// (keep-alive) to keep to keep per-host.  If zero,
+	// DefaultMaxIdleConnsPerHost is used.
+	MaxIdleConnsPerHost int
 }
 
 // RoundTrip implements the RoundTripper interface.
@@ -147,7 +157,7 @@ func (cm *connectMethod) proxyAuth() string {
 func (t *Transport) putIdleConn(pconn *persistConn) {
 	t.lk.Lock()
 	defer t.lk.Unlock()
-	if t.DisableKeepAlives {
+	if t.DisableKeepAlives || t.MaxIdleConnsPerHost < 0 {
 		pconn.close()
 		return
 	}
@@ -155,6 +165,14 @@ func (t *Transport) putIdleConn(pconn *persistConn) {
 		return
 	}
 	key := pconn.cacheKey
+	max := t.MaxIdleConnsPerHost
+	if max == 0 {
+		max = DefaultMaxIdleConnsPerHost
+	}
+	if len(t.idleConn[key]) >= max {
+		pconn.close()
+		return
+	}
 	t.idleConn[key] = append(t.idleConn[key], pconn)
 }
 
diff --git a/src/pkg/http/transport_test.go b/src/pkg/http/transport_test.go
index 5c3e1cdb58..69a17df856 100644
--- a/src/pkg/http/transport_test.go
+++ b/src/pkg/http/transport_test.go
@@ -164,7 +164,7 @@ func TestTransportIdleCacheKeys(t *testing.T) {
 	}
 
 	if e := "|http|" + ts.Listener.Addr().String(); keys[0] != e {
-		t.Logf("Expected idle cache key %q; got %q", e, keys[0])
+		t.Errorf("Expected idle cache key %q; got %q", e, keys[0])
 	}
 
 	tr.CloseIdleConnections()
@@ -173,6 +173,58 @@ func TestTransportIdleCacheKeys(t *testing.T) {
 	}
 }
 
+func TestTransportMaxPerHostIdleConns(t *testing.T) {
+	ch := make(chan string)
+	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
+		fmt.Fprintf(w, "%s", <-ch)
+	}))
+	defer ts.Close()
+	maxIdleConns := 2
+	tr := &Transport{DisableKeepAlives: false, MaxIdleConnsPerHost: maxIdleConns}
+	c := &Client{Transport: tr}
+
+	// Start 3 outstanding requests (will hang until we write to
+	// ch)
+	donech := make(chan bool)
+	doReq := func() {
+		c.Get(ts.URL)
+		donech <- true
+	}
+	go doReq()
+	go doReq()
+	go doReq()
+
+	if e, g := 0, len(tr.IdleConnKeysForTesting()); e != g {
+		t.Fatalf("Before writes, expected %d idle conn cache keys; got %d", e, g)
+	}
+
+	ch <- "res1"
+	<-donech
+	keys := tr.IdleConnKeysForTesting()
+	if e, g := 1, len(keys); e != g {
+		t.Fatalf("after first response, expected %d idle conn cache keys; got %d", e, g)
+	}
+	cacheKey := "|http|" + ts.Listener.Addr().String()
+	if keys[0] != cacheKey {
+		t.Fatalf("Expected idle cache key %q; got %q", cacheKey, keys[0])
+	}
+	if e, g := 1, tr.IdleConnCountForTesting(cacheKey); e != g {
+		t.Errorf("after first response, expected %d idle conns; got %d", e, g)
+	}
+
+	ch <- "res2"
+	<-donech
+	if e, g := 2, tr.IdleConnCountForTesting(cacheKey); e != g {
+		t.Errorf("after second response, expected %d idle conns; got %d", e, g)
+	}
+
+	ch <- "res3"
+	<-donech
+	if e, g := maxIdleConns, tr.IdleConnCountForTesting(cacheKey); e != g {
+		t.Errorf("after third response, still expected %d idle conns; got %d", e, g)
+	}
+}
+
 func TestTransportServerClosingUnexpectedly(t *testing.T) {
 	ts := httptest.NewServer(hostPortHandler)
 	defer ts.Close()