core/bazel: fix race

core/bazel/BUILD.bazel

@@ -1,3 +1,4 @@
+load("@rules_go//extras:gomock.bzl", "gomock")
 load("@rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
@@ -15,7 +16,6 @@ go_library(
         "@com_github_bazelbuild_buildtools//build_proto",
         "@org_golang_google_protobuf//encoding/protodelim",
         "@org_golang_google_protobuf//proto",
-        "@org_golang_x_sync//errgroup",
         "@org_uber_go_zap//:zap",
     ],
 )
@@ -24,17 +24,27 @@ go_test(
     name = "bazel_test",
     srcs = [
         "bazel_test.go",
+        "command_test.go",
+        "mock_commander_test.go",
         "query_test.go",
     ],
     embed = [":bazel"],
     deps = [
-        "//core/bazel/commandermock",
+        "//core/execcmd",
         "@com_github_bazelbuild_buildtools//build_proto",
+        "@com_github_golang_mock//gomock",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_protobuf//encoding/protodelim",
         "@org_uber_go_goleak//:goleak",
-        "@org_uber_go_mock//gomock",
         "@org_uber_go_zap//:zap",
     ],
 )
+
+gomock(
+    name = "mock_commander",
+    out = "mock_commander_test.go",
+    interfaces = ["Commander"],
+    library = ":bazel",
+    package = "bazel",
+)

core/bazel/bazel.go

@@ -55,8 +55,9 @@ type BazelClient struct {
 	envVarsMap         map[string]string
 	bazelCommand       string
 	logger             *zap.SugaredLogger
-	execCommandContext func(ctx context.Context, name string, arg ...string) commander
+	execCommandContext func(ctx context.Context, name string, arg ...string) Commander
 	queryTimeout       time.Duration
+	tempDir            string
 	streamLogs         bool
 }
 
