diff --git a/cache.go b/cache.go index 27fa1fb..c8efb10 100644 --- a/cache.go +++ b/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. diff --git a/internal_test.go b/internal_test.go index 4c945fc..17275e0 100644 --- a/internal_test.go +++ b/internal_test.go @@ -1,6 +1,7 @@ package starlet import ( + "context" "reflect" "strings" "testing" @@ -8,6 +9,29 @@ import ( "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{ diff --git a/module.go b/module.go index c3497e5..ef7e9fa 100644 --- a/module.go +++ b/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 diff --git a/run.go b/run.go index c8d12c2..a75dd22 100644 --- a/run.go +++ b/run.go @@ -100,11 +100,11 @@ 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{}) @@ -112,7 +112,7 @@ func (m *Machine) watchContextCancel(ctx context.Context) (stop 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", diff --git a/run_test.go b/run_test.go index f0b4a5b..2518a5d 100644 --- a/run_test.go +++ b/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).