Commit 1a63f116 authored by Tom Bergan's avatar Tom Bergan Committed by Brad Fitzpatrick

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: default avatarBrad Fitzpatrick <bradfitz@golang.org>
parent 6374a660
...@@ -5232,7 +5232,8 @@ func testServerShutdown(t *testing.T, h2 bool) { ...@@ -5232,7 +5232,8 @@ func testServerShutdown(t *testing.T, h2 bool) {
defer afterTest(t) defer afterTest(t)
var doShutdown func() // set later var doShutdown func() // set later
var shutdownRes = make(chan error, 1) 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() go doShutdown()
// Shutdown is graceful, so it should not interrupt // Shutdown is graceful, so it should not interrupt
// this in-flight response. Add a tiny sleep here to // this in-flight response. Add a tiny sleep here to
...@@ -5240,7 +5241,10 @@ func testServerShutdown(t *testing.T, h2 bool) { ...@@ -5240,7 +5241,10 @@ func testServerShutdown(t *testing.T, h2 bool) {
// bugs. // bugs.
time.Sleep(20 * time.Millisecond) time.Sleep(20 * time.Millisecond)
io.WriteString(w, r.RemoteAddr) io.WriteString(w, r.RemoteAddr)
})) })
cst := newClientServerTest(t, h2, handler, func(srv *httptest.Server) {
srv.Config.RegisterOnShutdown(func() { gotOnShutdown <- struct{}{} })
})
defer cst.close() defer cst.close()
doShutdown = func() { doShutdown = func() {
...@@ -5251,6 +5255,11 @@ func testServerShutdown(t *testing.T, h2 bool) { ...@@ -5251,6 +5255,11 @@ func testServerShutdown(t *testing.T, h2 bool) {
if err := <-shutdownRes; err != nil { if err := <-shutdownRes; err != nil {
t.Fatalf("Shutdown: %v", err) 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) res, err := cst.c.Get(cst.ts.URL)
if err == nil { if err == nil {
......
...@@ -2395,6 +2395,7 @@ type Server struct { ...@@ -2395,6 +2395,7 @@ type Server struct {
listeners map[net.Listener]struct{} listeners map[net.Listener]struct{}
activeConn map[*conn]struct{} activeConn map[*conn]struct{}
doneChan chan struct{} doneChan chan struct{}
onShutdown []func()
} }
func (s *Server) getDoneChan() <-chan struct{} { func (s *Server) getDoneChan() <-chan struct{} {
...@@ -2475,6 +2476,9 @@ func (srv *Server) Shutdown(ctx context.Context) error { ...@@ -2475,6 +2476,9 @@ func (srv *Server) Shutdown(ctx context.Context) error {
srv.mu.Lock() srv.mu.Lock()
lnerr := srv.closeListenersLocked() lnerr := srv.closeListenersLocked()
srv.closeDoneChanLocked() srv.closeDoneChanLocked()
for _, f := range srv.onShutdown {
go f()
}
srv.mu.Unlock() srv.mu.Unlock()
ticker := time.NewTicker(shutdownPollInterval) ticker := time.NewTicker(shutdownPollInterval)
...@@ -2491,6 +2495,17 @@ func (srv *Server) Shutdown(ctx context.Context) error { ...@@ -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 // closeIdleConns closes all idle connections and reports whether the
// server is quiescent. // server is quiescent.
func (s *Server) closeIdleConns() bool { func (s *Server) closeIdleConns() bool {
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment