From 052c17f4f69bdedd8fae91c75e95cb70a5c6385f Mon Sep 17 00:00:00 2001
From: Jianjun Liao <36503113+Leavrth@users.noreply.github.com>
Date: Fri, 16 Jun 2023 16:47:10 +0800
Subject: [PATCH] br: fix checksum incomplete when context is done (#44583)

close pingcap/tidb#44472
---
 br/pkg/checksum/executor.go      | 13 ++++++++++++-
 br/pkg/checksum/executor_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/br/pkg/checksum/executor.go b/br/pkg/checksum/executor.go
index 15a478d2da335..0159fe43b1d0c 100644
--- a/br/pkg/checksum/executor.go
+++ b/br/pkg/checksum/executor.go
@@ -365,5 +365,16 @@ func (exec *Executor) Execute(
 		updateChecksumResponse(checksumResp, resp)
 		updateFn()
 	}
-	return checksumResp, nil
+	return checksumResp, checkContextDone(ctx)
+}
+
+// The coprocessor won't return the error if the context is done,
+// so sometimes BR would get the incomplete result.
+// checkContextDone makes sure the result is not affected by CONTEXT DONE.
+func checkContextDone(ctx context.Context) error {
+	ctxErr := ctx.Err()
+	if ctxErr != nil {
+		return errors.Annotate(ctxErr, "context is cancelled by other error")
+	}
+	return nil
 }
diff --git a/br/pkg/checksum/executor_test.go b/br/pkg/checksum/executor_test.go
index 5d198b88d951e..a432d29234997 100644
--- a/br/pkg/checksum/executor_test.go
+++ b/br/pkg/checksum/executor_test.go
@@ -29,6 +29,35 @@ func getTableInfo(t *testing.T, mock *mock.Cluster, db, table string) *model.Tab
 	return tableInfo.Meta()
 }
 
+func TestChecksumContextDone(t *testing.T) {
+	mock, err := mock.NewCluster()
+	require.NoError(t, err)
+	require.NoError(t, mock.Start())
+	defer mock.Stop()
+
+	tk := testkit.NewTestKit(t, mock.Storage)
+	tk.MustExec("use test")
+
+	tk.MustExec("drop table if exists t1;")
+	tk.MustExec("create table t1 (a int, b int, key i1(a, b), primary key (a));")
+	tk.MustExec("insert into t1 values (10, 10);")
+	tableInfo1 := getTableInfo(t, mock, "test", "t1")
+	exe, err := checksum.NewExecutorBuilder(tableInfo1, math.MaxUint64).
+		SetConcurrency(variable.DefChecksumTableConcurrency).
+		Build()
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	cctx, cancel := context.WithCancel(ctx)
+
+	cancel()
+
+	resp, err := exe.Execute(cctx, mock.Storage.GetClient(), func() { t.Log("request done") })
+	t.Log(err)
+	t.Log(resp)
+	require.Error(t, err)
+}
+
 func TestChecksum(t *testing.T) {
 	mock, err := mock.NewCluster()
 	require.NoError(t, err)