Files
2026-06-22 08:55:57 +08:00

169 lines
4.8 KiB
Go

package proc
import (
"bytes"
"context"
"os/exec"
"time"
)
// ExecSpec 描述一次一次性(one-shot)命令执行的全部输入。
// 调用方(run.ExecService)负责前置鉴权、cwd / env 构造与命令白名单校验。
type ExecSpec struct {
Command string
Args []string
Dir string
Env []string // 完整环境(宿主白名单 + 覆盖),nil 表示空环境而非继承宿主
// Timeout 是硬上限:超时即对整组发起 TerminateGroup(树终止)并置 TimedOut。
// <=0 时不设独立超时,仅受 ctx 约束。
Timeout time.Duration
// KillGrace 是树终止时的优雅宽限(Unix SIGTERM → grace → SIGKILL)。
KillGrace time.Duration
// MaxOutputBytes 是 stdout / stderr 各自的捕获上限;超出截断并置 truncated。
// <=0 时使用 defaultMaxOutputBytes。
MaxOutputBytes int
}
// ExecResult 是一次性执行的结构化结果。
type ExecResult struct {
ExitCode int
Stdout string
Stderr string
Duration time.Duration
TimedOut bool
Truncated bool // stdout 或 stderr 任一被截断时为 true
StdoutSize int // 截断前的实际字节数
StderrSize int // 截断前的实际字节数
}
const defaultMaxOutputBytes = 1 << 20 // 1 MiB
// RunOnce 同步执行 spec 描述的命令直至退出(或超时 / ctx 取消),捕获 stdout /
// stderr(各自封顶 MaxOutputBytes),返回退出码、时长、超时标志。
//
// 进程树管理沿用 RunSupervisor.Spawn 的次序:Prepare(Start 前)→ Start →
// Adopt(Start 后)→ 退出或超时时 TerminateGroup 整树清理。这样测试运行器派生的
// 子孙进程(go test 编译出的子进程、node 子进程等)在超时时也会被一并终止。
//
// 与长驻 profile 的差异:本函数阻塞到进程退出后才返回,不接入 relay;输出在内存
// 缓冲并随返回值一次性交回。错误仅在无法启动进程(exec.Start 失败)时返回;命令
// 以非零码退出不算 error,体现在 ExitCode 上。
func RunOnce(ctx context.Context, spec ExecSpec) (ExecResult, error) {
maxOut := spec.MaxOutputBytes
if maxOut <= 0 {
maxOut = defaultMaxOutputBytes
}
cmd := exec.Command(spec.Command, spec.Args...)
cmd.Dir = spec.Dir
cmd.Env = spec.Env
outBuf := &cappedBuffer{limit: maxOut}
errBuf := &cappedBuffer{limit: maxOut}
cmd.Stdout = outBuf
cmd.Stderr = errBuf
group := NewGroup()
group.Prepare(cmd)
started := time.Now()
if err := cmd.Start(); err != nil {
return ExecResult{}, err
}
if err := group.Adopt(cmd); err != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
group.Close()
return ExecResult{}, err
}
// waitCh 携带 Wait() 的错误;done 在 Wait 返回后 close,供 TerminateGroup
// 感知组长已被回收(语义对齐 RunSupervisor.monitorProcess 的 done channel)。
waitCh := make(chan error, 1)
done := make(chan struct{})
go func() {
err := cmd.Wait()
waitCh <- err
close(done)
}()
// 超时 / ctx-cancel watchdog:触发整树终止;result 由下方 select 收尾。
var timeoutTimer <-chan time.Time
if spec.Timeout > 0 {
t := time.NewTimer(spec.Timeout)
defer t.Stop()
timeoutTimer = t.C
}
var timedOut bool
var waitErr error
loop:
for {
select {
case waitErr = <-waitCh:
break loop
case <-timeoutTimer:
timedOut = true
group.TerminateGroup(done, spec.KillGrace)
waitErr = <-waitCh
break loop
case <-ctx.Done():
group.TerminateGroup(done, spec.KillGrace)
waitErr = <-waitCh
break loop
}
}
group.Close()
res := ExecResult{
ExitCode: exitCodeOf(waitErr),
Stdout: outBuf.String(),
Stderr: errBuf.String(),
Duration: time.Since(started),
TimedOut: timedOut,
Truncated: outBuf.truncated || errBuf.truncated,
StdoutSize: outBuf.total,
StderrSize: errBuf.total,
}
return res, nil
}
// exitCodeOf 从 Cmd.Wait 的返回错误提取退出码。正常退出(nil)= 0;
// *exec.ExitError 取其 ExitCode(被信号杀死时为 -1);其它错误统一 -1。
func exitCodeOf(waitErr error) int {
if waitErr == nil {
return 0
}
if ee, ok := waitErr.(*exec.ExitError); ok {
return ee.ExitCode()
}
return -1
}
// cappedBuffer 是一个带上限的 io.Writer:累计写入 limit 字节后丢弃后续数据并置
// truncated,同时仍记录 total(截断前的总字节数)便于上层提示。
type cappedBuffer struct {
buf bytes.Buffer
limit int
total int
truncated bool
}
func (c *cappedBuffer) Write(p []byte) (int, error) {
c.total += len(p)
if c.buf.Len() >= c.limit {
c.truncated = true
return len(p), nil // 谎报全写入:避免 cmd 因 short write 报错而中断
}
remain := c.limit - c.buf.Len()
if len(p) > remain {
c.buf.Write(p[:remain])
c.truncated = true
return len(p), nil
}
c.buf.Write(p)
return len(p), nil
}
func (c *cappedBuffer) String() string { return c.buf.String() }