core/server/sandbox/batch_module.go

// 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
	}
}