fix(executor/sandbox): wait grpc
This commit is contained in:
parent
e897d2a78c
commit
b4e94c38f3
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user