Make Sleep check for cancelled context every second

Signed-off-by: Juanjo Alvarez <[email protected]>

Author: Juanjo Alvarez

Diff for: server/context.go

@@ -24,7 +24,7 @@ func DefaultSessionBuilder(c *mysql.Conn, addr string) sql.Session {
 
 // SessionManager is in charge of creating new sessions for the given
 // connections and keep track of which sessions are in each connection, so
-// they can be cancelled is the connection is closed.
+// they can be cancelled if the connection is closed.
 type SessionManager struct {
 	addr     string
 	tracer   opentracing.Tracer

Diff for: sql/expression/function/sleep.go

@@ -37,7 +37,28 @@ func (s *Sleep) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
 		return nil, err
 	}
 
-	time.Sleep(time.Duration(child.(float64)*1000) * time.Millisecond)
+	fchild := child.(float64)
+	if fchild <= 0 {
+		return 0, nil
+	}
+
+	// Wake up every second to check if the context was cancelled
+	remaining := fchild * 1000.0
+	for remaining >= 0 {
+		toSleep := 1000.0
+		if remaining < 1000 {
+			toSleep = remaining
+		}
+		remaining -= 1000
+
+		select {
+		case <-ctx.Done():
+			goto End
+		case <-time.After(time.Duration(toSleep) * time.Millisecond):
+		}
+	}
+End:
+
 	return 0, nil
 }