[fix] cancel a load() thread when the run's context fires

cache.go

@@ -1,6 +1,7 @@
 package starlet
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io/fs"
@@ -33,6 +34,11 @@ type cache struct {
 	// nil the load path falls back to a bare thread — a module loaded that way
 	// would inherit none of the Machine's guards.
 	newThread func(load func(*starlark.Thread, string) (starlark.StringDict, error)) *starlark.Thread
+	// watchCancel, when set, arms cancellation of the load thread from the run's
+	// context and returns a stop func to disarm it once loading finishes, so a
+	// module executed by load() honours the run timeout instead of running past
+	// it on its own thread. When nil the load thread is not cancelled by context.
+	watchCancel func(thread *starlark.Thread) (stop func())
 }
 
 type entry struct {
@@ -108,17 +114,31 @@ func (c *cache) get(cc *cycleChecker, module string) (starlark.StringDict, error
 			c.remove(module)
 			panic(r)
 		}()
-		e.globals, e.err = c.doLoad(cc, module)
+		var loadThread *starlark.Thread
+		e.globals, e.err = c.doLoad(cc, module, &loadThread)
 		loaded = true
 		e.setOwner(nil)
 
 		// Broadcast that the entry is now ready.
 		close(e.ready)
+
+		// A load aborted because the run's context was cancelled is transient
+		// (run-scoped), not a property of the module: evict it so a later run
+		// reusing this machine without Reset re-executes instead of replaying
+		// the stale cancellation error. Mirrors the step-budget panic eviction.
+		// Waiters already blocked on e.ready still observe this run's error; only
+		// fresh loads after eviction re-execute.
+		if e.err != nil && loadContextCancelled(loadThread) {
+			c.remove(module)
+		}
 	}
 	return e.globals, e.err
 }
 
