From d875e5e342b7aa863abd26d32009a2eea0e16b22 Mon Sep 17 00:00:00 2001
From: Ilia Babanov <ilia.babanov@databricks.com>
Date: Mon, 12 Feb 2024 15:26:24 +0100
Subject: [PATCH] Avoid race-conditions while executing sub-commands

`executor.Exec` now uses `cmd.CombinedOutput`. Previous implementation was hanging
on my windows VM during `bundle deploy` on the `ReadAll(MultiReader(stdout, stderr))` line.

Seems to be related to the fact the MultiReader reads sequentially.

Also noticed that Exec was not removing `scriptFile` after itself, fixed that too.
---
 libs/exec/exec.go      | 30 ++++++++++++++++--------------
 libs/exec/exec_test.go | 15 ++++++++++++---
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/libs/exec/exec.go b/libs/exec/exec.go
index 9767c199a2..8e4633271a 100644
--- a/libs/exec/exec.go
+++ b/libs/exec/exec.go
@@ -90,18 +90,25 @@ func NewCommandExecutorWithExecutable(dir string, execType ExecutableType) (*Exe
 	}, nil
 }
 
-func (e *Executor) StartCommand(ctx context.Context, command string) (Command, error) {
+func (e *Executor) prepareCommand(ctx context.Context, command string) (*osexec.Cmd, *execContext, error) {
 	ec, err := e.shell.prepare(command)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return e.start(ctx, ec)
-}
-
-func (e *Executor) start(ctx context.Context, ec *execContext) (Command, error) {
 	cmd := osexec.CommandContext(ctx, ec.executable, ec.args...)
 	cmd.Dir = e.dir
+	return cmd, ec, nil
+}
+
+func (e *Executor) StartCommand(ctx context.Context, command string) (Command, error) {
+	cmd, ec, err := e.prepareCommand(ctx, command)
+	if err != nil {
+		return nil, err
+	}
+	return e.start(ctx, cmd, ec)
+}
 
+func (e *Executor) start(ctx context.Context, cmd *osexec.Cmd, ec *execContext) (Command, error) {
 	stdout, err := cmd.StdoutPipe()
 	if err != nil {
 		return nil, err
@@ -116,17 +123,12 @@ func (e *Executor) start(ctx context.Context, ec *execContext) (Command, error)
 }
 
 func (e *Executor) Exec(ctx context.Context, command string) ([]byte, error) {
-	cmd, err := e.StartCommand(ctx, command)
-	if err != nil {
-		return nil, err
-	}
-
-	res, err := io.ReadAll(io.MultiReader(cmd.Stdout(), cmd.Stderr()))
+	cmd, ec, err := e.prepareCommand(ctx, command)
 	if err != nil {
 		return nil, err
 	}
-
-	return res, cmd.Wait()
+	defer os.Remove(ec.scriptFile)
+	return cmd.CombinedOutput()
 }
 
 func (e *Executor) ShellType() ExecutableType {
diff --git a/libs/exec/exec_test.go b/libs/exec/exec_test.go
index 0730638e3f..ad54601d08 100644
--- a/libs/exec/exec_test.go
+++ b/libs/exec/exec_test.go
@@ -32,6 +32,15 @@ func TestExecutorWithComplexInput(t *testing.T) {
 	assert.Equal(t, "Hello\nWorld\n", string(out))
 }
 
+func TestExecutorWithStderr(t *testing.T) {
+	executor, err := NewCommandExecutor(".")
+	assert.NoError(t, err)
+	out, err := executor.Exec(context.Background(), "echo 'Hello' && >&2 echo 'Error'")
+	assert.NoError(t, err)
+	assert.NotNil(t, out)
+	assert.Equal(t, "Hello\nError\n", string(out))
+}
+
 func TestExecutorWithInvalidCommand(t *testing.T) {
 	executor, err := NewCommandExecutor(".")
 	assert.NoError(t, err)
@@ -108,16 +117,16 @@ func TestExecutorCleanupsTempFiles(t *testing.T) {
 	executor, err := NewCommandExecutor(".")
 	assert.NoError(t, err)
 
-	ec, err := executor.shell.prepare("echo 'Hello'")
+	cmd, ec, err := executor.prepareCommand(context.Background(), "echo 'Hello'")
 	assert.NoError(t, err)
 
-	cmd, err := executor.start(context.Background(), ec)
+	command, err := executor.start(context.Background(), cmd, ec)
 	assert.NoError(t, err)
 
 	fileName := ec.args[1]
 	assert.FileExists(t, fileName)
 
-	err = cmd.Wait()
+	err = command.Wait()
 	assert.NoError(t, err)
 	assert.NoFileExists(t, fileName)
 }