Add a ServeWithContext() method to httputil

It is sometimes useful to control an HTTP server with a
Context (e.g. when using an errgroup). NOTE that this function does
not install a SIGTERM handler.

serverutil/http.go

@@ -104,13 +104,20 @@ func (config *ServerConfig) buildHTTPHandler(h http.Handler) (http.Handler, *tls
 	return h, tlsConfig, nil
 }
 
-// Serve HTTP(S) content on the specified address. If config.TLS is
-// not nil, enable HTTPS and TLS authentication.
-//
-// This function will return an error if there are problems creating
-// the listener, otherwise it will handle graceful termination on
-// SIGINT or SIGTERM and return nil.
-func Serve(h http.Handler, config *ServerConfig, addr string) error {
+func buildListener(addr string, tlsConfig *tls.Config) (net.Listener, error) {
+	// Create the net.Listener first, so we can detect
+	// initialization-time errors safely.
+	l, err := net.Listen("tcp", addr)
+	if err != nil {
+		return nil, err
+	}
+	if tlsConfig != nil {
+		l = tls.NewListener(l, tlsConfig)
+	}
+	return l, nil
+}
+
+func buildServer(h http.Handler, config *ServerConfig, addr string) (*http.Server, error) {
 	// Wrap with tracing handler (exclude metrics and other
 	// debugging endpoints).
 	h = tracing.WrapHandler(h, guessEndpointName(addr))
@@ -118,7 +125,7 @@ func Serve(h http.Handler, config *ServerConfig, addr string) error {
 	// Create the top-level HTTP handler with all our additions.
 	hh, tlsConfig, err := config.buildHTTPHandler(h)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	// These are not meant to be external-facing servers, so we
@@ -131,14 +138,24 @@ func Serve(h http.Handler, config *ServerConfig, addr string) error {
 		TLSConfig:         tlsConfig,
 	}
 
-	// Create the net.Listener first, so we can detect
-	// initialization-time errors safely.
-	l, err := net.Listen("tcp", addr)
+	return srv, nil
+}
+
+// Serve HTTP(S) content on the specified address. If config.TLS is
+// not nil, enable HTTPS and TLS authentication.
+//
+// This function will return an error if there are problems creating
+// the listener, otherwise it will handle graceful termination on
+// SIGINT or SIGTERM and return nil.
+func Serve(h http.Handler, config *ServerConfig, addr string) error {
+	srv, err := buildServer(h, config, addr)
 	if err != nil {
 		return err
 	}
-	if srv.TLSConfig != nil {
-		l = tls.NewListener(l, srv.TLSConfig)
+
+	l, err := buildListener(addr, srv.TLSConfig)
+	if err != nil {
+		return err
 	}
 
 	// Install a signal handler for gentle process termination.
@@ -176,6 +193,38 @@ func Serve(h http.Handler, config *ServerConfig, addr string) error {
 	return nil
 }
 
+// ServeWithContext operates like Serve but with a controlling Context
+// that can be used to stop the HTTP server.
+func ServeWithContext(ctx context.Context, h http.Handler, config *ServerConfig, addr string) error {
+	srv, err := buildServer(h, config, addr)
+	if err != nil {
+		return err
+	}
+
+	l, err := buildListener(addr, srv.TLSConfig)
+	if err != nil {
+		return err
+	}
+
+	go func() {
+		<-ctx.Done()
+
+		sctx, cancel := context.WithTimeout(context.Background(), gracefulShutdownTimeout)
+		srv.Shutdown(sctx) // nolint: errcheck
+		srv.Close()
+		cancel()
+	}()
+
+	daemon.SdNotify(false, "READY=1") // nolint
+
+	err = srv.Serve(l)
+	if err == http.ErrServerClosed {
+		err = nil
+	}
+
+	return err
+}
+
 func addDefaultHandlers(h http.Handler) http.Handler {
 	root := http.NewServeMux()