recover from panic

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/gin-contrib/pprof v1.4.0
 	github.com/gin-gonic/gin v1.8.1
 	github.com/go-mysql-org/go-mysql v1.6.0
+	github.com/go-sql-driver/mysql v1.7.0
 	github.com/pingcap/tidb v1.1.0-beta.0.20230103132820-3ccff46aa3bc
 	github.com/pingcap/tidb/parser v0.0.0-20230103132820-3ccff46aa3bc
 	github.com/pingcap/tiproxy/lib v0.0.0-00010101000000-000000000000

go.sum

@@ -235,6 +235,7 @@ github.com/go-playground/validator/v10 v10.10.0/go.mod h1:74x4gJWsvQexRdW8Pn3dXS
 github.com/go-sql-driver/mysql v1.3.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
 github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ7YPc=
+github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo=
 github.com/gobwas/pool v0.2.0/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=

lib/util/logger/logger.go

@@ -6,6 +6,7 @@ package logger
 import (
 	"bytes"
 	"fmt"
+	"sync"
 	"testing"
 
 	"go.uber.org/zap"
@@ -14,15 +15,20 @@ import (
 
 type testingLog struct {
 	*testing.T
+	sync.Mutex
 	buf bytes.Buffer
 }
 
 func (t *testingLog) Write(b []byte) (int, error) {
+	t.Lock()
+	defer t.Unlock()
 	t.Logf("%s", b)
 	return t.buf.Write(b)
 }
 
 func (t *testingLog) String() string {
+	t.Lock()
+	defer t.Unlock()
 	return t.buf.String()
 }
 

pkg/proxy/proxy.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pingcap/tiproxy/pkg/proxy/client"
 	"github.com/pingcap/tiproxy/pkg/proxy/keepalive"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
+	"github.com/pingcap/tiproxy/pkg/util"
 	"go.uber.org/zap"
 )
 
@@ -118,7 +119,7 @@ func (s *SQLServer) Run(ctx context.Context, cfgch <-chan *config.Config) {
 				}
 
 				s.wg.Run(func() {
-					s.onConn(ctx, conn)
+					util.WithRecovery(func() { s.onConn(ctx, conn) }, nil, s.logger)
 				})
 			}
 		}

pkg/proxy/proxy_test.go

@@ -4,14 +4,22 @@
 package proxy
 
 import (
+	"context"
+	"database/sql"
 	"net"
+	"strings"
 	"testing"
 	"time"
 
+	_ "github.com/go-sql-driver/mysql"
 	"github.com/pingcap/tiproxy/lib/config"
+	"github.com/pingcap/tiproxy/lib/util/errors"
 	"github.com/pingcap/tiproxy/lib/util/logger"
+	"github.com/pingcap/tiproxy/pkg/manager/cert"
+	"github.com/pingcap/tiproxy/pkg/manager/router"
 	"github.com/pingcap/tiproxy/pkg/proxy/backend"
 	"github.com/pingcap/tiproxy/pkg/proxy/client"
+	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/stretchr/testify/require"
 )
 
@@ -98,3 +106,49 @@ func TestGracefulShutdown(t *testing.T) {
 	case <-finish:
 	}
 }
+
+func TestRecoverPanic(t *testing.T) {
+	lg, text := logger.CreateLoggerForTest(t)
+	certManager := cert.NewCertManager()
+	err := certManager.Init(&config.Config{}, lg, nil)
+	require.NoError(t, err)
+	server, err := NewSQLServer(lg, config.ProxyServer{
+		Addr: "0.0.0.0:6000",
+	}, certManager, &panicHsHandler{})
+	require.NoError(t, err)
+	server.Run(context.Background(), nil)
+
+	mdb, err := sql.Open("mysql", "root@tcp(localhost:6000)/test")
+	require.NoError(t, err)
+	// The first connection encounters panic.
+	require.ErrorContains(t, mdb.Ping(), "invalid connection")
+	require.Eventually(t, func() bool {
+		return strings.Contains(text.String(), "panic")
+	}, 3*time.Second, 10*time.Millisecond)
+	// The second connection gets a server error, which means the server is still running.
+	require.ErrorContains(t, mdb.Ping(), "no router")
+	require.NoError(t, mdb.Close())
+	require.NoError(t, server.Close())
+	certManager.Close()
+}
+
+type panicHsHandler struct {
+	backend.DefaultHandshakeHandler
+}
+
+// HandleHandshakeResp only panics for the first connections.
+func (handler *panicHsHandler) HandleHandshakeResp(ctx backend.ConnContext, _ *pnet.HandshakeResp) error {
+	if ctx.Value(backend.ConnContextKeyConnID).(uint64) == 0 {
+		panic("HandleHandshakeResp panic")
+	}
+	return nil
+}
+
+func (handler *panicHsHandler) GetServerVersion() string {
+	return "5.7"
+}
+
+// GetRouter returns an error for the second connection.
+func (handler *panicHsHandler) GetRouter(backend.ConnContext, *pnet.HandshakeResp) (router.Router, error) {
+	return nil, errors.New("no router")
+}

pkg/util/misc.go

@@ -0,0 +1,23 @@
+// Copyright 2023 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package util
+
+import (
+	"go.uber.org/zap"
+)
+
+func WithRecovery(exec func(), recoverFn func(r interface{}), logger *zap.Logger) {
+	defer func() {
+		r := recover()
+		if recoverFn != nil {
+			recoverFn(r)
+		}
+		if r != nil {
+			logger.Error("panic in the recoverable goroutine",
+				zap.Reflect("r", r),
+				zap.Stack("stack trace"))
+		}
+	}()
+	exec()
+}

pkg/util/misc_test.go

@@ -0,0 +1,30 @@
+// Copyright 2023 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package util
+
+import (
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/pingcap/tiproxy/lib/util/errors"
+	"github.com/pingcap/tiproxy/lib/util/logger"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithRecovery(t *testing.T) {
+	lg, text := logger.CreateLoggerForTest(t)
+	ch := make(chan error, 1)
+	go WithRecovery(func() {
+		panic("mock panic")
+	}, func(r interface{}) {
+		if r != nil {
+			ch <- errors.Errorf("%v", r)
+		}
+	}, lg)
+	require.Error(t, <-ch)
+	require.Eventually(t, func() bool {
+		return strings.Contains(text.String(), "mock panic")
+	}, time.Second, 10*time.Millisecond)
+}