diff --git a/expression/builtin_miscellaneous.go b/expression/builtin_miscellaneous.go index 9308b73e3ea02..239489acd2272 100644 --- a/expression/builtin_miscellaneous.go +++ b/expression/builtin_miscellaneous.go @@ -19,6 +19,7 @@ import ( "math" "net" "strings" + "sync/atomic" "time" "github.com/google/uuid" @@ -131,12 +132,25 @@ func (b *builtinSleepSig) evalInt(row chunk.Row) (int64, bool, error) { if val > math.MaxFloat64/float64(time.Second.Nanoseconds()) { return 0, false, errIncorrectArgs.GenWithStackByArgs("sleep") } + dur := time.Duration(val * float64(time.Second.Nanoseconds())) - select { - case <-time.After(dur): - // TODO: Handle Ctrl-C is pressed in `mysql` client. - // return 1 when SLEEP() is KILLed + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + start := time.Now() + finish := false + for !finish { + select { + case now := <-ticker.C: + if now.Sub(start) > dur { + finish = true + } + default: + if atomic.CompareAndSwapUint32(&sessVars.Killed, 1, 0) { + return 1, false, nil + } + } } + return 0, false, nil } diff --git a/expression/evaluator_test.go b/expression/evaluator_test.go index 8c04dd8bb59ac..db1bbd54d02b5 100644 --- a/expression/evaluator_test.go +++ b/expression/evaluator_test.go @@ -14,6 +14,7 @@ package expression import ( + "sync/atomic" "testing" "time" @@ -176,23 +177,21 @@ func (s *testEvaluatorSuite) TestSleep(c *C) { sub := time.Since(start) c.Assert(sub.Nanoseconds(), GreaterEqual, int64(0.5*1e9)) - // quit when context canceled. - // TODO: recover it. - // d[0].SetFloat64(2) - // f, err = fc.getFunction(ctx, s.datumsToConstants(d)) - // c.Assert(err, IsNil) - // start = time.Now() - // go func() { - // time.Sleep(1 * time.Second) - // ctx.Cancel() - // }() - // ret, isNull, err = f.evalInt(chunk.Row{}) - // sub = time.Since(start) - // c.Assert(err, IsNil) - // c.Assert(isNull, IsFalse) - // c.Assert(ret, Equals, int64(1)) - // c.Assert(sub.Nanoseconds(), LessEqual, int64(2*1e9)) - // c.Assert(sub.Nanoseconds(), GreaterEqual, int64(1*1e9)) + d[0].SetFloat64(3) + f, err = fc.getFunction(ctx, s.datumsToConstants(d)) + c.Assert(err, IsNil) + start = time.Now() + go func() { + time.Sleep(1 * time.Second) + atomic.CompareAndSwapUint32(&ctx.GetSessionVars().Killed, 0, 1) + }() + ret, isNull, err = f.evalInt(chunk.Row{}) + sub = time.Since(start) + c.Assert(err, IsNil) + c.Assert(isNull, IsFalse) + c.Assert(ret, Equals, int64(1)) + c.Assert(sub.Nanoseconds(), LessEqual, int64(2*1e9)) + c.Assert(sub.Nanoseconds(), GreaterEqual, int64(1*1e9)) } func (s *testEvaluatorSuite) TestBinopComparison(c *C) {