diff options
author | alandonovan <adonovan@google.com> | 2020-08-21 10:29:38 -0400 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-08-21 10:29:38 -0400 |
commit | 949cc6f4b09716061bc5c308049b0da811aeca10 (patch) | |
tree | 122707c2d3c8c9855f260a822fd4e3d038cae7de | |
parent | 4379bb3f9ac0b7f6e21b1441663122a26f96be31 (diff) | |
download | starlark-go-949cc6f4b09716061bc5c308049b0da811aeca10.tar.gz |
starlark: support thread cancellation and limits on computation (#298)
This change adds two related features:
(1) asynchronous cancellation of Starlark threads (Thread.Cancel), and
(2) recording and limiting of the number of abstract computation steps.
Fixes issue #252
Fixes issue #160
-rw-r--r-- | starlark/eval.go | 48 | ||||
-rw-r--r-- | starlark/eval_test.go | 61 | ||||
-rw-r--r-- | starlark/interp.go | 11 |
3 files changed, 120 insertions, 0 deletions
diff --git a/starlark/eval.go b/starlark/eval.go index e7e827d..6aa3ead 100644 --- a/starlark/eval.go +++ b/starlark/eval.go @@ -13,9 +13,11 @@ import ( "math/big" "sort" "strings" + "sync/atomic" "time" "unicode" "unicode/utf8" + "unsafe" "go.starlark.net/internal/compile" "go.starlark.net/internal/spell" @@ -46,6 +48,12 @@ type Thread struct { // See example_test.go for some example implementations of Load. Load func(thread *Thread, module string) (StringDict, error) + // steps counts abstract computation steps executed by this thread. + steps, maxSteps uint64 + + // cancelReason records the reason from the first call to Cancel. + cancelReason *string + // locals holds arbitrary "thread-local" Go values belonging to the client. // They are accessible to the client but not to any Starlark program. locals map[string]interface{} @@ -54,6 +62,38 @@ type Thread struct { proftime time.Duration } +// ExecutionSteps returns a count of abstract computation steps executed +// by this thread. It is incremented by the interpreter. It may be used +// as a measure of the approximate cost of Starlark execution, by +// computing the difference in its value before and after a computation. +// +// The precise meaning of "step" is not specified and may change. +func (thread *Thread) ExecutionSteps() uint64 { + return thread.steps +} + +// SetMaxExecutionSteps sets a limit on the number of Starlark +// computation steps that may be executed by this thread. If the +// thread's step counter exceeds this limit, the interpreter calls +// thread.Cancel("too many steps"). +func (thread *Thread) SetMaxExecutionSteps(max uint64) { + thread.maxSteps = max +} + +// Cancel causes execution of Starlark code in the specified thread to +// promptly fail with an EvalError that includes the specified reason. +// There may be a delay before the interpreter observes the cancellation +// if the thread is currently in a call to a built-in function. +// +// Cancellation cannot be undone. +// +// Unlike most methods of Thread, it is safe to call Cancel from any +// goroutine, even if the thread is actively executing. +func (thread *Thread) Cancel(reason string) { + // Atomically set cancelReason, preserving earlier reason if any. + atomic.CompareAndSwapPointer((*unsafe.Pointer)(unsafe.Pointer(&thread.cancelReason)), nil, unsafe.Pointer(&reason)) +} + // SetLocal sets the thread-local value associated with the specified key. // It must not be called after execution begins. func (thread *Thread) SetLocal(key string, value interface{}) { @@ -1077,6 +1117,14 @@ func Call(thread *Thread, fn Value, args Tuple, kwargs []Tuple) (Value, error) { if fr == nil { fr = new(frame) } + + if thread.stack == nil { + // one-time initialization of thread + if thread.maxSteps == 0 { + thread.maxSteps-- // (MaxUint64) + } + } + thread.stack = append(thread.stack, fr) // push fr.callable = c diff --git a/starlark/eval_test.go b/starlark/eval_test.go index e7aef9d..a57ccad 100644 --- a/starlark/eval_test.go +++ b/starlark/eval_test.go @@ -834,3 +834,64 @@ func TestREPLChunk(t *testing.T) { t.Fatalf("chunk2: got %s, want %s", got, want) } } + +func TestCancel(t *testing.T) { + // A thread cancelled before it begins executes no code. + { + thread := new(starlark.Thread) + thread.Cancel("nope") + _, err := starlark.ExecFile(thread, "precancel.star", `x = 1//0`, nil) + if fmt.Sprint(err) != "Starlark computation cancelled: nope" { + t.Errorf("execution returned error %q, want cancellation", err) + } + + // cancellation is sticky + _, err = starlark.ExecFile(thread, "precancel.star", `x = 1//0`, nil) + if fmt.Sprint(err) != "Starlark computation cancelled: nope" { + t.Errorf("execution returned error %q, want cancellation", err) + } + } + // A thread cancelled during a built-in executes no more code. + { + thread := new(starlark.Thread) + predeclared := starlark.StringDict{ + "stopit": starlark.NewBuiltin("stopit", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + thread.Cancel(fmt.Sprint(args[0])) + return starlark.None, nil + }), + } + _, err := starlark.ExecFile(thread, "stopit.star", `msg = 'nope'; stopit(msg); x = 1//0`, predeclared) + if fmt.Sprint(err) != `Starlark computation cancelled: "nope"` { + t.Errorf("execution returned error %q, want cancellation", err) + } + } +} + +func TestExecutionSteps(t *testing.T) { + // A Thread records the number of computation steps. + thread := new(starlark.Thread) + countSteps := func(n int) (uint64, error) { + predeclared := starlark.StringDict{"n": starlark.MakeInt(n)} + steps0 := thread.ExecutionSteps() + _, err := starlark.ExecFile(thread, "steps.star", `squares = [x*x for x in range(n)]`, predeclared) + return thread.ExecutionSteps() - steps0, err + } + steps100, err := countSteps(1000) + if err != nil { + t.Errorf("execution failed: %v", err) + } + steps10000, err := countSteps(100000) + if err != nil { + t.Errorf("execution failed: %v", err) + } + if ratio := float64(steps10000) / float64(steps100); ratio < 99 || ratio > 101 { + t.Errorf("computation steps did not increase linearly: f(100)=%d, f(10000)=%d, ratio=%g, want ~100", steps100, steps10000, ratio) + } + + // Exceeding the step limit causes cancellation. + thread.SetMaxExecutionSteps(1000) + _, err = countSteps(1000) + if fmt.Sprint(err) != "Starlark computation cancelled: too many steps" { + t.Errorf("execution returned error %q, want cancellation", err) + } +} diff --git a/starlark/interp.go b/starlark/interp.go index 004763e..f00382c 100644 --- a/starlark/interp.go +++ b/starlark/interp.go @@ -5,6 +5,8 @@ package starlark import ( "fmt" "os" + "sync/atomic" + "unsafe" "go.starlark.net/internal/compile" "go.starlark.net/internal/spell" @@ -85,6 +87,15 @@ func (fn *Function) CallInternal(thread *Thread, args Tuple, kwargs []Tuple) (Va code := f.Code loop: for { + thread.steps++ + if thread.steps >= thread.maxSteps { + thread.Cancel("too many steps") + } + if reason := atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&thread.cancelReason))); reason != nil { + err = fmt.Errorf("Starlark computation cancelled: %s", *(*string)(reason)) + break loop + } + fr.pc = pc op := compile.Opcode(code[pc]) |