Fix shutdown waitgroup panic

Author: lunny

cmd/web.go

@@ -156,7 +156,6 @@ func serveInstall(cmd *cli.Command) error {
 	case <-graceful.GetManager().IsShutdown():
 		<-graceful.GetManager().Done()
 		log.Info("PID: %d Gitea Web Finished", os.Getpid())
-		log.GetManager().Close()
 		return err
 	default:
 	}
@@ -170,6 +169,16 @@ func serveInstalled(c *cli.Command) error {
 
 	showWebStartupMessage("Prepare to run web server")
 
+	// Check for shutdown signal during startup
+	select {
+	case <-graceful.GetManager().IsShutdown():
+		log.Info("Shutdown signal received during startup, aborting server start")
+		<-graceful.GetManager().Done()
+		log.Info("PID: %d Gitea Web Finished", os.Getpid())
+		return nil
+	default:
+	}
+
 	if setting.AppWorkPathMismatch {
 		log.Error("WORK_PATH from config %q doesn't match other paths from environment variables or command arguments. "+
 			"Only WORK_PATH in config should be set and used. Please make sure the path in config file is correct, "+
@@ -226,12 +235,21 @@ func serveInstalled(c *cli.Command) error {
 
 	gtprof.EnableBuiltinTracer(util.Iif(setting.IsProd, 2000*time.Millisecond, 100*time.Millisecond))
 
+	// Check for shutdown signal before starting web server
+	select {
+	case <-graceful.GetManager().IsShutdown():
+		log.Info("Shutdown signal received before starting web server, aborting")
+		<-graceful.GetManager().Done()
+		log.Info("PID: %d Gitea Web Finished", os.Getpid())
+		return nil
+	default:
+	}
+
 	// Set up Chi routes
 	webRoutes := routers.NormalRoutes()
 	err := listen(webRoutes, true)
 	<-graceful.GetManager().Done()
 	log.Info("PID: %d Gitea Web Finished", os.Getpid())
-	log.GetManager().Close()
 	return err
 }
 

modules/graceful/server.go

@@ -36,6 +36,8 @@ type Server struct {
 	wg                   sync.WaitGroup
 	state                state
 	lock                 *sync.RWMutex
+	connections          map[*wrappedConn]struct{}
+	connectionsLock      sync.RWMutex
 	BeforeBegin          func(network, address string)
 	OnShutdown           func()
 	PerWriteTimeout      time.Duration
@@ -53,6 +55,7 @@ func NewServer(network, address, name string) *Server {
 		wg:                   sync.WaitGroup{},
 		state:                stateInit,
 		lock:                 &sync.RWMutex{},
+		connections:          make(map[*wrappedConn]struct{}),
 		network:              network,
 		address:              address,
 		PerWriteTimeout:      setting.PerWriteTimeout,
@@ -178,6 +181,21 @@ func (srv *Server) setState(st state) {
 	srv.state = st
 }
 
+// closeAllConnections forcefully closes all active connections
+func (srv *Server) closeAllConnections() {
+	srv.connectionsLock.Lock()
+	connections := make([]*wrappedConn, 0, len(srv.connections))
+	for conn := range srv.connections {
+		connections = append(connections, conn)
+	}
+	srv.connectionsLock.Unlock()
+
+	// Close all connections outside the lock to avoid deadlock
+	for _, conn := range connections {
+		_ = conn.Conn.Close() // Force close the underlying connection
+	}
+}
+
 type filer interface {
 	File() (*os.File, error)
 }
@@ -216,16 +234,20 @@ func (wl *wrappedListener) Accept() (net.Conn, error) {
 
 	closed := int32(0)
 
-	c = &wrappedConn{
-		Conn:                 c,
-		server:               wl.server,
-		closed:               &closed,
-		perWriteTimeout:      wl.server.PerWriteTimeout,
-		perWritePerKbTimeout: wl.server.PerWritePerKbTimeout,
+	wc := &wrappedConn{
+		Conn:   c,
+		server: wl.server,
+		closed: &closed,
 	}
 
 	wl.server.wg.Add(1)
-	return c, nil
+
+	// Track the connection
+	wl.server.connectionsLock.Lock()
+	wl.server.connections[wc] = struct{}{}
+	wl.server.connectionsLock.Unlock()
+
+	return wc, nil
 }
 
 func (wl *wrappedListener) Close() error {
@@ -244,17 +266,15 @@ func (wl *wrappedListener) File() (*os.File, error) {
 
 type wrappedConn struct {
 	net.Conn
-	server               *Server
-	closed               *int32
-	deadline             time.Time
-	perWriteTimeout      time.Duration
-	perWritePerKbTimeout time.Duration
+	server   *Server
+	closed   *int32
+	deadline time.Time
 }
 
 func (w *wrappedConn) Write(p []byte) (n int, err error) {
-	if w.perWriteTimeout > 0 {
-		minTimeout := time.Duration(len(p)/1024) * w.perWritePerKbTimeout
-		minDeadline := time.Now().Add(minTimeout).Add(w.perWriteTimeout)
+	if w.server.PerWriteTimeout > 0 {
+		minTimeout := time.Duration(len(p)/1024) * w.server.PerWritePerKbTimeout
+		minDeadline := time.Now().Add(minTimeout).Add(w.server.PerWriteTimeout)
 
 		w.deadline = w.deadline.Add(minTimeout)
 		if minDeadline.After(w.deadline) {
@@ -278,6 +298,12 @@ func (w *wrappedConn) Close() error {
 				}
 			}
 		}()
+
+		// Remove from tracked connections
+		w.server.connectionsLock.Lock()
+		delete(w.server.connections, w)
+		w.server.connectionsLock.Unlock()
+
 		w.server.wg.Done()
 	}
 	return w.Conn.Close()

modules/graceful/server_hooks.go

@@ -5,7 +5,6 @@ package graceful
 
 import (
 	"os"
-	"runtime"
 
 	"code.gitea.io/gitea/modules/log"
 )
@@ -48,26 +47,14 @@ func (srv *Server) doShutdown() {
 }
 
 func (srv *Server) doHammer() {
-	defer func() {
-		// We call srv.wg.Done() until it panics.
-		// This happens if we call Done() when the WaitGroup counter is already at 0
-		// So if it panics -> we're done, Serve() will return and the
-		// parent will goroutine will exit.
-		if r := recover(); r != nil {
-			log.Error("WaitGroup at 0: Error: %v", r)
-		}
-	}()
 	if srv.getState() != stateShuttingDown {
 		return
 	}
-	log.Warn("Forcefully shutting down parent")
-	for {
-		if srv.getState() == stateTerminate {
-			break
-		}
-		srv.wg.Done()
+	log.Warn("Forcefully closing all connections")
 
-		// Give other goroutines a chance to finish before we forcibly stop them.
-		runtime.Gosched()
-	}
+	// Close all active connections. Each connection's Close() method
+	// will call wg.Done() exactly once, maintaining WaitGroup integrity.
+	// This will allow wg.Wait() in Serve() to complete, and Serve() will
+	// then set the state to stateTerminate.
+	srv.closeAllConnections()
 }