net/http: support passing Server Shutdown context to HTTP/3 implementations

Currently, when Server.Shutdown is called with a given context, an
HTTP/3 server implementation does not have access to said context. This
forces HTTP/3 server implementations to guess a reasonable fixed timeout
when attempting to gracefully shutdown, which is not ideal.

Therefore, add ShutdownContext method to http3ServerHandler to handle
this.

For #77440

Change-Id: Ib15b615f646fd08788981eb06f3a70606a6a6964
Reviewed-on: https://go-review.googlesource.com/c/go/+/761480
Reviewed-by: Nicholas Husin <husin@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Damien Neil <dneil@google.com>

1 files changed, 28 insertions, 11 deletions

diff --git a/src/net/http/server.go b/src/net/http/server.go
index c25e83b704..4cd18e9faf 100644
--- a/src/net/http/server.go
+++ b/src/net/http/server.go
@@ -3121,6 +3121,7 @@ type Server struct {
 	activeConn map[*conn]struct{}
 	onShutdown []func()
 	h2         *http2Server
+	h3         *http3ServerHandler
 
 	listenerGroup sync.WaitGroup
 }
@@ -3188,6 +3189,9 @@ func (s *Server) Shutdown(ctx context.Context) error {
 	s.inShutdown.Store(true)
 
 	s.mu.Lock()
+	if s.h3 != nil {
+		s.h3.shutdownCtx = ctx
+	}
 	lnerr := s.closeListenersLocked()
 	for _, f := range s.onShutdown {
 		go f()
@@ -3751,43 +3755,53 @@ func ListenAndServeTLS(addr, certFile, keyFile string, handler Handler) error {
 // us to test our HTTP/3 implementation againts tests in net/http. HTTP/3 is
 // not yet accessible to end-users.
 type http3ServerHandler struct {
-	handler   serverHandler
-	tlsConfig *tls.Config
-	baseCtx   context.Context
-	errc      chan error
+	handler     serverHandler
+	tlsConfig   *tls.Config
+	baseCtx     context.Context
+	errc        chan error
+	shutdownCtx context.Context
 }
 
 // ServeHTTP ensures that http3ServerHandler implements the Handler interface,
 // and gives an HTTP/3 server implementation access to the net/http handler.
-func (h http3ServerHandler) ServeHTTP(w ResponseWriter, r *Request) {
+func (h *http3ServerHandler) ServeHTTP(w ResponseWriter, r *Request) {
 	h.handler.ServeHTTP(w, r)
 }
 
 // Addr gives an HTTP/3 server implementation the address that it should listen
 // on.
-func (h http3ServerHandler) Addr() string {
+func (h *http3ServerHandler) Addr() string {
 	return h.handler.srv.Addr
 }
 
 // TLSConfig gives an HTTP/3 server implementation the *tls.Config that it
 // should use.
-func (h http3ServerHandler) TLSConfig() *tls.Config {
+func (h *http3ServerHandler) TLSConfig() *tls.Config {
 	return h.tlsConfig
 }
 
 // BaseContext gives an HTTP/3 server implementation the base context to use
 // for server requests.
-func (h http3ServerHandler) BaseContext() context.Context {
+func (h *http3ServerHandler) BaseContext() context.Context {
 	return h.baseCtx
 }
 
 // ListenErrHook should be called by an HTTP/3 server implementation to
 // propagate any error it encounters when trying to listen, if any, to
 // net/http.
-func (h http3ServerHandler) ListenErrHook(err error) {
+func (h *http3ServerHandler) ListenErrHook(err error) {
 	h.errc <- err
 }
 
+// ShutdownContext gives an HTTP/3 server implementation the context that is
+// used when [Server.Shutdown] is called. This allows an HTTP/3 server
+// implementation to know how long it can take to gracefully shutdown in the
+// function it registers with [Server.RegisterOnShutdown]. Callers must not use
+// this method for any other purpose.
+func (h *http3ServerHandler) ShutdownContext() context.Context {
+	return h.shutdownCtx
+}
+
 // ListenAndServeTLS listens on the TCP network address s.Addr and
 // then calls [ServeTLS] to handle requests on incoming TLS connections.
 // Accepted connections are configured to enable TCP keep-alives.
@@ -3823,12 +3837,15 @@ func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
 			return err
 		}
 		errc := make(chan error, 1)
-		go fn(s, nil, http3ServerHandler{
+		s.mu.Lock()
+		s.h3 = &http3ServerHandler{
 			handler:   serverHandler{s},
 			tlsConfig: config,
 			baseCtx:   context.WithValue(context.Background(), ServerContextKey, s),
 			errc:      errc,
-		})
+		}
+		s.mu.Unlock()
+		go fn(s, nil, s.h3)
 		if err := <-errc; err != nil {
 			return err
 		}