feat: configure application transports and unify graceful shutdown

transport/native/server.go

@@ -46,6 +46,7 @@ type Server struct {
 	ReadTimeout           time.Duration
 	WriteTimeout          time.Duration
 	IdleTimeout           time.Duration
+	ShutdownTimeout       time.Duration
 	ReadBufferSize        int
 	MaxHeaderBytes        int
 	MaxBodyBytes          int64
@@ -252,15 +253,14 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
 	s.mu.Unlock()
 	defer cancel()
 
-	shutdownDone := make(chan struct{})
+	stopWatcher, watcherDone := make(chan struct{}), make(chan struct{})
 	go func() {
+		defer close(watcherDone)
 		select {
 		case <-ctx.Done():
-			timeout := s.shutdownTimeout()
-			shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
-			_ = s.Shutdown(shutdownCtx)
-			cancel()
-		case <-shutdownDone:
+			// Wake Accept; the serving goroutine owns draining and its error.
+			_ = listener.Close()
+		case <-stopWatcher:
 		}
 	}()
 
@@ -307,20 +307,8 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
 		go s.serveConnection(baseCtx, state)
 	}
 
-	close(shutdownDone)
-	s.mu.Lock()
-	if s.listener == listener {
-		s.listener = nil
-	}
-	s.mu.Unlock()
-	if ctx.Err() != nil || s.isClosing() {
-		timeout := s.shutdownTimeout()
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
-		shutdownErr := s.Shutdown(shutdownCtx)
-		cancel()
-		return errors.Join(shutdownErr, serveErr)
-	}
-	_ = listener.Close()
+	close(stopWatcher)
+	<-watcherDone
 	timeout := s.shutdownTimeout()
 	shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
 	shutdownErr := s.Shutdown(shutdownCtx)
@@ -328,17 +316,15 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
 	return errors.Join(shutdownErr, serveErr)
 }
 
-// Shutdown stops accepting requests, cancels active request contexts and waits
-// for connections to finish. On timeout, active sockets are closed.
+// Shutdown stops accepting requests and waits for active handlers to finish.
+// On timeout, request contexts are canceled and active sockets are closed.
 func (s *Server) Shutdown(ctx context.Context) error {
 	if ctx == nil {
 		ctx = context.Background()
 	}
 	s.mu.Lock()
 	s.closing = true
-	if s.cancel != nil {
-		s.cancel()
-	}
+	cancelRequests := s.cancel
 	listener := s.listener
 	states := make([]*connection, 0, len(s.conns))
 	for state := range s.conns {
@@ -365,6 +351,9 @@ func (s *Server) Shutdown(ctx context.Context) error {
 	case <-done:
 		return nil
 	case <-ctx.Done():
+		if cancelRequests != nil {
+			cancelRequests()
+		}
 		for _, state := range states {
 			state.close()
 		}
@@ -374,6 +363,9 @@ func (s *Server) Shutdown(ctx context.Context) error {
 }
 
 func (s *Server) shutdownTimeout() time.Duration {
+	if s.ShutdownTimeout > 0 {
+		return s.ShutdownTimeout
+	}
 	if value := s.WriteTimeout; value > 0 {
 		return value
 	}
@@ -432,9 +424,15 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 		}
 		_ = state.conn.SetReadDeadline(time.Time{})
 		firstRequest = false
+		s.mu.Lock()
+		if s.closing || state.closed.Load() {
+			s.mu.Unlock()
+			return
+		}
+		state.active.Store(true)
+		s.mu.Unlock()
 		requestContext, cancel := context.WithCancel(baseCtx)
 		request = request.WithContext(requestContext)
-		state.active.Store(true)
 		s.activeRequests.Add(1)
 		writer := s.writerPool.Get().(*responseWriter)
 		writer.reset(request.Method)