From 10f59ce53951001c41555b3bfc7bb014377cdcbc Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Wed, 11 Jun 2025 21:09:24 +0200 Subject: [PATCH] tart exec support: fix StdoutPipe() and StderrPipe() usage --- go.mod | 1 + go.sum | 2 + internal/rpc/exec.go | 135 ++++++++++++++++++++----------------------- 3 files changed, 65 insertions(+), 73 deletions(-) diff --git a/go.mod b/go.mod index bfe3aa5..b6c0739 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 566c9ee..fbd2dc3 100644 --- a/go.sum +++ b/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= diff --git a/internal/rpc/exec.go b/internal/rpc/exec.go index 55f2d15..ef74656 100644 --- a/internal/rpc/exec.go +++ b/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]), - } - } -}