core/bazel: fix race

core/bazel/bazel.go

@@ -32,6 +32,13 @@ import (
 const (
 	// default query timeout if not provided in config
 	_queryTimeout = 15 * time.Minute
+	// _pipeUnblockDelay is how long after the command context ends a query
+	// waits before force-closing the stdout/stderr pipes to unblock stream
+	// reads that never saw EOF (a descendant of the killed process can inherit
+	// the pipe write ends and hold them open, see go.dev/issue/23019). It
+	// exceeds execcmd.GracePeriod so the SIGTERM→SIGKILL sequence gets a
+	// chance to end the streams with a natural EOF first.
+	_pipeUnblockDelay = execcmd.GracePeriod + 5*time.Second
 )
 
 type QueryRequest struct {
@@ -56,6 +63,7 @@ type BazelClient struct {
 	logger             *zap.SugaredLogger
 	execCommandContext func(ctx context.Context, name string, arg ...string) commander
 	queryTimeout       time.Duration
+	pipeUnblockDelay   time.Duration
 	streamLogs         bool
 }
 
@@ -98,6 +106,7 @@ func NewBazelClient(ctx context.Context, p Params) (*BazelClient, error) {
 		logger:             p.Logger,
 		execCommandContext: execCmd,
 		queryTimeout:       timeout,
+		pipeUnblockDelay:   _pipeUnblockDelay,
 		streamLogs:         p.StreamLogs,
 	}, nil
 }

core/bazel/query.go

@@ -22,6 +22,7 @@ import (
 	"io"
 	"os"
 	"strings"
+	"time"
 
 	buildpb "github.com/bazelbuild/buildtools/build_proto"
 	"go.uber.org/zap"
@@ -74,8 +75,6 @@ func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, st
 	if b.streamLogs {
 		stderrSink = os.Stderr
 	}
-	// orchestrate `allOfFailFast`
-	// create a `g` group and a new `gCtx` derived from our 15 minute timeout `ctx`.
 	g, gCtx := errgroup.WithContext(cmdCtx)
 
 	// Start the process
@@ -84,16 +83,47 @@ func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, st
 	}
 	// stream and parse targets
 	g.Go(func() error {
-		var err error
-		queryResults, err = streamAndParseTargets(gCtx, stdout, &stdoutBuf)
-		return err
+		res, err := getQueryResult(gCtx, stdout, &stdoutBuf)
+		if err != nil {
+			return err
+		}
+		queryResults = res
+		return nil
 	})
 	// stream stderr
 	g.Go(func() error {
 		return streamOutput(gCtx, stderr, stderrSink)
 	})
-	waitErr := cmd.Wait()
+
+	// The stream reads normally end with EOF when the process dies and the
+	// kernel closes the pipe write ends. But a descendant of a killed process
+	// that inherited the pipes can hold the write ends open indefinitely
+	// (go.dev/issue/23019). cmd.Wait would force-close the parent ends, but we
+	// only call Wait after the reads finish, so this watchdog is what
+	// guarantees the reads can't block forever.
+	readsDone := make(chan struct{})
+	go func() {
+		select {
+		case <-readsDone:
+			return
+		case <-cmdCtx.Done():
+		}
+		timer := time.NewTimer(b.pipeUnblockDelay)
+		defer timer.Stop()
+		select {
+		case <-readsDone:
+		case <-timer.C:
+			stdout.Close()
+			stderr.Close()
+		}
+	}()
+
+	// Drain the streams before reaping the process: cmd.Wait closes the parent
+	// pipe ends as soon as the process exits, which would truncate whatever
+	// the readers hadn't consumed yet and race on the output buffers.
 	streamErr := g.Wait()
+	close(readsDone)
+	waitErr := cmd.Wait()
 	if waitErr != nil {
 		return queryResults, b.wrapQueryFailure("bazel query failed", waitErr, &stderrBuf)
 	}

core/bazel/query_test.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"strings"
 	"testing"
@@ -175,6 +176,174 @@ func TestExecuteQueryInternal_ContextTimeout(t *testing.T) {
 	assert.Contains(t, err.Error(), "deadline exceeded")
 }
 
+// TestExecuteQueryInternal_PipesHeldOpenAfterCancel simulates a descendant of
+// the killed bazel process inheriting the stdout/stderr pipes and holding
+// their write ends open past the kill (go.dev/issue/23019). Since cmd.Wait —
+// which would force-close the parent ends — only runs after the stream reads
+// finish, the pipe watchdog must unblock the reads or the query would hang.
+func TestExecuteQueryInternal_PipesHeldOpenAfterCancel(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	ctrl := gomock.NewController(t)
+	mockCmd := commandermock.NewMockcommander(ctrl)
+
+	// The write ends are intentionally never closed by the "process".
+	prStdout, pwStdout := io.Pipe()
+	prStderr, pwStderr := io.Pipe()
+	defer pwStdout.Close()
+	defer pwStderr.Close()
+
+	var cmdCtx context.Context
+	gomock.InOrder(
+		mockCmd.EXPECT().StdoutPipe().Return(prStdout, nil),
+		mockCmd.EXPECT().StderrPipe().Return(prStderr, nil),
+		mockCmd.EXPECT().Start().Return(nil),
+		mockCmd.EXPECT().Wait().DoAndReturn(func() error {
+			// The process itself dies with the context; only its orphaned
+			// descendant lives on, holding the pipes.
+			<-cmdCtx.Done()
+			return context.DeadlineExceeded
+		}),
+	)
+
+	client, err := NewBazelClient(context.Background(), Params{
+		BazelCommand:  "bazel",
+		WorkspacePath: "/tmp/test",
+		EnvVarsMap:    map[string]string{},
+		Logger:        zap.NewNop().Sugar(),
+		QueryTimeout:  10 * time.Millisecond,
+		ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander {
+			cmdCtx = ctx
+			return mockCmd
+		},
+	})
+	require.NoError(t, err)
+	client.pipeUnblockDelay = 20 * time.Millisecond
+
+	type queryOutcome struct {
+		result *buildpb.QueryResult
+		err    error
+	}
+	done := make(chan queryOutcome, 1)
+	go func() {
+		result, err := client.executeQueryInternal(context.Background(), "//...", nil)
+		done <- queryOutcome{result: result, err: err}
+	}()
+
+	select {
+	case out := <-done:
+		require.Error(t, out.err)
+		assert.ErrorIs(t, out.err, context.DeadlineExceeded)
+		assert.Nil(t, out.result)
+	case <-time.After(10 * time.Second):
+		t.Fatal("executeQueryInternal never returned: pipe reads were not unblocked")
+	}
+}
+
+// TestExecuteQueryInternal_DrainsStreamsBeforeWait verifies that cmd.Wait runs
+// only after the stream reads finish. The mock Waits emulate exec.Cmd.Wait,
+// which force-closes the parent pipe ends when it returns — had Wait run
+// before the reads completed, the still-streaming output would be truncated.
+func TestExecuteQueryInternal_DrainsStreamsBeforeWait(t *testing.T) {
+	t.Run("stdout is fully parsed", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+		ctrl := gomock.NewController(t)
+		mockCmd := commandermock.NewMockcommander(ctrl)
+
+		prStdout, pwStdout := io.Pipe()
+		prStderr, pwStderr := io.Pipe()
+
+		// io.Pipe is synchronous, so every write below blocks until the query's
+		// stream reader consumes it: the "process" is still producing output
+		// while the query runs.
+		const targetCount = 100
+		go func() {
+			defer pwStdout.Close()
+			defer pwStderr.Close()
+			ruleClass := "go_library"
+			for i := 0; i < targetCount; i++ {
+				name := fmt.Sprintf("//pkg:target%d", i)
+				target := &buildpb.Target{
+					Type: buildpb.Target_RULE.Enum(),
+					Rule: &buildpb.Rule{Name: &name, RuleClass: &ruleClass},
+				}
+				if _, err := protodelim.MarshalTo(pwStdout, target); err != nil {
+					return
+				}
+			}
+		}()
+
+		gomock.InOrder(
+			mockCmd.EXPECT().StdoutPipe().Return(prStdout, nil),
+			mockCmd.EXPECT().StderrPipe().Return(prStderr, nil),
+			mockCmd.EXPECT().Start().Return(nil),
+			mockCmd.EXPECT().Wait().DoAndReturn(func() error {
+				prStdout.Close()
+				prStderr.Close()
+				return nil
+			}),
+		)
+
+		client, err := NewBazelClient(context.Background(), Params{
+			BazelCommand:  "bazel",
+			WorkspacePath: "/tmp/test",
+			EnvVarsMap:    map[string]string{},
+			Logger:        zap.NewNop().Sugar(),
+			ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander {
+				return mockCmd
+			},
+		})
+		require.NoError(t, err)
+
+		result, err := client.executeQueryInternal(context.Background(), "//...", nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+		assert.Len(t, result.Target, targetCount)
+	})
+
+	t.Run("stderr is fully captured in the failure error", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+		ctrl := gomock.NewController(t)
+		mockCmd := commandermock.NewMockcommander(ctrl)
+
+		prStderr, pwStderr := io.Pipe()
+
+		const stderrTail = "FINAL STDERR LINE"
+		go func() {
+			defer pwStderr.Close()
+			_, _ = io.WriteString(pwStderr, strings.Repeat("bazel progress line\n", 200)+stderrTail)
+		}()
+
+		gomock.InOrder(
+			mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(strings.NewReader("")), nil),
+			mockCmd.EXPECT().StderrPipe().Return(prStderr, nil),
+			mockCmd.EXPECT().Start().Return(nil),
+			mockCmd.EXPECT().Wait().DoAndReturn(func() error {
+				prStderr.Close()
+				return errors.New("exit status 7")
+			}),
+		)
+
+		client, err := NewBazelClient(context.Background(), Params{
+			BazelCommand:  "bazel",
+			WorkspacePath: "/tmp/test",
+			EnvVarsMap:    map[string]string{},
+			Logger:        zap.NewNop().Sugar(),
+			ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander {
+				return mockCmd
+			},
+		})
+		require.NoError(t, err)
+
+		result, err := client.executeQueryInternal(context.Background(), "//...", nil)
+		require.Error(t, err)
+		require.NotNil(t, result)
+		// This asserts on the payload the error carries (the captured stderr),
+		// not on the error's own wording: the tail marker only appears if the
+		// whole stream was drained before the process was reaped.
+		assert.Contains(t, err.Error(), stderrTail)
+	})
+}
+
 func TestExecuteQueryInternal_Failures(t *testing.T) {
 	tests := []struct {
 		name            string

core/bazel/stream.go

@@ -23,39 +23,17 @@ import (
 	"google.golang.org/protobuf/encoding/protodelim"
 )
 
+// streamOutput copies src into dst until the stream ends. The copy blocks
+// until it is unblocked externally — by the process exiting and closing the
+// pipe write end, or by the caller force-closing the read end — so when ctx
+// has ended, any copy error is just fallout from that shutdown and the
+// context error is reported instead.
 func streamOutput(ctx context.Context, src io.Reader, dst io.Writer) error {
-	done := make(chan error, 1)
-	go func() {
-		_, err := io.Copy(dst, src)
-		done <- err
-	}()
-
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case err := <-done:
-		return err
-	}
-}
-
-func streamAndParseTargets(ctx context.Context, src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) {
-	type result struct {
-		queryResult *buildpb.QueryResult
-		err         error
-	}
-	done := make(chan result, 1)
-
-	go func() {
-		queryResult, err := getQueryResult(ctx, src, dst)
-		done <- result{queryResult: queryResult, err: err}
-	}()
-
-	select {
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	case res := <-done:
-		return res.queryResult, res.err
+	_, err := io.Copy(dst, src)
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		return ctxErr
 	}
+	return err
 }
 
 // cancelCheckInterval is how often we poll ctx.Err() inside per-target hot loops.
@@ -92,6 +70,10 @@ func getQueryResult(ctx context.Context, src io.Reader, dst io.Writer) (*buildpb
 		}
 		result.Target = append(result.Target, &target)
 	}