@@ -65,22 +66,22 @@ type Params struct {
 	WorkspacePath      string
 	EnvVarsMap         map[string]string
 	Logger             *zap.SugaredLogger
-	ExecCommandContext func(ctx context.Context, name string, arg ...string) commander
+	ExecCommandContext func(ctx context.Context, name string, arg ...string) Commander
 	QueryTimeout       time.Duration
 	StreamLogs         bool
 }
 
 func NewBazelClient(ctx context.Context, p Params) (*BazelClient, error) {
 	execCmd := p.ExecCommandContext
 	if execCmd == nil {
-		execCmd = func(ctx context.Context, name string, arg ...string) commander {
+		execCmd = func(ctx context.Context, name string, arg ...string) Commander {
 			cmd := execcmd.CommandContext(ctx, name, arg...)
 			cmd.Dir = p.WorkspacePath
 			for key, value := range p.EnvVarsMap {
 				cmd.Env = append(cmd.Env, key+"="+value)
 			}
 			cmd.Stdin = nil
-			return cmd
+			return &execCommander{Cmd: cmd}
 		}
 	}
 	timeout := p.QueryTimeout
@@ -99,6 +100,7 @@ func NewBazelClient(ctx context.Context, p Params) (*BazelClient, error) {
 		logger:             p.Logger,
 		execCommandContext: execCmd,
 		queryTimeout:       timeout,
+		tempDir:            os.TempDir(),
 		streamLogs:         p.StreamLogs,
 	}, nil
 }

core/bazel/bazel_test.go

@@ -16,7 +16,7 @@ package bazel
 
 import (
 	"context"
-	"os/exec"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -36,7 +36,7 @@ func TestNewBazelClient(t *testing.T) {
 				WorkspacePath: "/tmp/test",
 				EnvVarsMap:    map[string]string{"FOO": "bar"},
 				Logger:        zap.NewNop().Sugar(),
-				ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander {
+				ExecCommandContext: func(ctx context.Context, name string, arg ...string) Commander {
 					return nil
 				},
 			},
@@ -61,6 +61,7 @@ func TestNewBazelClient(t *testing.T) {
 			assert.Equal(t, tt.params.WorkspacePath, client.workspacePath)
 			assert.Equal(t, tt.params.EnvVarsMap, client.envVarsMap)
 			assert.Equal(t, tt.params.Logger, client.logger)
+			assert.Equal(t, os.TempDir(), client.tempDir)
 			assert.NotNil(t, client.execCommandContext)
 		})
 	}
@@ -79,8 +80,10 @@ func TestNewBazelClient_WithNilExecCommand(t *testing.T) {
 
 	cmd := client.execCommandContext(context.Background(), "test", "arg1")
 	require.NotNil(t, cmd)
-	execCmd, ok := cmd.(*exec.Cmd)
+	execCmd, ok := cmd.(*execCommander)
 	require.True(t, ok)
 	assert.Equal(t, "/workspace", execCmd.Dir)
 	assert.Contains(t, execCmd.Env, "KEY=value")
+	assert.NotNil(t, execCmd.Cancel)
+	assert.Positive(t, execCmd.WaitDelay)
 }

core/bazel/command.go

@@ -16,11 +16,24 @@ package bazel
 
 import (
 	"io"
+	"os/exec"
 )
 
-type commander interface {
-	StdoutPipe() (io.ReadCloser, error)
-	StderrPipe() (io.ReadCloser, error)
-	Start() error
-	Wait() error
+// Commander runs a command with the supplied output destinations.
+type Commander interface {
+	Run(stdout, stderr io.Writer) error
+}
+
+type execCommander struct {
+	*exec.Cmd
+}
+
+// Run attaches the supplied writers before starting the command. Query
+// execution intentionally supplies *os.File values: os/exec passes those file
+// descriptors directly to the child instead of creating copy goroutines that
+// Wait must drain.
+func (c *execCommander) Run(stdout, stderr io.Writer) error {
+	c.Stdout = stdout
+	c.Stderr = stderr
+	return c.Cmd.Run()
 }

core/bazel/command_test.go

@@ -0,0 +1,101 @@
+// Copyright (c) 2025 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package bazel
+
+import (
+	"context"
+	"io"
+	"os"
+	"os/exec"
+	"os/signal"
+	"strconv"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/uber/tango/core/execcmd"
+	"go.uber.org/goleak"
+)
+
+func TestExecCommander_OutputFilesDoNotWaitForDescendant(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	stdout, err := os.CreateTemp(t.TempDir(), "stdout-*")
+	require.NoError(t, err)
+	defer stdout.Close()
+	stderr, err := os.CreateTemp(t.TempDir(), "stderr-*")
+	require.NoError(t, err)
+	defer stderr.Close()
+
+	cmd := &execCommander{Cmd: exec.Command("/bin/sh", "-c", "sleep 5 & echo $! >&2")}
+	started := time.Now()
+	require.NoError(t, cmd.Run(stdout, stderr))
+	elapsed := time.Since(started)
+
+	contents, err := os.ReadFile(stderr.Name())
+	require.NoError(t, err)
+	pid, err := strconv.Atoi(strings.TrimSpace(string(contents)))
+	require.NoError(t, err)
+	process, err := os.FindProcess(pid)
+	require.NoError(t, err)
+	_ = process.Kill()
+	require.Less(t, elapsed, 2*time.Second)
+}
+
+func TestExecCommander_ContextCancellation(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	readyReader, readyWriter, err := os.Pipe()
+	require.NoError(t, err)
+	defer readyReader.Close()
+	defer readyWriter.Close()
+
+	cmd := execcmd.CommandContext(ctx, os.Args[0], "-test.run=TestBazelCommandHelper")
+	cmd.Env = append(os.Environ(), "TANGO_BAZEL_COMMAND_HELPER=1")
+	cmd.ExtraFiles = []*os.File{readyWriter}
+	runner := &execCommander{Cmd: cmd}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- runner.Run(io.Discard, io.Discard)
+	}()
+	ready := make([]byte, 1)
+	_, err = io.ReadFull(readyReader, ready)
+	require.NoError(t, err)
+
+	cancel()
+	select {
+	case err := <-done:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(2 * time.Second):
+		t.Fatal("command did not exit after context cancellation")
+	}
+}
+
+func TestBazelCommandHelper(t *testing.T) {
+	if os.Getenv("TANGO_BAZEL_COMMAND_HELPER") != "1" {
+		return
+	}
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGTERM)
+
+	ready := os.NewFile(3, "ready")
+	_, _ = ready.Write([]byte{1})
+	_ = ready.Close()
+
+	<-signals
+}