net/http: Add Server.RegisterOnShutdown

This will be used to allow http2 servers to register a shutdown function
so that net/http.Server.Shutdown will work when the http2 server is
configured via a manual call to http2.ConfigureServer. Currently, Shutdown
only works when the http2 server is configured automatically by the
net/http package.

Updates #20302
Updates #18471

Change-Id: Ifc2b5f3126126a106b49ea4a7e999279852b9cc9
Reviewed-on: https://go-review.googlesource.com/44003
Run-TryBot: Tom Bergan <tombergan@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

diff --git a/src/net/http/serve_test.go b/src/net/http/serve_test.go
index 1d541a8e46f0664b2f11f84486dc606f9123a5ee..80fcc8c407125ec6b6ef3b3ad3c0cce5b2acf229 100644 (file)
--- a/src/net/http/serve_test.go
+++ b/src/net/http/serve_test.go
@@ -5232,7 +5232,8 @@ func testServerShutdown(t *testing.T, h2 bool) {
        defer afterTest(t)
        var doShutdown func() // set later
        var shutdownRes = make(chan error, 1)
-       cst := newClientServerTest(t, h2, HandlerFunc(func(w ResponseWriter, r *Request) {
+       var gotOnShutdown = make(chan struct{}, 1)
+       handler := HandlerFunc(func(w ResponseWriter, r *Request) {
                go doShutdown()
                // Shutdown is graceful, so it should not interrupt
                // this in-flight response. Add a tiny sleep here to
@@ -5240,7 +5241,10 @@ func testServerShutdown(t *testing.T, h2 bool) {
                // bugs.
                time.Sleep(20 * time.Millisecond)
                io.WriteString(w, r.RemoteAddr)
-       }))
+       })
+       cst := newClientServerTest(t, h2, handler, func(srv *httptest.Server) {
+               srv.Config.RegisterOnShutdown(func() { gotOnShutdown <- struct{}{} })
+       })
        defer cst.close()
 
        doShutdown = func() {
@@ -5251,6 +5255,11 @@ func testServerShutdown(t *testing.T, h2 bool) {
        if err := <-shutdownRes; err != nil {
                t.Fatalf("Shutdown: %v", err)
        }
+       select {
+       case <-gotOnShutdown:
+       case <-time.After(5 * time.Second):
+               t.Errorf("onShutdown callback not called, RegisterOnShutdown broken?")
+       }
 
        res, err := cst.c.Get(cst.ts.URL)
        if err == nil {

diff --git a/src/net/http/server.go b/src/net/http/server.go
index 71f46a74f97720de8f803afe5afe08d8b43c1629..b60bd2481e4248c6d97ae4ef5892542023a6b538 100644 (file)
--- a/src/net/http/server.go
+++ b/src/net/http/server.go
@@ -2395,6 +2395,7 @@ type Server struct {
        listeners  map[net.Listener]struct{}
        activeConn map[*conn]struct{}
        doneChan   chan struct{}
+       onShutdown []func()
 }
 
 func (s *Server) getDoneChan() <-chan struct{} {
@@ -2475,6 +2476,9 @@ func (srv *Server) Shutdown(ctx context.Context) error {
        srv.mu.Lock()
        lnerr := srv.closeListenersLocked()
        srv.closeDoneChanLocked()
+       for _, f := range srv.onShutdown {
+               go f()
+       }
        srv.mu.Unlock()
 
        ticker := time.NewTicker(shutdownPollInterval)
@@ -2491,6 +2495,17 @@ func (srv *Server) Shutdown(ctx context.Context) error {
        }
 }
 
+// RegisterOnShutdown registers a function to call on Shutdown.
+// This can be used to gracefully shutdown connections that have
+// undergone NPN/ALPN protocol upgrade or that have been hijacked.
+// This function should start protocol-specific graceful shutdown,
+// but should not wait for shutdown to complete.
+func (srv *Server) RegisterOnShutdown(f func()) {
+       srv.mu.Lock()
+       srv.onShutdown = append(srv.onShutdown, f)
+       srv.mu.Unlock()
+}
+
 // closeIdleConns closes all idle connections and reports whether the
 // server is quiescent.
 func (s *Server) closeIdleConns() bool {