add context to script run

This commit is contained in:
Daniel Kang 2019-01-20 04:59:43 -08:00
parent 47418fddac
commit 04bb0f5e62
3 changed files with 61 additions and 0 deletions

View file

@ -1,6 +1,8 @@
package script
import (
"context"
"github.com/d5/tengo/compiler"
"github.com/d5/tengo/objects"
"github.com/d5/tengo/runtime"
@ -20,6 +22,25 @@ func (c *Compiled) Run() error {
return c.machine.Run()
}
// RunContext is like Run but includes a context.
func (c *Compiled) RunContext(ctx context.Context) (err error) {
ch := make(chan error, 1)
go func() {
ch <- c.machine.Run()
}()
select {
case <-ctx.Done():
c.machine.Abort()
<-ch
err = ctx.Err()
case err = <-ch:
}
return
}
// IsDefined returns true if the variable name is defined (has value) before or after the execution.
func (c *Compiled) IsDefined(name string) bool {
symbol, _, ok := c.symbolTable.Resolve(name)

View file

@ -1,7 +1,9 @@
package script_test
import (
"context"
"testing"
"time"
"github.com/d5/tengo/assert"
"github.com/d5/tengo/script"
@ -45,6 +47,31 @@ func TestCompiled_IsDefined(t *testing.T) {
compiledIsDefined(t, c, "b", false)
}
func TestCompiled_RunContext(t *testing.T) {
// machine completes normally
c := compile(t, `a := 5`, nil)
err := c.RunContext(context.Background())
assert.NoError(t, err)
compiledGet(t, c, "a", int64(5))
// cancelled
c = compile(t, `for true {}`, nil)
ctx, cancel := context.WithCancel(context.Background())
go func() {
time.Sleep(1 * time.Millisecond)
cancel()
}()
err = c.RunContext(ctx)
assert.Equal(t, context.Canceled, err)
// timeout
c = compile(t, `for true {}`, nil)
ctx, cancel = context.WithTimeout(context.Background(), 1*time.Millisecond)
defer cancel()
err = c.RunContext(ctx)
assert.Equal(t, context.DeadlineExceeded, err)
}
func compile(t *testing.T, input string, vars M) *script.Compiled {
s := script.New([]byte(input))
for vn, vv := range vars {

View file

@ -1,6 +1,7 @@
package script
import (
"context"
"fmt"
"github.com/d5/tengo/compiler"
@ -87,6 +88,18 @@ func (s *Script) Run() (compiled *Compiled, err error) {
return
}
// RunContext is like Run but includes a context.
func (s *Script) RunContext(ctx context.Context) (compiled *Compiled, err error) {
compiled, err = s.Compile()
if err != nil {
return
}
err = compiled.RunContext(ctx)
return
}
func (s *Script) prepCompile() (symbolTable *compiler.SymbolTable, globals []*objects.Object) {
var names []string
for name := range s.variables {