Use net/http.Server.SetKeepAlivesEnabled in Go 1.3

This feature can be used in place of the pile of hacks in middleware.go,
and doesn't involve awkwardly shimming out a http.ResponseWriter. Sounds
like a win-win!

example/main.go

@@ -13,6 +13,7 @@ import (
 	"net/http"
 	"regexp"
 	"strconv"
+	"time"
 
 	"github.com/goji/param"
 	"github.com/zenazn/goji"
@@ -58,6 +59,9 @@ func main() {
 	// Use a custom 404 handler
 	goji.NotFound(NotFound)
 
+	// Sometimes requests take a long time.
+	goji.Get("/waitforit", WaitForIt)
+
 	// Call Serve() at the bottom of your main() function, and it'll take
 	// care of everything else for you, including binding to a socket (with
 	// automatic support for systemd and Einhorn) and supporting graceful
@@ -135,6 +139,20 @@ func GetGreet(c web.C, w http.ResponseWriter, r *http.Request) {
 	greet.Write(w)
 }
 
+// WaitForIt is a particularly slow handler (GET "/waitforit"). Try loading this
+// endpoint and initiating a graceful shutdown (Ctrl-C) or Einhorn reload. The
+// old server will stop accepting new connections and will attempt to kill
+// outstanding idle (keep-alive) connections, but will patiently stick around
+// for this endpoint to finish. How kind of it!
+func WaitForIt(w http.ResponseWriter, r *http.Request) {
+	io.WriteString(w, "This is going to be legend... (wait for it)\n")
+	if fl, ok := w.(http.Flusher); ok {
+		fl.Flush()
+	}
+	time.Sleep(15 * time.Second)
+	io.WriteString(w, "...dary! Legendary!\n")
+}
+
 // AdminRoot is root (GET "/admin/root"). Much secret. Very administrate. Wow.
 func AdminRoot(w http.ResponseWriter, r *http.Request) {
 	io.WriteString(w, "Gritter\n======\n\nSuper secret admin page!\n")

graceful/graceful.go

@@ -22,7 +22,6 @@ import (
 	"crypto/tls"
 	"net"
 	"net/http"
-	"time"
 )
 
 /*
@@ -41,38 +40,6 @@ ListenAndServe and ListenAndServeTLS here more-or-less verbatim. "Oh well!"
 // implementations of its methods.
 type Server http.Server
 
-func (srv *Server) Serve(l net.Listener) (err error) {
-	go func() {
-		<-kill
-		l.Close()
-	}()
-	l = WrapListener(l)
-
-	// Spawn a shadow http.Server to do the actual servering. We do this
-	// because we need to sketch on some of the parameters you passed in,
-	// and it's nice to keep our sketching to ourselves.
-	shadow := *(*http.Server)(srv)
-
-	if shadow.ReadTimeout == 0 {
-		shadow.ReadTimeout = forever
-	}
-	shadow.Handler = Middleware(shadow.Handler)
-
-	err = shadow.Serve(l)
-
-	// We expect an error when we close the listener, so we indiscriminately
-	// swallow Serve errors when we're in a shutdown state.
-	select {
-	case <-kill:
-		return nil
-	default:
-		return err
-	}
-}
-
-// About 200 years, also known as "forever"
-const forever time.Duration = 200 * 365 * 24 * time.Hour
-
 func (srv *Server) ListenAndServe() error {
 	addr := srv.Addr
 	if addr == "" {

graceful/middleware.go

@@ -1,3 +1,5 @@
+// +build !go1.3
+
 package graceful
 
 import (

graceful/middleware13.go

@@ -0,0 +1,12 @@
+// +build go1.3
+
+package graceful
+
+import "net/http"
+
+// Middleware is a stub that does nothing. When used with versions of Go before
+// Go 1.3, it provides functionality similar to net/http.Server's
+// SetKeepAlivesEnabled.
+func Middleware(h http.Handler) http.Handler {
+	return h
+}

graceful/middleware_test.go

@@ -1,3 +1,5 @@
+// +build !go1.3
+
 package graceful
 
 import (

graceful/net.go

@@ -162,6 +162,9 @@ func (c *conn) Read(b []byte) (n int, err error) {
 
 		if c.state == csWaiting {
 			c.state = csWorking
+		} else if c.state == csDead {
+			n = 0
+			err = io.EOF
 		}
 	}()
 

graceful/serve.go

@@ -0,0 +1,41 @@
+// +build !go1.3
+
+package graceful
+
+import (
+	"net"
+	"net/http"
+	"time"
+)
+
+// About 200 years, also known as "forever"
+const forever time.Duration = 200 * 365 * 24 * time.Hour
+
+func (srv *Server) Serve(l net.Listener) error {
+	go func() {
+		<-kill
+		l.Close()
+	}()
+	l = WrapListener(l)
+
+	// Spawn a shadow http.Server to do the actual servering. We do this
+	// because we need to sketch on some of the parameters you passed in,
+	// and it's nice to keep our sketching to ourselves.
+	shadow := *(*http.Server)(srv)
+
+	if shadow.ReadTimeout == 0 {
+		shadow.ReadTimeout = forever
+	}
+	shadow.Handler = Middleware(shadow.Handler)
+
+	err := shadow.Serve(l)
+
+	// We expect an error when we close the listener, so we indiscriminately
+	// swallow Serve errors when we're in a shutdown state.
+	select {
+	case <-kill:
+		return nil
+	default:
+		return err
+	}
+}

graceful/serve13.go

@@ -0,0 +1,42 @@
+// +build go1.3
+
+package graceful
+
+import (
+	"net"
+	"net/http"
+	"time"
+)
+
+// About 200 years, also known as "forever"
+const forever time.Duration = 200 * 365 * 24 * time.Hour
+
+func (srv *Server) Serve(l net.Listener) error {
+	l = WrapListener(l)
+
+	// Spawn a shadow http.Server to do the actual servering. We do this
+	// because we need to sketch on some of the parameters you passed in,
+	// and it's nice to keep our sketching to ourselves.
+	shadow := *(*http.Server)(srv)
+
+	if shadow.ReadTimeout == 0 {
+		shadow.ReadTimeout = forever
+	}
+
+	go func() {
+		<-kill
+		shadow.SetKeepAlivesEnabled(false)
+		l.Close()
+	}()
+
+	err := shadow.Serve(l)
+
+	// We expect an error when we close the listener, so we indiscriminately
+	// swallow Serve errors when we're in a shutdown state.
+	select {
+	case <-kill:
+		return nil
+	default:
+		return err
+	}
+}