br: fix checksum incomplete when context is done (#44583) (#44735)

close #44472

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
 }

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)