*: recover goroutines when they panic (#488)

Author: djshow832

lib/util/waitgroup/waitgroup.go

@@ -5,6 +5,8 @@ package waitgroup
 
 import (
 	"sync"
+
+	"go.uber.org/zap"
 )
 
 // WaitGroup is a wrapper for sync.WaitGroup
@@ -27,15 +29,25 @@ func (w *WaitGroup) Run(exec func()) {
 // and call done when function return. it will dump current goroutine stack into log if catch any recover result.
 // exec is that execute logic function. recoverFn is that handler will be called after recover and before dump stack,
 // passing `nil` means noop.
-func (w *WaitGroup) RunWithRecover(exec func(), recoverFn func(r interface{})) {
+func (w *WaitGroup) RunWithRecover(exec func(), recoverFn func(r interface{}), logger *zap.Logger) {
 	w.Add(1)
 	go func() {
 		defer func() {
 			r := recover()
+			defer func() {
+				// If it panics again in recovery, quit ASAP.
+				_ = recover()
+			}()
+			if r != nil && logger != nil {
+				logger.Error("panic in the recoverable goroutine",
+					zap.Reflect("r", r),
+					zap.Stack("stack trace"))
+			}
+			// Call Done() before recoverFn because recoverFn normally calls `Close()`, which may call `wg.Wait()`.
+			w.Done()
 			if r != nil && recoverFn != nil {
 				recoverFn(r)
 			}
-			w.Done()
 		}()
 		exec()
 	}()

lib/util/waitgroup/waitgroup_test.go

@@ -0,0 +1,31 @@
+// Copyright 2024 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package waitgroup
+
+import (
+	"testing"
+
+	"github.com/pingcap/tiproxy/lib/util/logger"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithRecoveryLog(t *testing.T) {
+	lg, text := logger.CreateLoggerForTest(t)
+	var wg WaitGroup
+	wg.RunWithRecover(func() {
+		panic("mock panic")
+	}, nil, lg)
+	wg.Wait()
+	require.Contains(t, text.String(), "mock panic")
+}
+
+func TestWithRecoveryPanic(t *testing.T) {
+	var wg WaitGroup
+	wg.RunWithRecover(func() {
+		panic("mock panic1")
+	}, func(r interface{}) {
+		panic("mock panic2")
+	}, nil)
+	wg.Wait()
+}

pkg/manager/cert/manager.go

@@ -98,6 +98,7 @@ func (cm *CertManager) SQLTLS() *tls.Config {
 // The proxy periodically reloads certs. If it fails, we will retry in the next round.
 // If configuration changes, it only affects new connections by returning new *tls.Config.
 func (cm *CertManager) reloadLoop(ctx context.Context, cfgch <-chan *config.Config) {
+	// Failing to reload certs may cause even more serious problems than TiProxy reboot, so we don't recover panics.
 	cm.wg.Run(func() {
 		for {
 			select {

pkg/manager/config/manager.go

@@ -82,7 +82,7 @@ func (e *ConfigManager) Init(ctx context.Context, logger *zap.Logger, configFile
 			return err
 		}
 
-		e.wg.Run(func() {
+		e.wg.RunWithRecover(func() {
 			// Read the file periodically and reload the config if it changes.
 			//
 			// We tried other ways to watch file:
@@ -100,7 +100,7 @@ func (e *ConfigManager) Init(ctx context.Context, logger *zap.Logger, configFile
 					}
 				}
 			}
-		})
+		}, nil, e.logger)
 	} else {
 		if err := e.SetTOMLConfig(nil); err != nil {
 			return err

pkg/manager/infosync/info.go

@@ -109,9 +109,9 @@ func (is *InfoSyncer) Init(ctx context.Context, cfg *config.Config, certMgr *cer
 
 	childCtx, cancelFunc := context.WithCancel(ctx)
 	is.cancelFunc = cancelFunc
-	is.wg.Run(func() {
+	is.wg.RunWithRecover(func() {
 		is.updateTopologyLivenessLoop(childCtx, topologyInfo)
-	})
+	}, nil, is.lg)
 	return nil
 }
 

pkg/manager/logger/manager.go

@@ -47,9 +47,9 @@ func (lm *LoggerManager) Init(cfgch <-chan *config.Config) {
 	ctx, cancel := context.WithCancel(context.Background())
 	lm.cancel = cancel
 
-	lm.wg.Run(func() {
+	lm.wg.RunWithRecover(func() {
 		lm.watchCfg(ctx, cfgch)
-	})
+	}, nil, lm.logger)
 }
 
 func (lm *LoggerManager) watchCfg(ctx context.Context, cfgch <-chan *config.Config) {

pkg/manager/router/backend_observer.go

@@ -123,6 +123,7 @@ func NewBackendObserver(logger *zap.Logger, eventReceiver BackendEventReceiver,
 func (bo *BackendObserver) Start() {
 	childCtx, cancelFunc := context.WithCancel(context.Background())
 	bo.cancelFunc = cancelFunc
+	// Failing to observe backends may cause even more serious problems than TiProxy reboot, so we don't recover panics.
 	bo.wg.Run(func() {
 		bo.observe(childCtx)
 	})

pkg/manager/router/router_score.go

@@ -50,6 +50,7 @@ func (r *ScoreBasedRouter) Init(fetcher BackendFetcher, hc HealthCheck, cfg *con
 	r.observer = observer
 	childCtx, cancelFunc := context.WithCancel(context.Background())
 	r.cancelFunc = cancelFunc
+	// Failing to rebalance backends may cause even more serious problems than TiProxy reboot, so we don't recover panics.
 	r.wg.Run(func() {
 		r.rebalanceLoop(childCtx)
 	})

pkg/metrics/metrics.go

@@ -79,9 +79,9 @@ func (mm *MetricsManager) setupMonitor(ctx context.Context) {
 			KeepAliveCounter.Inc()
 		}
 	}
-	mm.wg.Run(func() {
+	mm.wg.RunWithRecover(func() {
 		systimemon.StartMonitor(ctx, mm.logger, time.Now, systimeErrHandler, successCallBack)
-	})
+	}, nil, mm.logger)
 }
 
 // registerProxyMetrics registers metrics.

pkg/proxy/backend/backend_conn_mgr.go

@@ -203,9 +203,13 @@ func (mgr *BackendConnManager) Connect(ctx context.Context, clientIO *pnet.Packe
 	childCtx, cancelFunc := context.WithCancel(ctx)
 	mgr.cancelFunc = cancelFunc
 	mgr.lastActiveTime = endTime
-	mgr.wg.Run(func() {
+	mgr.wg.RunWithRecover(func() {
 		mgr.processSignals(childCtx)
-	})
+	}, func(_ any) {
+		// If we do not clean up, the router may retain the connection forever and the TiDB won't be released in the
+		// Serverless Tier.
+		_ = mgr.Close()
+	}, mgr.logger)
 	return nil
 }
 
@@ -422,22 +426,26 @@ func (mgr *BackendConnManager) processSignals(ctx context.Context) {
 	for {
 		select {
 		case s := <-mgr.signalReceived:
-			// Redirect the session immediately just in case the session is finishedTxn.
-			mgr.processLock.Lock()
-			switch s {
-			case signalTypeGracefulClose:
-				mgr.tryGracefulClose(ctx)
-			case signalTypeRedirect:
-				mgr.tryRedirect(ctx)
-			}
-			mgr.processLock.Unlock()
+			func() {
+				// Redirect the session immediately just in case the session is finishedTxn.
+				mgr.processLock.Lock()
+				defer mgr.processLock.Unlock()
+				switch s {
+				case signalTypeGracefulClose:
+					mgr.tryGracefulClose(ctx)
+				case signalTypeRedirect:
+					mgr.tryRedirect(ctx)
+				}
+			}()
 		case rs := <-mgr.redirectResCh:
 			mgr.notifyRedirectResult(ctx, rs)
 		case <-checkBackendTicker.C:
-			mgr.checkBackendActive()
-			mgr.processLock.Lock()
-			mgr.setKeepAlive()
-			mgr.processLock.Unlock()
+			func() {
+				mgr.checkBackendActive()
+				mgr.processLock.Lock()
+				defer mgr.processLock.Unlock()
+				mgr.setKeepAlive()
+			}()
 		case <-ctx.Done():
 			checkBackendTicker.Stop()
 			return