fix(executor/sandbox): wait grpc
Some checks failed
submodules sync / sync (push) Successful in 28s
build / build (push) Failing after 1m8s
build / trigger-build-image (push) Has been skipped

This commit is contained in:
张泊明518370910136 2026-07-22 22:15:04 -07:00
parent e897d2a78c
commit b4e94c38f3
GPG Key ID: D47306D7062CDA9D

View File

@ -92,7 +92,6 @@ func (e *Sandbox) runUnary(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) {
tarCmd := stage.Cmd{
// FIXME: `/w` won't work if go-judge change default working directory
Args: []string{"/bin/tar", "xf", "-", "-C", "/w"},
Env: cmds[0].Env,
Stdin: &stage.CmdFile{StreamIn: true},
@ -110,11 +109,7 @@ func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.Ex
pbReq := &pb.Request{}
pbReq.SetCmd(pbCmds)
slog.Debug(
"sandbox stream execute",
"pbReq size", proto.Size(pbReq),
"tarSize", len(tarData),
)
slog.Debug("sandbox stream execute", "pbReq size", proto.Size(pbReq), "tarSize", len(tarData))
stream, err := e.execClient.ExecStream(context.TODO())
if err != nil {
@ -127,11 +122,39 @@ func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.Ex
return nil, fmt.Errorf("stream send request: %w", err)
}
errCh := make(chan error, 1)
respCh := make(chan *pb.Response, 1)
go func() {
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,
@ -140,6 +163,11 @@ func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.Ex
}).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)
}
}
@ -148,44 +176,27 @@ func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.Ex
return nil, fmt.Errorf("stream close send: %w", err)
}
var finalResponse *pb.Response
for {
resp, err := stream.Recv()
if err == io.EOF {
break
select {
case finalResponse := <-respCh:
results := convertPBResult(finalResponse.GetResults())
if len(results) == 0 {
return nil, fmt.Errorf("stream: empty results")
}
if err != nil {
return nil, fmt.Errorf("stream recv: %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,
)
}
if resp.HasExecResponse() {
finalResponse = resp.GetExecResponse()
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)
}
if finalResponse == nil {
return nil, fmt.Errorf("stream: no exec response received")
}
if finalResponse.GetError() != "" {
return nil, fmt.Errorf("stream execute error: %s", finalResponse.GetError())
}
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
}
func estimateCopyInSize(cmd *stage.Cmd) int {