allow context cancellation in WAIT (#524)

* allow context cancellation in WAIT

* cancel context before running the program

pkg/runtime/program_test.go

@@ -32,23 +32,11 @@ func TestProgram(t *testing.T) {
 		c := compiler.New()
 		p := c.MustCompile(`WAIT(1000) RETURN TRUE`)
 
-		out := make(chan Result)
-
 		ctx, cancel := context.WithCancel(context.Background())
-
-		go func() {
-			v, err := p.Run(ctx)
-
-			out <- Result{
-				Value: v,
-				Error: err,
-			}
-		}()
-
 		cancel()
 
-		o := <-out
+		_, err := p.Run(ctx)
 
-		So(o.Error, ShouldEqual, core.ErrTerminated)
+		So(err, ShouldEqual, core.ErrTerminated)
 	})
 }

pkg/stdlib/utils/wait.go

@@ -10,7 +10,7 @@ import (
 
 // Wait pauses the execution for a given period.
 // @param timeout (Float|Int) - Number value which indicates for how long to stop an execution.
-func Wait(_ context.Context, args ...core.Value) (core.Value, error) {
+func Wait(ctx context.Context, args ...core.Value) (core.Value, error) {
 	err := core.ValidateArgs(args, 1, 1)
 
 	if err != nil {
@@ -19,7 +19,13 @@ func Wait(_ context.Context, args ...core.Value) (core.Value, error) {
 
 	arg := values.ToInt(args[0])
 
-	time.Sleep(time.Millisecond * time.Duration(arg))
+	timer := time.NewTimer(time.Millisecond * time.Duration(arg))
+	select {
+	case <-ctx.Done():
+		timer.Stop()
+		return values.None, ctx.Err()
+	case <-timer.C:
+	}
 
 	return values.None, nil
 }