113 lines
3.0 KiB
Go
113 lines
3.0 KiB
Go
package execute
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
)
|
|
|
|
// runStage executes one stage of an execute_code pipeline, feeding stdinData as stdin.
|
|
// Parallel stages fan out concurrently and concatenate results in submission order.
|
|
func (e *Server) runStage(ctx context.Context, idx int, stage CodeStep, timeout int, sandboxName, stdinData string) (string, error) {
|
|
if len(stage.Parallel) > 0 {
|
|
results := make([]string, len(stage.Parallel))
|
|
errs := make([]error, len(stage.Parallel))
|
|
var wg sync.WaitGroup
|
|
for i, step := range stage.Parallel {
|
|
wg.Add(1)
|
|
go func(j int, s CodeStep) {
|
|
defer wg.Done()
|
|
if s.Elevated {
|
|
e.wdMu.RLock()
|
|
dir := e.cwd
|
|
e.wdMu.RUnlock()
|
|
results[j], errs[j] = e.executeElevated(ctx, s.Code, dir, timeout, s.Detach)
|
|
return
|
|
}
|
|
code, lang, trusted, err := resolveCodeStep(s)
|
|
if err != nil {
|
|
errs[j] = err
|
|
return
|
|
}
|
|
results[j], errs[j] = e.executeWithStdin(ctx, code, lang, timeout, sandboxName, trusted, stdinData)
|
|
}(i, step)
|
|
}
|
|
wg.Wait()
|
|
var out string
|
|
for i, r := range results {
|
|
out += r
|
|
if errs[i] != nil {
|
|
return out, errs[i]
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
if stage.Elevated {
|
|
e.wdMu.RLock()
|
|
dir := e.cwd
|
|
e.wdMu.RUnlock()
|
|
return e.executeElevated(ctx, stage.Code, dir, timeout, stage.Detach)
|
|
}
|
|
|
|
if stage.Tool != "" || stage.Code != "" {
|
|
code, lang, trusted, err := resolveCodeStep(stage)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return e.executeWithStdin(ctx, code, lang, timeout, sandboxName, trusted, stdinData, stage.Detach)
|
|
}
|
|
|
|
return "", fmt.Errorf("stage %d: requires code, tool, or parallel", idx)
|
|
}
|
|
|
|
// runReadBatch runs a contiguous slice of read-safe stages concurrently,
|
|
// concatenating their outputs in submission order. All stages receive the
|
|
// same stdinData (read-safe tools don't depend on stdin chaining).
|
|
// If e.lockDir is set, each goroutine acquires LOCK_SH before running.
|
|
func (e *Server) runReadBatch(ctx context.Context, startIdx int, stages []CodeStep, timeout int, sandboxName, stdinData string) (string, error) {
|
|
if len(stages) == 1 {
|
|
lf, err := acquireFlock(e.lockDir, "rw", false)
|
|
if err != nil {
|
|
return "", fmt.Errorf("step %d: lock: %w", startIdx, err)
|
|
}
|
|
if lf != nil {
|
|
defer lf.Close()
|
|
}
|
|
out, err := e.runStage(ctx, startIdx, stages[0], timeout, sandboxName, stdinData)
|
|
if err != nil {
|
|
return out, fmt.Errorf("step %d: %w", startIdx, err)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
results := make([]string, len(stages))
|
|
errs := make([]error, len(stages))
|
|
var wg sync.WaitGroup
|
|
for i, stage := range stages {
|
|
wg.Add(1)
|
|
go func(j int, s CodeStep) {
|
|
defer wg.Done()
|
|
lf, lerr := acquireFlock(e.lockDir, "rw", false)
|
|
if lerr != nil {
|
|
errs[j] = fmt.Errorf("step %d: lock: %w", startIdx+j, lerr)
|
|
return
|
|
}
|
|
if lf != nil {
|
|
defer lf.Close()
|
|
}
|
|
results[j], errs[j] = e.runStage(ctx, startIdx+j, s, timeout, sandboxName, stdinData)
|
|
}(i, stage)
|
|
}
|
|
wg.Wait()
|
|
|
|
var out string
|
|
for i, r := range results {
|
|
out += r
|
|
if errs[i] != nil {
|
|
return out, errs[i]
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|