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>

1 files changed, 15 insertions, 0 deletions

diff --git a/src/net/http/server.go b/src/net/http/server.go
index 71f46a74f9..b60bd2481e 100644
--- 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 {