From 089cba4853d42cd435bc7831f0fcd3a2f651e50a Mon Sep 17 00:00:00 2001
From: guo-shaoge <shaoge1994@163.com>
Date: Thu, 28 Mar 2024 15:06:48 +0800
Subject: [PATCH] pkg/util: fix panic caused by logging grpc err (#52179)

close pingcap/tidb#51301
---
 pkg/util/topsql/reporter/BUILD.bazel    | 1 +
 pkg/util/topsql/reporter/pubsub.go      | 6 ++++++
 pkg/util/topsql/reporter/pubsub_test.go | 5 +++++
 3 files changed, 12 insertions(+)

diff --git a/pkg/util/topsql/reporter/BUILD.bazel b/pkg/util/topsql/reporter/BUILD.bazel
index 0b92142fa43a9..7f612bb2a412d 100644
--- a/pkg/util/topsql/reporter/BUILD.bazel
+++ b/pkg/util/topsql/reporter/BUILD.bazel
@@ -53,6 +53,7 @@ go_test(
         "//pkg/util/topsql/reporter/mock",
         "//pkg/util/topsql/state",
         "//pkg/util/topsql/stmtstats",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_tipb//go-tipb",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
diff --git a/pkg/util/topsql/reporter/pubsub.go b/pkg/util/topsql/reporter/pubsub.go
index 0210df87d9212..a9ce1fe7e9b8c 100644
--- a/pkg/util/topsql/reporter/pubsub.go
+++ b/pkg/util/topsql/reporter/pubsub.go
@@ -19,6 +19,7 @@ import (
 	"errors"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	reporter_metrics "github.com/pingcap/tidb/pkg/util/topsql/reporter/metrics"
@@ -97,6 +98,10 @@ func (ds *pubSubDataSink) OnReporterClosing() {
 
 func (ds *pubSubDataSink) run() error {
 	defer func() {
+		if r := recover(); r != nil {
+			// To catch panic when log grpc error. https://github.com/pingcap/tidb/issues/51301.
+			logutil.BgLogger().Error("[top-sql] got panic in pub sub data sink, just ignore", zap.Error(util.GetRecoverError(r)))
+		}
 		ds.registerer.Deregister(ds)
 		ds.cancel()
 	}()
@@ -133,6 +138,7 @@ func (ds *pubSubDataSink) run() error {
 				return ctx.Err()
 			}
 
+			failpoint.Inject("mockGrpcLogPanic", nil)
 			if err != nil {
 				logutil.BgLogger().Warn(
 					"[top-sql] pubsub datasink failed to send data to subscriber",
diff --git a/pkg/util/topsql/reporter/pubsub_test.go b/pkg/util/topsql/reporter/pubsub_test.go
index f82229c5ab53f..b25d23269052c 100644
--- a/pkg/util/topsql/reporter/pubsub_test.go
+++ b/pkg/util/topsql/reporter/pubsub_test.go
@@ -20,8 +20,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tipb/go-tipb"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/metadata"
 )
 
@@ -85,6 +87,8 @@ func TestPubSubDataSink(t *testing.T) {
 		_ = ds.run()
 	}()
 
+	panicPath := "github.com/pingcap/tidb/pkg/util/topsql/reporter/mockGrpcLogPanic"
+	require.NoError(t, failpoint.Enable(panicPath, "panic"))
 	err := ds.TrySend(&ReportData{
 		DataRecords: []tipb.TopSQLRecord{{
 			SqlDigest:  []byte("S1"),
@@ -117,4 +121,5 @@ func TestPubSubDataSink(t *testing.T) {
 	mockStream.Unlock()
 
 	ds.OnReporterClosing()
+	require.NoError(t, failpoint.Disable(panicPath))
 }