Feat v0.7.12 batch.exec concurrent execution primitive

Add batch.exec(callables, timeout=10) — a general-purpose concurrent execution primitive that runs arbitrary Starlark callables in parallel, each in its own thread with independent step budget. Enables extensions to fan out library function calls (frozen exports from lib.require) without decomposing them back into raw http.post parameters. New permission: batch.exec. Max 8 callables, timeout 1-30s. Nesting prohibited via atomic flag. 12 new tests, all pass with -race. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-02 23:52:01 +00:00
parent c2d52f50c5
commit 8957e99610
7 changed files with 981 additions and 5 deletions
--- a/server/sandbox/batch_module.go
+++ b/server/sandbox/batch_module.go
@@ -0,0 +1,140 @@
+// Package sandbox — batch_module.go
+//
+// Permission: batch.exec
+//
+// Starlark API:
+//   results, errors = batch.exec([fn_a, fn_b, fn_c], timeout=10)
+//
+// Runs a list of zero-arg callables concurrently, each in its own
+// starlark.Thread with an independent step budget. Results and errors
+// are returned in input order. Max 8 callables per call.
+//
+// Each branch gets a fresh Sandbox (thread + step counter) and a
+// child context with the per-branch timeout. The callable's closed-over
+// bindings (frozen library structs, module builtins) are used as-is —
+// the Go resources underneath (*sql.DB, http.Client) are concurrent-safe.
+//
+// lib.require() is not available inside branches (Option A).
+// print() output from branches is discarded.
+// Nested batch.exec() is prohibited via an atomic flag.
+package sandbox
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
+)
+
+const (
+	batchMaxCallables   = 8
+	batchDefaultTimeout = 10
+	batchMaxTimeout     = 30
+)
+
+// BuildBatchModule creates the "batch" module.
+// The runner/packageID/manifest/rc/lc parameters are stored for future
+// expansion (Option B: per-branch lib.require) but are not used in v1.
+func BuildBatchModule(
+	ctx context.Context,
+	runner *Runner,
+	packageID string,
+	manifest map[string]any,
+	rc *RunContext,
+	lc *libContext,
+) *starlarkstruct.Module {
+	var running atomic.Bool
+
+	return MakeModule("batch", starlark.StringDict{
+		"exec": starlark.NewBuiltin("batch.exec", batchExec(ctx, &running)),
+	})
+}
+
+// batchExec returns the batch.exec builtin implementation.
+func batchExec(
+	ctx context.Context,
+	running *atomic.Bool,
+) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
+	return func(
+		thread *starlark.Thread, b *starlark.Builtin,
+		args starlark.Tuple, kwargs []starlark.Tuple,
+	) (starlark.Value, error) {
+		// Nesting guard: parent thread is blocked in wg.Wait(), so only
+		// a branch goroutine could re-enter. The captured builtin reference
+		// makes this reachable even though branches don't get fresh modules.
+		if !running.CompareAndSwap(false, true) {
+			return nil, fmt.Errorf("batch.exec: nested calls are not allowed")
+		}
+		defer running.Store(false)
+
+		var callableList *starlark.List
+		var timeout int = batchDefaultTimeout
+
+		if err := starlark.UnpackArgs(b.Name(), args, kwargs,
+			"callables", &callableList,
+			"timeout?", &timeout,
+		); err != nil {
+			return nil, err
+		}
+
+		n := callableList.Len()
+		if n == 0 {
+			return nil, fmt.Errorf("batch.exec: callables list is empty")
+		}
+		if n > batchMaxCallables {
+			return nil, fmt.Errorf("batch.exec: max %d callables, got %d", batchMaxCallables, n)
+		}
+		if timeout < 1 || timeout > batchMaxTimeout {
+			timeout = batchDefaultTimeout
+		}
+
+		// Validate all elements are callable before spawning goroutines.
+		callables := make([]starlark.Callable, n)
+		for i := 0; i < n; i++ {
+			c, ok := callableList.Index(i).(starlark.Callable)
+			if !ok {
+				return nil, fmt.Errorf("batch.exec: element %d is %s, not callable",
+					i, callableList.Index(i).Type())
+			}
+			callables[i] = c
+		}
+
+		// Execute concurrently.
+		results := make([]starlark.Value, n)
+		errors := make([]starlark.Value, n)
+		var wg sync.WaitGroup
+
+		for i, callable := range callables {
+			wg.Add(1)
+			go func(idx int, fn starlark.Callable) {
+				defer wg.Done()
+
+				branchCtx, cancel := context.WithTimeout(ctx,
+					time.Duration(timeout)*time.Second)
+				defer cancel()
+
+				sb := New(DefaultConfig())
+				val, _, callErr := sb.Call(branchCtx, fn, nil, nil)
+
+				if callErr != nil {
+					results[idx] = starlark.None
+					errors[idx] = starlark.String(callErr.Error())
+				} else {
+					results[idx] = val
+					errors[idx] = starlark.None
+				}
+			}(i, callable)
+		}
+
+		wg.Wait()
+
+		return starlark.Tuple{
+			starlark.NewList(results),
+			starlark.NewList(errors),
+		}, nil
+	}
+}
--- a/server/sandbox/batch_module_test.go
+++ b/server/sandbox/batch_module_test.go
@@ -0,0 +1,324 @@
+package sandbox
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"go.starlark.net/starlark"
+)
+
+// execWithBatch runs a Starlark script with the batch module available.
+func execWithBatch(t *testing.T, script string) (*Result, error) {
+	t.Helper()
+	ctx := context.Background()
+	batchMod := BuildBatchModule(ctx, nil, "test-pkg", nil, nil, nil)
+	sb := New(DefaultConfig())
+	return sb.Exec(ctx, "test.star", script, map[string]starlark.Value{
+		"batch": batchMod,
+	})
+}
+
+// ─── Parallel Ordering ──────────────────────
+
+func TestBatchExec_ParallelOrdering(t *testing.T) {
+	result, err := execWithBatch(t, `
+results, errors = batch.exec([
+    lambda: 10,
+    lambda: 20,
+    lambda: 30,
+])
+`)
+	if err != nil {
+		t.Fatalf("exec error: %v", err)
+	}
+
+	// Verify results are in input order.
+	globals := result.Globals
+	results := globals["results"].(*starlark.List)
+	errors := globals["errors"].(*starlark.List)
+
+	if results.Len() != 3 {
+		t.Fatalf("expected 3 results, got %d", results.Len())
+	}
+
+	for i, want := range []int{10, 20, 30} {
+		v, _ := starlark.AsInt32(results.Index(i))
+		got := int(v)
+		if got != want {
+			t.Errorf("results[%d] = %d, want %d", i, got, want)
+		}
+		if errors.Index(i) != starlark.None {
+			t.Errorf("errors[%d] = %v, want None", i, errors.Index(i))
+		}
+	}
+}
+
+// ─── Partial Failure ────────────────────────
+
+func TestBatchExec_PartialFailure(t *testing.T) {
+	result, err := execWithBatch(t, `
+def ok_a():
+    return "a"
+
+def bad():
+    return 1 // 0
+
+def ok_c():
+    return "c"
+
+results, errors = batch.exec([ok_a, bad, ok_c])
+`)
+	if err != nil {
+		t.Fatalf("exec error: %v", err)
+	}
+
+	results := result.Globals["results"].(*starlark.List)
+	errors := result.Globals["errors"].(*starlark.List)
+
+	// results[0] = "a", results[2] = "c"
+	if s, _ := starlark.AsString(results.Index(0)); s != "a" {
+		t.Errorf("results[0] = %v, want 'a'", results.Index(0))
+	}
+	if s, _ := starlark.AsString(results.Index(2)); s != "c" {
+		t.Errorf("results[2] = %v, want 'c'", results.Index(2))
+	}
+
+	// results[1] = None (failed)
+	if results.Index(1) != starlark.None {
+		t.Errorf("results[1] = %v, want None", results.Index(1))
+	}
+
+	// errors[0] and errors[2] = None (succeeded)
+	if errors.Index(0) != starlark.None {
+		t.Errorf("errors[0] = %v, want None", errors.Index(0))
+	}
+	if errors.Index(2) != starlark.None {
+		t.Errorf("errors[2] = %v, want None", errors.Index(2))
+	}
+
+	// errors[1] should be a string containing the error
+	errStr, ok := starlark.AsString(errors.Index(1))
+	if !ok || errStr == "" {
+		t.Errorf("errors[1] = %v, want non-empty error string", errors.Index(1))
+	}
+}
+
+// ─── Timeout Per-Branch ─────────────────────
+
+func TestBatchExec_TimeoutPerBranch(t *testing.T) {
+	// Use a tight timeout. The slow lambda burns steps in an infinite loop,
+	// which will hit either the step limit or the context timeout.
+	result, err := execWithBatch(t, `
+def fast():
+    return "done"
+
+def slow():
+    x = 0
+    for i in range(999999999):
+        x = x + 1
+    return x
+
+results, errors = batch.exec([fast, slow], timeout=1)
+`)
+	if err != nil {
+		t.Fatalf("exec error: %v", err)
+	}
+
+	results := result.Globals["results"].(*starlark.List)
+	errors := result.Globals["errors"].(*starlark.List)
+
+	// Fast branch should succeed.
+	if s, _ := starlark.AsString(results.Index(0)); s != "done" {
+		t.Errorf("results[0] = %v, want 'done'", results.Index(0))
+	}
+	if errors.Index(0) != starlark.None {
+		t.Errorf("errors[0] = %v, want None", errors.Index(0))
+	}
+
+	// Slow branch should fail (step limit or timeout).
+	if results.Index(1) != starlark.None {
+		t.Errorf("results[1] = %v, want None", results.Index(1))
+	}
+	errStr, _ := starlark.AsString(errors.Index(1))
+	if errStr == "" {
+		t.Errorf("errors[1] should contain error string")
+	}
+}
+
+// ─── Parent Cancellation ────────────────────
+
+func TestBatchExec_ParentCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+
+	batchMod := BuildBatchModule(ctx, nil, "test-pkg", nil, nil, nil)
+	sb := New(DefaultConfig())
+	result, err := sb.Exec(ctx, "test.star", `
+results, errors = batch.exec([lambda: 1, lambda: 2])
+`, map[string]starlark.Value{"batch": batchMod})
+
+	// Either the Exec itself fails (context cancelled before script runs)
+	// or the batch.exec branches all fail.
+	if err != nil {
+		// Script-level cancellation — acceptable.
+		return
+	}
+
+	errors := result.Globals["errors"].(*starlark.List)
+	allErrored := true
+	for i := 0; i < errors.Len(); i++ {
+		if errors.Index(i) == starlark.None {
+			allErrored = false
+		}
+	}
+	if !allErrored {
+		t.Errorf("expected all branches to error on cancelled context, errors=%v", errors)
+	}
+}
+
+// ─── Cap Enforcement ────────────────────────
+
+func TestBatchExec_CapEnforcement(t *testing.T) {
+	_, err := execWithBatch(t, `
+fns = [lambda: i for i in range(9)]
+batch.exec(fns)
+`)
+	if err == nil {
+		t.Fatal("expected error for 9 callables, got nil")
+	}
+	if !strings.Contains(err.Error(), "max 8") {
+		t.Errorf("error = %q, want to contain 'max 8'", err.Error())
+	}
+}
+
+// ─── Empty List ─────────────────────────────
+
+func TestBatchExec_EmptyList(t *testing.T) {
+	_, err := execWithBatch(t, `batch.exec([])`)
+	if err == nil {
+		t.Fatal("expected error for empty list, got nil")
+	}
+	if !strings.Contains(err.Error(), "empty") {
+		t.Errorf("error = %q, want to contain 'empty'", err.Error())
+	}
+}
+
+// ─── Non-Callable Element ───────────────────
+
+func TestBatchExec_NonCallableElement(t *testing.T) {
+	_, err := execWithBatch(t, `batch.exec([1, 2])`)
+	if err == nil {
+		t.Fatal("expected error for non-callable, got nil")
+	}
+	if !strings.Contains(err.Error(), "not callable") {
+		t.Errorf("error = %q, want to contain 'not callable'", err.Error())
+	}
+}
+
+// ─── Permission Gating ──────────────────────
+
+func TestBatchExec_PermissionGating(t *testing.T) {
+	sb := New(DefaultConfig())
+	_, err := sb.Exec(context.Background(), "test.star",
+		`batch.exec([lambda: 1])`,
+		map[string]starlark.Value{}, // no batch module
+	)
+	if err == nil {
+		t.Fatal("expected error when batch module not available")
+	}
+	if !strings.Contains(err.Error(), "batch") {
+		t.Errorf("error = %q, want to mention 'batch'", err.Error())
+	}
+}
+
+// ─── Frozen Struct Sharing ──────────────────
+
+func TestBatchExec_FrozenStructSharing(t *testing.T) {
+	// Two branches read the same frozen tuple. Should not race.
+	result, err := execWithBatch(t, `
+data = (1, 2, 3)  # tuples are frozen/immutable
+
+results, errors = batch.exec([
+    lambda: data[0] + data[1],
+    lambda: data[1] + data[2],
+])
+`)
+	if err != nil {
+		t.Fatalf("exec error: %v", err)
+	}
+
+	results := result.Globals["results"].(*starlark.List)
+
+	got0, _ := starlark.AsInt32(results.Index(0))
+	got1, _ := starlark.AsInt32(results.Index(1))
+	if got0 != 3 {
+		t.Errorf("results[0] = %d, want 3", got0)
+	}
+	if got1 != 5 {
+		t.Errorf("results[1] = %d, want 5", got1)
+	}
+}
+
+// ─── Nested batch.exec ──────────────────────
+
+func TestBatchExec_NestedBatchExec(t *testing.T) {
+	result, err := execWithBatch(t, `
+results, errors = batch.exec([
+    lambda: batch.exec([lambda: 1]),
+])
+`)
+	if err != nil {
+		t.Fatalf("exec error: %v", err)
+	}
+
+	// The inner batch.exec should fail with nesting error.
+	errors := result.Globals["errors"].(*starlark.List)
+	errStr, _ := starlark.AsString(errors.Index(0))
+	if !strings.Contains(errStr, "nested") {
+		t.Errorf("errors[0] = %q, want to contain 'nested'", errStr)
+	}
+}
+
+// ─── Default Timeout ────────────────────────
+
+func TestBatchExec_DefaultTimeout(t *testing.T) {
+	// Call without explicit timeout — should succeed.
+	result, err := execWithBatch(t, `
+results, errors = batch.exec([lambda: 42])
+`)
+	if err != nil {
+		t.Fatalf("exec error: %v", err)
+	}
+
+	results := result.Globals["results"].(*starlark.List)
+	got, _ := starlark.AsInt32(results.Index(0))
+	if got != 42 {
+		t.Errorf("results[0] = %d, want 42", got)
+	}
+}
+
+// ─── Timeout Clamping ───────────────────────
+
+func TestBatchExec_TimeoutClamp(t *testing.T) {
+	// timeout=0 should be clamped to default, not crash.
+	start := time.Now()
+	result, err := execWithBatch(t, `
+results, errors = batch.exec([lambda: "ok"], timeout=0)
+`)
+	elapsed := time.Since(start)
+	if err != nil {
+		t.Fatalf("exec error: %v", err)
+	}
+
+	results := result.Globals["results"].(*starlark.List)
+	if s, _ := starlark.AsString(results.Index(0)); s != "ok" {
+		t.Errorf("results[0] = %v, want 'ok'", results.Index(0))
+	}
+
+	// Should complete quickly (clamped to default 10s, but lambda is instant).
+	if elapsed > 2*time.Second {
+		t.Errorf("took %v, expected fast completion", elapsed)
+	}
+}
--- a/server/sandbox/runner.go
+++ b/server/sandbox/runner.go
@@ -338,6 +338,7 @@ func (r *Runner) buildModulesWithLibCtx(ctx context.Context, packageID string, m

 	// Track db permission level: 0=none, 1=read, 2=write
 	dbLevel := 0
+	hasBatchExec := false

 	for _, perm := range granted {
 		switch perm {
@@ -374,6 +375,9 @@ func (r *Runner) buildModulesWithLibCtx(ctx context.Context, packageID string, m
 			if r.bus != nil {
 				modules["realtime"] = BuildRealtimeModule(ctx, r.bus, packageID)
 			}
+
+		case models.ExtPermBatchExec:
+			hasBatchExec = true
 		}
 	}

@@ -404,5 +408,12 @@ func (r *Runner) buildModulesWithLibCtx(ctx context.Context, packageID string, m
 		modules["lib"] = BuildLibModule(ctx, r, packageID, rc, lc)
 	}

+	// Concurrent execution primitive — wired after all other modules
+	// so the runner reference can construct per-branch module sets if
+	// needed in future (Option B).
+	if hasBatchExec {
+		modules["batch"] = BuildBatchModule(ctx, r, packageID, manifest, rc, lc)
+	}
+
 	return modules, nil
 }