expression: make sleep function response to the kill stateme… (#11028)

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
 }
 

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) {