expression: make sleep function response to the kill statement (#10959

) (#12159)
pingcap · Sep 12, 2019 · c260994 · c260994
1 parent 603f0b9
commit c260994
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 25 deletions.
diff --git 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"
@@ -124,12 +125,26 @@ 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:
+			fmt.Println("================")
+			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
@@ -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) {

diff --git a/server/conn_test.go b/server/conn_test.go
@@ -221,17 +221,17 @@ func (ts ConnTestSuite) TestConnExecutionTimeout(c *C) {
 	_, err = se.Execute(context.Background(), "set @@max_execution_time = 500;")
 	c.Assert(err, IsNil)
 
-	err = cc.handleQuery(context.Background(), "select * FROM testTable2 WHERE SLEEP(1);")
-	c.Assert(err, NotNil)
+	err = cc.handleQuery(context.Background(), "select * FROM testTable2 WHERE SLEEP(3);")
+	c.Assert(err, IsNil)
 
 	_, err = se.Execute(context.Background(), "set @@max_execution_time = 0;")
 	c.Assert(err, IsNil)
 
 	err = cc.handleQuery(context.Background(), "select * FROM testTable2 WHERE SLEEP(1);")
 	c.Assert(err, IsNil)
 
-	err = cc.handleQuery(context.Background(), "select /*+ MAX_EXECUTION_TIME(100)*/  * FROM testTable2 WHERE  SLEEP(1);")
-	c.Assert(err, NotNil)
+	err = cc.handleQuery(context.Background(), "select /*+ MAX_EXECUTION_TIME(100)*/  * FROM testTable2 WHERE  SLEEP(3);")
+	c.Assert(err, IsNil)
 
 	c.Assert(failpoint.Disable("github.com/pingcap/tidb/server/FakeClientConn"), IsNil)
 }