diff --git a/transport/http2/server.go b/transport/http2/server.go index 1ca5651..ea0493a 100644 --- a/transport/http2/server.go +++ b/transport/http2/server.go @@ -51,6 +51,7 @@ type Server struct { closing bool cancel context.CancelFunc conns map[net.Conn]struct{} + protocolConns map[net.Conn]struct{} connSlots chan struct{} requestSlots chan struct{} accepted atomic.Uint64 @@ -204,6 +205,9 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error { if s.conns == nil { s.conns = make(map[net.Conn]struct{}) } + if s.protocolConns == nil { + s.protocolConns = make(map[net.Conn]struct{}) + } if s.MaxConnections > 0 { s.connSlots = make(chan struct{}, s.MaxConnections) } else { @@ -330,6 +334,7 @@ func (s *Server) serveConn(ctx context.Context, protocolServer *xhttp2.Server, b defer func() { s.mu.Lock() delete(s.conns, conn) + delete(s.protocolConns, conn) s.activeConns.Add(-1) s.mu.Unlock() if s.connSlots != nil { @@ -344,6 +349,13 @@ func (s *Server) serveConn(ctx context.Context, protocolServer *xhttp2.Server, b } _ = conn.SetDeadline(time.Time{}) } + s.mu.Lock() + if s.closing { + s.mu.Unlock() + return + } + s.protocolConns[conn] = struct{}{} + s.mu.Unlock() protocolServer.ServeConn(conn, &xhttp2.ServeConnOpts{ Context: ctx, Handler: handler, BaseConfig: baseConfig, }) @@ -386,6 +398,7 @@ func (s *Server) Shutdown(ctx context.Context) error { for conn := range s.conns { connections = append(connections, conn) } + protocolConnectionCount := len(s.protocolConns) s.mu.Unlock() if listener != nil { _ = listener.Close() @@ -394,7 +407,7 @@ func (s *Server) Shutdown(ctx context.Context) error { // streams finish. The base request context remains live until the timeout // path below, where connections are forcibly closed. if shutdownConfig != nil && !alreadyClosing { - s.goAwaySent.Add(uint64(len(connections))) + s.goAwaySent.Add(uint64(protocolConnectionCount)) _ = shutdownConfig.Shutdown(context.Background()) }