diff --git a/internal/executor/sandbox/executor.go b/internal/executor/sandbox/executor.go index ed8ede3..97daa8e 100644 --- a/internal/executor/sandbox/executor.go +++ b/internal/executor/sandbox/executor.go @@ -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 {