feat(executor/sandbox): add file with grpc
Some checks failed
submodules sync / sync (push) Successful in 48s
build / build (push) Failing after 1m13s
build / trigger-build-image (push) Has been skipped

This commit is contained in:
张泊明518370910136 2026-07-22 22:43:42 -07:00
parent b4e94c38f3
commit 3b112e0219
GPG Key ID: D47306D7062CDA9D

View File

@ -16,10 +16,7 @@ import (
"google.golang.org/protobuf/proto" "google.golang.org/protobuf/proto"
) )
const ( const tarThreshold = 128 * 1024
tarStreamThreshold = 128 * 1024
streamChunkSize = 32 * 1024
)
func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
var err error var err error
@ -30,7 +27,6 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
return nil, err return nil, err
} }
} }
// cannot use range loop since we need to change the value
for i := 0; i < len(cmds); i += 1 { for i := 0; i < len(cmds); i += 1 {
cmd := &cmds[i] cmd := &cmds[i]
if cmd.CopyIn == nil { if cmd.CopyIn == nil {
@ -42,19 +38,19 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
} }
} }
} }
if needStream, tarData := prepareTarStream(cmds); needStream { if needTar, tarData := prepareTar(cmds); needTar {
return e.runWithTarStream(cmds, tarData) return e.runWithTar(cmds, tarData)
} }
return e.runUnary(cmds) return e.runUnary(cmds)
} }
func prepareTarStream(cmds []stage.Cmd) (bool, []byte) { func prepareTar(cmds []stage.Cmd) (bool, []byte) {
if len(cmds) == 0 { if len(cmds) == 0 {
return false, nil return false, nil
} }
for i := range cmds { for i := range cmds {
if cmds[i].CopyInDir != "" && if cmds[i].CopyInDir != "" &&
estimateCopyInSize(&cmds[i]) >= tarStreamThreshold { estimateCopyInSize(&cmds[i]) >= tarThreshold {
tarData, keysInTar := createCopyInTar(&cmds[i]) tarData, keysInTar := createCopyInTar(&cmds[i])
for j := range cmds { for j := range cmds {
cmds[j].CopyInDir = "" cmds[j].CopyInDir = ""
@ -62,6 +58,9 @@ func prepareTarStream(cmds []stage.Cmd) (bool, []byte) {
delete(cmds[j].CopyIn, k) delete(cmds[j].CopyIn, k)
} }
} }
if tarData == nil {
return false, nil
}
return true, tarData return true, tarData
} }
} }
@ -90,11 +89,20 @@ func (e *Sandbox) runUnary(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
return results, nil return results, nil
} }
func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) { func (e *Sandbox) runWithTar(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) {
fc := &pb.FileContent{}
fc.SetContent(tarData)
fileIDResp, err := e.execClient.FileAdd(context.TODO(), fc)
if err != nil {
return nil, fmt.Errorf("file add tar: %w", err)
}
fid := fileIDResp.GetFileID()
slog.Debug("sandbox tar uploaded", "fileID", fid, "tarSize", len(tarData))
tarFileName := "/w/__joj3_copyin.tar"
tarCmd := stage.Cmd{ tarCmd := stage.Cmd{
Args: []string{"/bin/tar", "xf", "-", "-C", "/w"}, Args: []string{"/bin/tar", "xf", tarFileName, "-C", "/w"},
Env: cmds[0].Env, Env: cmds[0].Env,
Stdin: &stage.CmdFile{StreamIn: true},
CPULimit: cmds[0].CPULimit, CPULimit: cmds[0].CPULimit,
ClockLimit: cmds[0].ClockLimit, ClockLimit: cmds[0].ClockLimit,
MemoryLimit: cmds[0].MemoryLimit, MemoryLimit: cmds[0].MemoryLimit,
@ -102,6 +110,10 @@ func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.Ex
ProcLimit: cmds[0].ProcLimit, ProcLimit: cmds[0].ProcLimit,
CPURateLimit: cmds[0].CPURateLimit, CPURateLimit: cmds[0].CPURateLimit,
CPUSetLimit: cmds[0].CPUSetLimit, CPUSetLimit: cmds[0].CPUSetLimit,
CopyIn: map[string]stage.CmdFile{
tarFileName: {FileID: &fid},
},
CopyOutCached: []string{tarFileName},
} }
allCmds := append([]stage.Cmd{tarCmd}, cmds...) allCmds := append([]stage.Cmd{tarCmd}, cmds...)
@ -109,94 +121,40 @@ func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.Ex
pbReq := &pb.Request{} pbReq := &pb.Request{}
pbReq.SetCmd(pbCmds) pbReq.SetCmd(pbCmds)
slog.Debug("sandbox stream execute", "pbReq size", proto.Size(pbReq), "tarSize", len(tarData)) slog.Debug("sandbox tar exec", "pbReq size", proto.Size(pbReq), "tarSize", len(tarData))
stream, err := e.execClient.ExecStream(context.TODO()) pbRet, err := e.execClient.Exec(context.TODO(), pbReq)
if err != nil { if err != nil {
return nil, fmt.Errorf("exec stream: %w", err) return nil, fmt.Errorf("tar exec: %w", err)
}
if pbRet.GetError() != "" {
return nil, fmt.Errorf("sandbox execute error: %s", pbRet.GetError())
} }
sr := &pb.StreamRequest{} results := convertPBResult(pbRet.GetResults())
sr.SetExecRequest(pbReq) if len(results) == 0 {
if err := stream.Send(sr); err != nil { return nil, fmt.Errorf("tar exec: empty results")
return nil, fmt.Errorf("stream send request: %w", err) }
tarResult := results[0]
if tarResult.Status != stage.StatusAccepted || tarResult.ExitStatus != 0 {
return nil, fmt.Errorf(
"tar extraction failed: status=%v, exit=%d, error=%s",
tarResult.Status, tarResult.ExitStatus, tarResult.Error,
)
} }
errCh := make(chan error, 1) results = results[1:]
respCh := make(chan *pb.Response, 1) for _, result := range results {
go func() { maps.Copy(e.cachedMap, result.FileIDs)
for {
resp, recvErr := stream.Recv()
if recvErr == io.EOF {
return
}
if recvErr != nil {
errCh <- recvErr
return
}
if resp.HasExecResponse() {
er := resp.GetExecResponse()
if er.GetError() != "" {
errCh <- fmt.Errorf("server error: %s", er.GetError())
return
}
respCh <- er
}
}
}()
for offset := 0; offset < len(tarData); offset += streamChunkSize {
end := offset + streamChunkSize
if end > len(tarData) {
end = len(tarData)
}
select {
case e := <-errCh:
return nil, fmt.Errorf("stream send input: %w", e)
default:
}
input := pb.StreamRequest_builder{
ExecInput: (&pb.StreamRequest_Input_builder{
Index: 0,
Fd: 0,
Content: tarData[offset:end],
}).Build(),
}.Build()
if err := stream.Send(input); err != nil {
select {
case e := <-errCh:
return nil, fmt.Errorf("stream send input: %w (server: %v)", err, e)
default:
}
return nil, fmt.Errorf("stream send input: %w", err)
}
} }
if err := stream.CloseSend(); err != nil { deleteReq := &pb.FileID{}
return nil, fmt.Errorf("stream close send: %w", err) deleteReq.SetFileID(fid)
if _, err := e.execClient.FileDelete(context.TODO(), deleteReq); err != nil {
slog.Warn("sandbox tar file delete", "fileID", fid, "error", err)
} }
select { return results, nil
case finalResponse := <-respCh:
results := convertPBResult(finalResponse.GetResults())
if len(results) == 0 {
return nil, fmt.Errorf("stream: empty results")
}
tarResult := results[0]
if tarResult.Status != stage.StatusAccepted || tarResult.ExitStatus != 0 {
return nil, fmt.Errorf(
"tar extraction failed: status=%v, exit=%d, error=%s",
tarResult.Status, tarResult.ExitStatus, tarResult.Error,
)
}
results = results[1:]
for _, result := range results {
maps.Copy(e.cachedMap, result.FileIDs)
}
return results, nil
case e := <-errCh:
return nil, fmt.Errorf("stream recv: %w", e)
}
} }
func estimateCopyInSize(cmd *stage.Cmd) int { func estimateCopyInSize(cmd *stage.Cmd) int {