From 916682ea367a290da9392ef38016af5b8fd6a3b7 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Fri, 28 Feb 2014 07:40:25 -0800
Subject: [PATCH] net/http: add Server.SetKeepAlivesEnabled

Part of graceful shutdown.

Update #4674

LGTM=adg, josharian
R=adg, josharian, r
CC=golang-codereviews
https://golang.org/cl/69670043
---
 src/pkg/net/http/serve_test.go | 16 ++++++++++++++++
 src/pkg/net/http/server.go     | 27 ++++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/pkg/net/http/serve_test.go b/src/pkg/net/http/serve_test.go
index 77d2b97a7d..c8ca03e07d 100644
--- a/src/pkg/net/http/serve_test.go
+++ b/src/pkg/net/http/serve_test.go
@@ -2358,6 +2358,22 @@ func mustGet(t *testing.T, url string, headers ...string) {
 	}
 }
 
+func TestServerKeepAlivesEnabled(t *testing.T) {
+	defer afterTest(t)
+	ts := httptest.NewUnstartedServer(HandlerFunc(func(w ResponseWriter, r *Request) {}))
+	ts.Config.SetKeepAlivesEnabled(false)
+	ts.Start()
+	defer ts.Close()
+	res, err := Get(ts.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer res.Body.Close()
+	if !res.Close {
+		t.Errorf("Body.Close == false; want true")
+	}
+}
+
 func BenchmarkClientServer(b *testing.B) {
 	b.ReportAllocs()
 	b.StopTimer()
diff --git a/src/pkg/net/http/server.go b/src/pkg/net/http/server.go
index 8ca48ab3ce..ba3a530adc 100644
--- a/src/pkg/net/http/server.go
+++ b/src/pkg/net/http/server.go
@@ -22,6 +22,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -703,6 +704,7 @@ func (cw *chunkWriter) writeHeader(p []byte) {
 	cw.wroteHeader = true
 
 	w := cw.res
+	keepAlivesEnabled := w.conn.server.doKeepAlives()
 	isHEAD := w.req.Method == "HEAD"
 
 	// header is written out to w.conn.buf below. Depending on the
@@ -750,7 +752,7 @@ func (cw *chunkWriter) writeHeader(p []byte) {
 
 	// If this was an HTTP/1.0 request with keep-alive and we sent a
 	// Content-Length back, we can make this a keep-alive response ...
-	if w.req.wantsHttp10KeepAlive() {
+	if w.req.wantsHttp10KeepAlive() && keepAlivesEnabled {
 		sentLength := header.get("Content-Length") != ""
 		if sentLength && header.get("Connection") == "keep-alive" {
 			w.closeAfterReply = false
@@ -769,7 +771,7 @@ func (cw *chunkWriter) writeHeader(p []byte) {
 		w.closeAfterReply = true
 	}
 
-	if header.get("Connection") == "close" {
+	if header.get("Connection") == "close" || !keepAlivesEnabled {
 		w.closeAfterReply = true
 	}
 
@@ -851,7 +853,7 @@ func (cw *chunkWriter) writeHeader(p []byte) {
 		return
 	}
 
-	if w.closeAfterReply && !hasToken(cw.header.get("Connection"), "close") {
+	if w.closeAfterReply && (!keepAlivesEnabled || !hasToken(cw.header.get("Connection"), "close")) {
 		delHeader("Connection")
 		if w.req.ProtoAtLeast(1, 1) {
 			setHeader.connection = "close"
@@ -1579,6 +1581,7 @@ func Serve(l net.Listener, handler Handler) error {
 }
 
 // A Server defines parameters for running an HTTP server.
+// The zero value for Server is a valid configuration.
 type Server struct {
 	Addr           string        // TCP address to listen on, ":http" if empty
 	Handler        Handler       // handler to invoke, http.DefaultServeMux if nil
@@ -1600,6 +1603,8 @@ type Server struct {
 	// called when a client connection changes state. See the
 	// ConnState type and associated constants for details.
 	ConnState func(net.Conn, ConnState)
+
+	disableKeepAlives int32 // accessed atomically.
 }
 
 // A ConnState represents the state of a client connection to a server.
@@ -1714,6 +1719,22 @@ func (srv *Server) Serve(l net.Listener) error {
 	}
 }
 
+func (s *Server) doKeepAlives() bool {
+	return atomic.LoadInt32(&s.disableKeepAlives) == 0
+}
+
+// SetKeepAlivesEnabled controls whether HTTP keep-alives are enabled.
+// By default, keep-alives are always enabled. Only very
+// resource-constrained environments or servers in the process of
+// shutting down should disable them.
+func (s *Server) SetKeepAlivesEnabled(v bool) {
+	if v {
+		atomic.StoreInt32(&s.disableKeepAlives, 0)
+	} else {
+		atomic.StoreInt32(&s.disableKeepAlives, 1)
+	}
+}
+
 // ListenAndServe listens on the TCP network address addr
 // and then calls Serve with handler to handle requests
 // on incoming connections.  Handler is typically nil,
-- 
2.48.1