feat: make HTTP2 shutdown graceful

transport/http2/server.go

@@ -42,6 +42,9 @@ type Server struct {
 	mu             sync.Mutex
 	handler        http.Handler
 	listener       net.Listener
+	shutdownConfig *http.Server
+	shutdownDone   chan struct{}
+	shutdownOnce   sync.Once
 	closing        bool
 	cancel         context.CancelFunc
 	conns          map[net.Conn]struct{}
@@ -176,9 +179,14 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
 		s.mu.Unlock()
 		return errors.New("go-omni/http2: server is already running")
 	}
-	baseCtx, cancel := context.WithCancel(ctx)
+	// Keep the connection context independent from the caller's context. A
+	// canceled parent starts graceful shutdown; it must not cancel active
+	// streams before the GOAWAY/drain phase has completed.
+	baseCtx, cancel := context.WithCancel(context.Background())
 	s.cancel = cancel
 	s.listener = listener
+	s.shutdownDone = make(chan struct{})
+	s.shutdownOnce = sync.Once{}
 	s.closing = false
 	if s.conns == nil {
 		s.conns = make(map[net.Conn]struct{})
@@ -200,24 +208,43 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
 		MaxConcurrentStreams: s.MaxConcurrentStreams,
 		IdleTimeout:          s.IdleTimeout,
 	}
-	shutdownDone := make(chan struct{})
-	go func() {
-		select {
-		case <-ctx.Done():
-			shutdownCtx, cancel := context.WithTimeout(context.Background(), s.shutdownTimeout())
-			_ = s.Shutdown(shutdownCtx)
-			cancel()
-		case <-shutdownDone:
-		}
-	}()
-
-	defer func() {
-		close(shutdownDone)
+	shutdownConfig := &http.Server{IdleTimeout: s.IdleTimeout}
+	if err := xhttp2.ConfigureServer(shutdownConfig, protocolServer); err != nil {
 		cancel()
 		s.mu.Lock()
 		s.listener = nil
 		s.cancel = nil
 		s.mu.Unlock()
+		return err
+	}
+	s.mu.Lock()
+	s.shutdownConfig = shutdownConfig
+	s.mu.Unlock()
+	contextWakeDone := make(chan struct{})
+	go func() {
+		select {
+		case <-ctx.Done():
+			// Wake Accept. Serve performs the graceful drain synchronously so
+			// its return value reports a shutdown timeout to the caller.
+			_ = listener.Close()
+		case <-contextWakeDone:
+		}
+	}()
+
+	defer func() {
+		close(contextWakeDone)
+		s.mu.Lock()
+		closing := s.closing
+		shutdownDone := s.shutdownDone
+		s.listener = nil
+		s.shutdownConfig = nil
+		s.cancel = nil
+		s.mu.Unlock()
+		if !closing {
+			cancel()
+		} else if shutdownDone != nil {
+			<-shutdownDone
+		}
 	}()
 
 	for {
@@ -251,7 +278,13 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
 		s.activeConns.Add(1)
 		s.wg.Add(1)
 		s.mu.Unlock()
-		go s.serveConn(baseCtx, protocolServer, gate, conn)
+		go s.serveConn(baseCtx, protocolServer, shutdownConfig, gate, conn)
+	}
+	if ctx.Err() != nil && !s.isClosing() {
+		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), s.shutdownTimeout())
+		err := s.Shutdown(shutdownCtx)
+		shutdownCancel()
+		return err
 	}
 	return nil
 }
@@ -263,7 +296,7 @@ func (s *Server) shutdownTimeout() time.Duration {
 	return defaultShutdownTimeout
 }
 
-func (s *Server) serveConn(ctx context.Context, protocolServer *xhttp2.Server, handler http.Handler, conn net.Conn) {
+func (s *Server) serveConn(ctx context.Context, protocolServer *xhttp2.Server, baseConfig *http.Server, handler http.Handler, conn net.Conn) {
 	defer s.wg.Done()
 	defer func() {
 		s.mu.Lock()
@@ -283,11 +316,7 @@ func (s *Server) serveConn(ctx context.Context, protocolServer *xhttp2.Server, h
 		_ = conn.SetDeadline(time.Time{})
 	}
 	protocolServer.ServeConn(conn, &xhttp2.ServeConnOpts{
-		Context: ctx,
-		Handler: handler,
-		BaseConfig: &http.Server{
-			IdleTimeout: s.IdleTimeout,
-		},
+		Context: ctx, Handler: handler, BaseConfig: baseConfig,
 	})
 }
 
@@ -321,6 +350,7 @@ func (s *Server) Shutdown(ctx context.Context) error {
 	}
 	s.closing = true
 	listener := s.listener
+	shutdownConfig := s.shutdownConfig
 	cancel := s.cancel
 	connections := make([]net.Conn, 0, len(s.conns))
 	for conn := range s.conns {
@@ -330,7 +360,12 @@ func (s *Server) Shutdown(ctx context.Context) error {
 	if listener != nil {
 		_ = listener.Close()
 	}
-	cancel()
+	// x/net/http2's configured shutdown callback sends GOAWAY and lets active
+	// streams finish. The base request context remains live until the timeout
+	// path below, where connections are forcibly closed.
+	if shutdownConfig != nil {
+		_ = shutdownConfig.Shutdown(context.Background())
+	}
 
 	done := make(chan struct{})
 	go func() {
@@ -339,15 +374,28 @@ func (s *Server) Shutdown(ctx context.Context) error {
 	}()
 	select {
 	case <-done:
+		cancel()
+		s.finishShutdown()
 		return nil
 	case <-ctx.Done():
 		for _, conn := range connections {
 			_ = conn.Close()
 		}
+		cancel()
+		s.finishShutdown()
 		return ctx.Err()
 	}
 }
 
+func (s *Server) finishShutdown() {
+	s.mu.Lock()
+	done := s.shutdownDone
+	s.mu.Unlock()
+	if done != nil {
+		s.shutdownOnce.Do(func() { close(done) })
+	}
+}
+
 func (s *Server) isClosing() bool {
 	s.mu.Lock()
 	closing := s.closing