Repository: google/starlark-go
Pull request #433 (2022-10-17 05:41:10)

Imlemented ResetStepsAndCancelReason()
Imlemented ResetStepsAndCancelReason() and a test
This ResetStepsAndCancelReason() is very helpful for limiting cpu usage per Call or CallInternal


Patch:
diff --git a/starlark/eval.go b/starlark/eval.go
index 4941d44..379d038 100644
--- a/starlark/eval.go
+++ b/starlark/eval.go
@@ -51,8 +51,13 @@ type Thread struct {
 	// The default behavior is to call thread.Cancel("too many steps").
 	OnMaxSteps func(thread *Thread)
 
-	// steps counts abstract computation steps executed by this thread.
-	steps, maxSteps uint64
+	// Steps 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.
+	Steps, maxSteps uint64
 
 	// cancelReason records the reason from the first call to Cancel.
 	cancelReason *string
@@ -65,14 +70,9 @@ 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.
+// ExecutionSteps returns the current value of Steps.
 func (thread *Thread) ExecutionSteps() uint64 {
-	return thread.steps
+	return thread.Steps
 }
 
 // SetMaxExecutionSteps sets a limit on the number of Starlark
@@ -84,12 +84,20 @@ func (thread *Thread) SetMaxExecutionSteps(max uint64) {
 	thread.maxSteps = max
 }
 
+// Uncancel resets the cancellation state.
+//
+// Unlike most methods of Thread, it is safe to call Uncancel from any
+// goroutine, even if the thread is actively executing.
+func (thread *Thread) Uncancel() {
+	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&thread.cancelReason)), nil)
+}
+
 // 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.
+// Call [Uncancel] to reset the cancellation state.
 //
 // Unlike most methods of Thread, it is safe to call Cancel from any
 // goroutine, even if the thread is actively executing.
diff --git a/starlark/interp.go b/starlark/interp.go
index 9961e2f..8b073bd 100644
--- a/starlark/interp.go
+++ b/starlark/interp.go
@@ -87,8 +87,8 @@ func (fn *Function) CallInternal(thread *Thread, args Tuple, kwargs []Tuple) (Va
 	code := f.Code
 loop:
 	for {
-		thread.steps++
-		if thread.steps >= thread.maxSteps {
+		thread.Steps++
+		if thread.Steps >= thread.maxSteps {
 			if thread.OnMaxSteps != nil {
 				thread.OnMaxSteps(thread)
 			} else {


Test patch:
diff --git a/starlark/eval_test.go b/starlark/eval_test.go
index 5974579..448401e 100644
--- a/starlark/eval_test.go
+++ b/starlark/eval_test.go
@@ -972,6 +972,13 @@ func TestExecutionSteps(t *testing.T) {
 	if fmt.Sprint(err) != "Starlark computation cancelled: too many steps" {
 		t.Errorf("execution returned error %q, want cancellation", err)
 	}
+
+	thread.Steps = 0
+	thread.Uncancel()
+	_, err = countSteps(1)
+	if err != nil {
+		t.Errorf("execution returned error %q, want nil", err)
+	}
 }
 
 // TestDeps fails if the interpreter proper (not the REPL, etc) sprouts new external dependencies.