-
+	// Like streamOutput, prefer the context error over whatever the shutdown
+	// did to the stream mid-parse.
+	if err := ctx.Err(); err != nil {
+		return result, err
+	}
 	return result, parseErr
 }

core/execcmd/execcmd.go

@@ -26,12 +26,15 @@ import (
 	"time"
 )
 
-// gracePeriod is how long the child has to exit after SIGTERM before the Go
-// runtime escalates to SIGKILL via Cmd.WaitDelay.
-const gracePeriod = 10 * time.Second
+// GracePeriod is how long the child has to exit after SIGTERM before the Go
+// runtime escalates to SIGKILL via Cmd.WaitDelay. Exported so callers that
+// layer their own shutdown timeouts on top of a command (e.g. force-closing
+// its I/O pipes) can schedule them to fire only after the kill sequence has
+// run its course.
+const GracePeriod = 10 * time.Second
 
 // CommandContext is a drop-in replacement for exec.CommandContext that sends
-// SIGTERM when ctx is canceled and escalates to SIGKILL after gracePeriod if
+// SIGTERM when ctx is canceled and escalates to SIGKILL after GracePeriod if
 // the process is still running. This gives child processes (e.g. git, bazel)
 // a chance to release lock files and disconnect cleanly before being killed.
 func CommandContext(ctx context.Context, name string, args ...string) *exec.Cmd {
@@ -43,6 +46,6 @@ func CommandContext(ctx context.Context, name string, args ...string) *exec.Cmd
 		}
 		return err
 	}
-	cmd.WaitDelay = gracePeriod
+	cmd.WaitDelay = GracePeriod
 	return cmd
 }