tart exec support: fix StdoutPipe() and StderrPipe() usage

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/stretchr/testify v1.10.0
 	go.uber.org/zap v1.27.0
 	golang.design/x/clipboard v0.7.0
+	golang.org/x/sync v0.12.0
 	golang.org/x/sys v0.32.0
 	google.golang.org/grpc v1.72.1
 	google.golang.org/protobuf v1.36.6

go.sum

@@ -80,6 +80,8 @@ golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.12.0 h1:MHc5BpPuC30uJk597Ri8TV3CNZcTLu6B6z4lJy+g6Jw=
+golang.org/x/sync v0.12.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

internal/rpc/exec.go

@@ -7,6 +7,7 @@ import (
 	"github.com/creack/pty"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"io"
 	"os"
@@ -21,11 +22,6 @@ const (
 	eofChar = 0x04
 )
 
-type standardStreamOutput struct {
-	Data []byte
-	Err  error
-}
-
 func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) error {
 	// Read the first exec request, it should describe a command to execute
 	firstExecRequest, err := stream.Recv()
@@ -154,74 +150,88 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse])
 		}
 	}()
 
+	group, _ := errgroup.WithContext(stream.Context())
+
 	// Handle standard output from the command
-	stdoutOutputCh := make(chan *standardStreamOutput, 1)
+	group.Go(func() error {
+		buf := make([]byte, standardStreamsBufferSize)
 
-	go streamStandardStream(stdout, stdoutOutputCh)
+		for {
+			n, err := stdout.Read(buf)
+			if err != nil {
+				if errors.Is(err, io.EOF) {
+					return nil
+				}
 
-	// Handle standard error from the command
-	//
-	// Note that it makes no sense to handle standard error when TTY is requested
-	// because in this case stdout and stderr will point to the same file descriptor
-	stderrOutputCh := make(chan *standardStreamOutput, 1)
-
-	if !firstExecRequestCommand.Command.Tty {
-		go streamStandardStream(stderr, stderrOutputCh)
-	}
-
-	// Wait for the command to finish
-	commandErrCh := make(chan error, 1)
-	go func() {
-		commandErrCh <- cmd.Wait()
-	}()
-
-	for {
-		select {
-		case stdoutOutput := <-stdoutOutputCh:
-			if err := stdoutOutput.Err; err != nil {
 				return err
 			}
 
 			if err := stream.Send(&ExecResponse{
 				Type: &ExecResponse_StandardOutput{
 					StandardOutput: &IOChunk{
-						Data: stdoutOutput.Data,
+						Data: slices.Clone(buf[:n]),
 					},
 				},
 			}); err != nil {
 				return err
 			}
-		case stderrOutput := <-stderrOutputCh:
-			if err := stderrOutput.Err; err != nil {
-				return err
-			}
+		}
+	})
 
-			if err := stream.Send(&ExecResponse{
-				Type: &ExecResponse_StandardError{
-					StandardError: &IOChunk{
-						Data: stderrOutput.Data,
+	// Handle standard error from the command
+	//
+	// Note that it makes no sense to handle standard error when TTY is requested
+	// because in this case stdout and stderr will point to the same file descriptor
+	if !firstExecRequestCommand.Command.Tty {
+		group.Go(func() error {
+			buf := make([]byte, standardStreamsBufferSize)
+
+			for {
+				n, err := stderr.Read(buf)
+				if err != nil {
+					if errors.Is(err, io.EOF) {
+						return nil
+					}
+
+					return err
+				}
+
+				if err := stream.Send(&ExecResponse{
+					Type: &ExecResponse_StandardError{
+						StandardError: &IOChunk{
+							Data: slices.Clone(buf[:n]),
+						},
 					},
-				},
-			}); err != nil {
-				return err
+				}); err != nil {
+					return err
+				}
 			}
-		case commandErr := <-commandErrCh:
-			exitCode := 0
+		})
+	}
 
-			var exitError *exec.ExitError
-			if errors.As(commandErr, &exitError) {
-				exitCode = exitError.ExitCode()
-			}
+	if err := group.Wait(); err != nil {
+		return err
+	}
 
-			return stream.Send(&ExecResponse{
-				Type: &ExecResponse_Exit_{
-					Exit: &ExecResponse_Exit{
-						Code: int32(exitCode),
-					},
-				},
-			})
+	// Wait for the command to finish
+	exitCode := 0
+
+	if err := cmd.Wait(); err != nil {
+		var exitError *exec.ExitError
+		if errors.As(err, &exitError) {
+			exitCode = exitError.ExitCode()
+		} else {
+			return err
 		}
 	}
+
+	return stream.Send(&ExecResponse{
+		Type: &ExecResponse_Exit_{
+			Exit: &ExecResponse_Exit{
+				Code: int32(exitCode),
+			},
+		},
+	})
 }
 
 func formatCommandAndArgs(name string, args []string) string {
@@ -236,24 +246,3 @@ func formatCommandAndArgs(name string, args []string) string {
 
 	return fmt.Sprintf("[%s]", strings.Join(all, ", "))
 }
-
-func streamStandardStream(standardStream io.Reader, outputCh chan *standardStreamOutput) {
-	buf := make([]byte, standardStreamsBufferSize)
-
-	for {
-		n, err := standardStream.Read(buf)
-		if err != nil {
-			if !errors.Is(err, io.EOF) {
-				outputCh <- &standardStreamOutput{
-					Err: err,
-				}
-			}
-
-			return
-		}
-
-		outputCh <- &standardStreamOutput{
-			Data: slices.Clone(buf[:n]),
-		}
-	}
-}