gitaly: Buffer command output on context doneps-ctx-done-cmd-buffering

This is an attempt to prevent a failure on reading from the
command output and context cancellation. When the context is done
the termination signal is sent to the command to stop it. The
handling of the signal by the command may delay, so it may
proceed its execution, produce an output and finish with 0 code.
That is why before discarding the full command output we now
store its 1024 bytes and discard anything else because of safety
reasons. This data can now be read from the command by another
goroutine.

Closes: https://gitlab.com/gitlab-org/gitaly/-/issues/5021

1 files changed, 27 insertions, 8 deletions

diff --git a/internal/command/command.go b/internal/command/command.go
index bbea588c8..2be1bdafc 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -10,6 +11,7 @@ import (
 	"path"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -131,7 +133,7 @@ const (
 // terminated and reaped automatically when the context.Context that
 // created it is canceled.
 type Command struct {
-	reader       io.Reader
+	reader       atomic.Pointer[io.ReadCloser]
 	writer       io.WriteCloser
 	stderrBuffer *stderrBuffer
 	cmd          *exec.Cmd
@@ -275,7 +277,7 @@ func New(ctx context.Context, nameAndArgs []string, opts ...Option) (*Command, e
 			return nil, fmt.Errorf("creating stdout pipe: %w", err)
 		}
 
-		command.reader = pipe
+		command.reader.Store(&pipe)
 	}
 
 	if cfg.stderr != nil {
@@ -339,13 +341,16 @@ func New(ctx context.Context, nameAndArgs []string, opts ...Option) (*Command, e
 	return command, nil
 }
 
-// Read calls Read() on the stdout pipe of the command.
+// Read calls Read() on the stdout pipe of the command or temporary buffer
+// that contains max 1024 bytes of the output in case the command execution
+// context is cancelled.
 func (c *Command) Read(p []byte) (int, error) {
-	if c.reader == nil {
+	readerPtr := c.reader.Load()
+	if readerPtr == nil || *readerPtr == nil {
 		panic("command has no reader")
 	}
 
-	return c.reader.Read(p)
+	return (*readerPtr).Read(p)
 }
 
 // Write calls Write() on the stdin pipe of the command.
@@ -379,16 +384,30 @@ func (c *Command) wait() {
 		c.writer.Close()
 	}
 
-	if c.reader != nil {
+	contextIsDone := c.context.Err() != nil
+
+	if readerPtr := c.reader.Load(); readerPtr != nil && *readerPtr != nil {
+		if contextIsDone {
+			// When context is done it is not yet mean the running command should fail.
+			// If the termination error is not yet processed by the command it may produce
+			// useful output that can be processed by the caller. That is why we consume
+			// max up to 1024 bytes from it and discard all the other info because of
+			// safety reasons. The consumed data can be read from the buffer by another
+			// goroutine. All remaining output of the command will be discarded.
+			buffer := &bytes.Buffer{}
+			_, _ = io.Copy(buffer, io.LimitReader(*readerPtr, 1024))
+			tmpBuffer := io.NopCloser(buffer)
+			c.reader.Store(&tmpBuffer)
+		}
 		// Prevent the command from blocking on writing to its stdout.
-		_, _ = io.Copy(io.Discard, c.reader)
+		_, _ = io.Copy(io.Discard, *readerPtr)
 	}
 
 	c.waitError = c.cmd.Wait()
 
 	// If the context is done, the process was likely terminated due to it. If so,
 	// we return the context error to correctly report the reason.
-	if c.context.Err() != nil {
+	if contextIsDone {
 		// The standard library sets exit status -1 if the process was terminated by a signal,
 		// such as the SIGTERM sent when context is done.
 		if exitCode, ok := ExitStatus(c.waitError); ok && exitCode == -1 {