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

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])