-func (c *cache) doLoad(cc *cycleChecker, module string) (starlark.StringDict, error) {
+// doLoad loads module. It publishes the thread the module executes on through
+// *loadThread so the caller can tell whether a failure was a transient
+// run-context cancellation (see get) without re-indenting the load body.
+func (c *cache) doLoad(cc *cycleChecker, module string, loadThread **starlark.Thread) (starlark.StringDict, error) {
 	// Tunnel the cycle-checker state for this "thread of loading".
 	loadFn := func(_ *starlark.Thread, module string) (starlark.StringDict, error) {
 		return c.get(cc, module)
@@ -135,6 +155,14 @@ func (c *cache) doLoad(cc *cycleChecker, module string) (starlark.StringDict, er
 			Load:  loadFn,
 		}
 	}
+	*loadThread = thread
+
+	// Cancel this load thread when the run's context fires, so a long
+	// computation inside a loaded module honours the run timeout rather than
+	// running to completion on its own thread. Disarmed once loading finishes.
+	if c.watchCancel != nil {
+		defer c.watchCancel(thread)()
+	}
 
 	// 1: load from built-in module, the first field returns nil if not found
 	m, err := c.loadMod(module)
@@ -172,6 +200,14 @@ func (c *cache) doLoad(cc *cycleChecker, module string) (starlark.StringDict, er
 	return starlark.ExecFileOptions(c.execOpts, thread, module, b, c.globals)
 }
 
+// loadContextCancelled reports whether the given load thread's run context
+// (carried as the "context" thread-local by newLoadThread) has been cancelled.
+// A single expression so it is fully covered by any load, cancelled or not.
+func loadContextCancelled(thread *starlark.Thread) bool {
+	ctx, ok := thread.Local("context").(context.Context)
+	return ok && ctx != nil && ctx.Err() != nil
+}
+
 // -- concurrent cycle checking --
 
 // A cycleChecker is used for concurrent deadlock detection.

internal_test.go

@@ -1,13 +1,37 @@
 package starlet
 
 import (
+	"context"
 	"reflect"
 	"strings"
 	"testing"
 
 	"go.starlark.net/starlark"
 )
 
+// TestWatchLoadThread covers the load-thread cancel watcher's two arms. A load
+// thread with no run context (e.g. a load during a REPL, which sets none) must
+// get a harmless no-op stop rather than watch a nil context; a thread carrying a
+// live, cancellable context gets a real watcher, and cancelling then stopping it
+// must be safe. The firing-on-timeout behaviour itself is covered end to end by
+// TestMachine_LoadInheritsMachineContext.
+func TestWatchLoadThread(t *testing.T) {
+	m := NewDefault()
+
+	// no context on the thread -> no-op stop.
+	stop := m.watchLoadThread(&starlark.Thread{})
+	stop() // must not panic
+
+	// a live context -> a real watcher; cancel then stop must be safe/idempotent.
+	ctx, cancel := context.WithCancel(context.Background())
+	th := &starlark.Thread{}
+	th.SetLocal("context", ctx)
+	stop = m.watchLoadThread(th)
+	cancel()
+	stop()
+	stop()
+}
+
 func TestCastStringDictToAnyMap(t *testing.T) {
 	// Create a starlark.StringDict
 	m := starlark.StringDict{

module.go

@@ -180,6 +180,13 @@ func MakeModuleLoaderFromMap(m StringAnyMap) ModuleLoader {
 }
 
 // MakeModuleLoaderFromString creates a module loader from the given source code.
+//
+// The returned loader executes the source on its own bare thread, so the module
+// body does NOT inherit a Machine's step budget or run-context cancellation. The
+// source is host-supplied (not script input), so this is a robustness caveat,
+// not a script-reachable DoS; but for a module whose body must be bounded by the
+// Machine's guards, load it through the file path instead (SetScript's scriptFS
+// plus a load(name) of the file), which runs on the guarded load thread.
 func MakeModuleLoaderFromString(name, source string, predeclared starlark.StringDict) ModuleLoader {
 	return func() (starlark.StringDict, error) {
 		if name == "" {
@@ -190,6 +197,11 @@ func MakeModuleLoaderFromString(name, source string, predeclared starlark.String
 }
 
 // MakeModuleLoaderFromReader creates a module loader from the given IO reader.
+//
+// Like MakeModuleLoaderFromString, the returned loader executes the source on a
+// bare thread and does NOT inherit a Machine's step budget or run-context
+// cancellation; load host-supplied source through the file path when it must be
+// bounded by the Machine's guards.
 func MakeModuleLoaderFromReader(name string, rd io.Reader, predeclared starlark.StringDict) ModuleLoader {
 	return func() (starlark.StringDict, error) {
 		if name == "" {
@@ -200,6 +212,12 @@ func MakeModuleLoaderFromReader(name string, rd io.Reader, predeclared starlark.
 }
 
 // MakeModuleLoaderFromFile creates a module loader from the given file.
+//
+// Like MakeModuleLoaderFromString, the returned loader executes the file on a
+// bare thread and does NOT inherit a Machine's step budget or run-context
+// cancellation. To have a Machine bound a loaded file with its guards, register
+// the filesystem via SetScript and load(name) the file so it runs on the guarded
+// load thread instead of registering it through this constructor.
 func MakeModuleLoaderFromFile(name string, fileSys fs.FS, predeclared starlark.StringDict) ModuleLoader {
 	return func() (starlark.StringDict, error) {
 		// read file content

run.go

@@ -100,19 +100,19 @@ func (m *Machine) RunWithContext(ctx context.Context, extras StringAnyMap) (Stri
 	return m.runInternal(ctx, extras, true)
 }
 
-// watchContextCancel cancels the machine's thread when ctx fires, until the
-// returned stop function is called. stop is idempotent and waits for the
-// watcher goroutine to exit, so callers can both defer it (panic safety)
-// and invoke it right after execution finishes.
-func (m *Machine) watchContextCancel(ctx context.Context) (stop func()) {
+// watchThreadCancelCtx cancels thread when ctx fires, until the returned stop
+// function is called. stop is idempotent and waits for the watcher goroutine to
+// exit, so callers can both defer it (panic safety) and invoke it right after
+// execution finishes.
+func watchThreadCancelCtx(ctx context.Context, thread *starlark.Thread) (stop func()) {
 	var wg sync.WaitGroup
 	wg.Add(1)
 	done := make(chan struct{})
 	go func() {
 		defer wg.Done()
 		select {
 		case <-ctx.Done():
-			m.thread.Cancel("context cancelled")
+			thread.Cancel("context cancelled")
 		case <-done:
 			// no action if execution has finished
 		}
@@ -126,6 +126,27 @@ func (m *Machine) watchContextCancel(ctx context.Context) (stop func()) {
 	}
 }
 
+// watchContextCancel cancels the machine's main thread when ctx fires, until
+// the returned stop function is called.
+func (m *Machine) watchContextCancel(ctx context.Context) (stop func()) {
+	return watchThreadCancelCtx(ctx, m.thread)
+}
+
+// watchLoadThread cancels a load thread when the run's context (carried as the
+// "context" thread-local by newLoadThread) fires, until the returned stop is
+// called. A module executed by load() runs on its own thread, so without this
+// the run's timeout/cancellation would interrupt only the main thread and a
+// long computation inside a loaded module would run past the deadline (bounded
+// only by its step budget, or unbounded when none is set). Returns a no-op stop
+// when the thread carries no cancellable context.
+func (m *Machine) watchLoadThread(thread *starlark.Thread) (stop func()) {
+	ctx, ok := thread.Local("context").(context.Context)
+	if !ok || ctx == nil {
+		return func() {}
+	}
+	return watchThreadCancelCtx(ctx, thread)
+}
+
 func (m *Machine) runInternal(ctx context.Context, extras StringAnyMap, allowCache bool) (out StringAnyMap, err error) {
 	defer func() {
 		if r := recover(); r != nil {
@@ -274,8 +295,9 @@ func (m *Machine) prepareThread(extras StringAnyMap) (err error) {
 			readFile: func(name string) ([]byte, error) {
 				return readScriptFile(name, m.scriptFS)
 			},
-			globals:   m.predeclared,
-			newThread: m.newLoadThread,
+			globals:     m.predeclared,
+			newThread:   m.newLoadThread,
+			watchCancel: m.watchLoadThread,
 		}
 		m.thread = &starlark.Thread{
 			Name:  "starlet",

run_test.go

@@ -81,6 +81,87 @@ func TestMachine_LoadInheritsMachineContext(t *testing.T) {
 		}
 	})
 
+	t.Run("a run timeout cancels a loaded module's long computation", func(t *testing.T) {
+		// A module executed by load() runs on its own thread. The run's timeout
+		// cancelled only the main thread; the load thread had the run context as
+		// a thread-local but no cancel watcher, so a long computation inside a
+		// loaded module ran to completion and defeated the timeout (bounded only
+		// by the step budget, or unbounded with none). The load thread is now
+		// cancelled too, so the run returns at the deadline instead of hanging.
+		sfs := fstest.MapFS{
+			// a long, non-allocating spin that runs THROUGH the interpreter (so
+			// cancellation is checked between iterations, unlike a single builtin
+			// like max()) while the always-false filter keeps the list empty, so
+			// it far outlasts the timeout without OOM when it is not cancelled.
+			"spin.star": &fstest.MapFile{Data: []byte(`x = len([i for i in range(2000000000) if i < 0])`)},
+		}
+		m := starlet.NewDefault()
+		m.SetScript("main.star", []byte(`load("spin.star", "x"); y = x`), sfs)
+		// No step budget: the timeout is the only guard, which is the case a
+		// host relying on RunWithTimeout (e.g. starbox RunTimeout) hits.
+
+		done := make(chan error, 1)
+		start := time.Now()
+		go func() {
+			_, err := m.RunWithTimeout(200*time.Millisecond, nil)
+			done <- err
+		}()
+		select {
+		case err := <-done:
+			if err == nil {
+				t.Fatal("expected a timeout error, got nil (the loaded loop was not interrupted)")
+			}
+			// generous margin: the timeout is 200ms; a runaway (uncancelled) loop
+			// of 2e9 interpreter iterations takes many seconds even on a fast
+			// machine, so this cleanly separates "cancelled" from "ran to the end"
+			// without flaking on a slow/contended CI runner under -race.
+			if elapsed := time.Since(start); elapsed > 4*time.Second {
+				t.Fatalf("run returned only after %v; the loaded module's loop was not cancelled by the timeout", elapsed)
+			}
+		case <-time.After(6 * time.Second):
+			t.Fatal("run did not return: the loaded module's loop ignored the timeout")
+		}
+	})
+
+	t.Run("a cancelled load does not poison the cache", func(t *testing.T) {
+		// Cancelling a load makes it return an error. That failure is transient
+		// (run-scoped), not a property of the module, so it must not stay cached:
+		// a second run reusing the same machine (without Reset) must re-execute
+		// the module and succeed instead of replaying the stale cancellation.
+		// Mirrors the step-budget panic eviction. The cancellation is driven
+		// deterministically here (a builtin cancels the run's context and its own
+		// load thread on the first run only) so the test does not depend on race
+		// timing between a wall-clock timeout and a loop.
+		firstRun := true
+		var cancelRun1 context.CancelFunc
+		cancelPoint := starlark.NewBuiltin("cancel_point", func(th *starlark.Thread, _ *starlark.Builtin, _ starlark.Tuple, _ []starlark.Tuple) (starlark.Value, error) {
+			if firstRun {
+				firstRun = false
+				cancelRun1()              // cancel run 1's context: the load is context-cancelled
+				th.Cancel("cancel_point") // stop this load thread now, deterministically
+			}
+			return starlark.None, nil
+		})
+		sfs := fstest.MapFS{
+			"work.star": &fstest.MapFile{Data: []byte("cancel_point()\nx = 1")},
+		}
+		m := starlet.NewDefault()
+		m.SetGlobals(starlet.StringAnyMap{"cancel_point": cancelPoint})
+		m.SetScript("main.star", []byte(`load("work.star", "x"); y = x`), sfs)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancelRun1 = cancel
+		if _, err := m.RunWithContext(ctx, nil); err == nil {
+			t.Fatal("first run: expected the cancelled load to fail, got nil")
+		}
+
+		// Second run on the SAME machine, no Reset, uncancelled: the load must
+		// re-execute (cancel_point is now a no-op) and succeed.
+		if _, err := m.Run(); err != nil {
+			t.Fatalf("second run: the cancelled load poisoned the cache: %v", err)
+		}
+	})
+
 	t.Run("a panic(nil) loader does not poison the cache", func(t *testing.T) {
 		// The cache cleanup keys on a completion sentinel, not recover() != nil,
 		// so it survives panic(nil) too (recover() is nil for it under go 1.19).