fix(executor/sandbox): wait grpc

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(
+	slog.Debug("sandbox stream execute", "pbReq size", proto.Size(pbReq), "tarSize", len(tarData))
-		"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
+	select {
-	for {
+	case finalResponse := <-respCh:
-		resp, err := stream.Recv()
+		results := convertPBResult(finalResponse.GetResults())
-		if err == io.EOF {
+		if len(results) == 0 {
-			break
+			return nil, fmt.Errorf("stream: empty results")
 		}
-		if err != nil {
+		tarResult := results[0]
-			return nil, fmt.Errorf("stream recv: %w", err)
+		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() {
+		results = results[1:]
-			finalResponse = resp.GetExecResponse()
+		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 {