diff --git a/executor/executor_test.go b/executor/executor_test.go index aefd677218bbd..5679158971ebe 100644 --- a/executor/executor_test.go +++ b/executor/executor_test.go @@ -19,6 +19,7 @@ import ( "context" "fmt" "math" + "math/rand" "path/filepath" "reflect" "runtime" @@ -6441,3 +6442,64 @@ func TestIssue51324(t *testing.T) { tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1364 Field 'a' doesn't have a default value", "Warning 1364 Field 'b' doesn't have a default value", "Warning 1364 Field 'c' doesn't have a default value", "Warning 1364 Field 'd' doesn't have a default value")) tk.MustQuery("select * from t order by id").Check(testkit.Rows("13 0 0", "21 1 0 1", "22 1 0 1")) } + +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() +}