diff --git a/pkg/executor/test/executor/executor_test.go b/pkg/executor/test/executor/executor_test.go index 4cb70182acb8a..95a575a0d7078 100644 --- a/pkg/executor/test/executor/executor_test.go +++ b/pkg/executor/test/executor/executor_test.go @@ -18,6 +18,7 @@ import ( "archive/zip" "context" "fmt" + "math/rand" "os" "path/filepath" "reflect" @@ -3013,3 +3014,64 @@ func TestIssue50308(t *testing.T) { tk.MustQuery("show warnings").Check(testkit.RowsWithSep("|", "Warning 1292 Incorrect timestamp value: '2099-01-01'")) tk.MustQuery("select * from t;").Check(testkit.Rows("0000-00-00 00:00:00")) } + +func TestQueryWithKill(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test;") + tk.MustExec("drop table if exists tkq;") + tk.MustExec("create table tkq (a int key, b int, index idx_b(b));") + tk.MustExec("insert into tkq values (1,1);") + var wg sync.WaitGroup + ch := make(chan context.CancelFunc, 1024) + testDuration := time.Second * 10 + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test;") + start := time.Now() + for { + ctx, cancel := context.WithCancel(context.Background()) + ch <- cancel + rs, err := tk.ExecWithContext(ctx, "select a from tkq where b = 1;") + if err == nil { + require.NotNil(t, rs) + rows, err := session.ResultSetToStringSlice(ctx, tk.Session(), rs) + if err == nil { + require.Equal(t, 1, len(rows)) + require.Equal(t, 1, len(rows[0])) + require.Equal(t, "1", fmt.Sprintf("%v", rows[0][0])) + } + } + if err != nil { + require.Equal(t, context.Canceled, err) + } + if rs != nil { + rs.Close() + } + if time.Since(start) > testDuration { + return + } + } + }() + } + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case cancel := <-ch: + // mock for random kill query + if len(ch) < 5 { + time.Sleep(time.Duration(rand.Intn(1000)) * time.Nanosecond) + } + cancel() + case <-time.After(time.Second): + return + } + } + }() + wg.Wait() +}