diff --git a/core/bazel/BUILD.bazel b/core/bazel/BUILD.bazel index 3fa695b7..0b8a3df3 100644 --- a/core/bazel/BUILD.bazel +++ b/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", +) diff --git a/core/bazel/bazel.go b/core/bazel/bazel.go index 217283aa..86c9e7ec 100644 --- a/core/bazel/bazel.go +++ b/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,7 +66,7 @@ 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 } @@ -73,14 +74,14 @@ type Params struct { 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 } diff --git a/core/bazel/bazel_test.go b/core/bazel/bazel_test.go index c76342aa..12d3146a 100644 --- a/core/bazel/bazel_test.go +++ b/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) } diff --git a/core/bazel/command.go b/core/bazel/command.go index 42dc5861..dddea87c 100644 --- a/core/bazel/command.go +++ b/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() } diff --git a/core/bazel/command_test.go b/core/bazel/command_test.go new file mode 100644 index 00000000..d5bcdd4e --- /dev/null +++ b/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 +} diff --git a/core/bazel/commandermock/BUILD.bazel b/core/bazel/commandermock/BUILD.bazel deleted file mode 100644 index b90b729e..00000000 --- a/core/bazel/commandermock/BUILD.bazel +++ /dev/null @@ -1,9 +0,0 @@ -load("@rules_go//go:def.bzl", "go_library") - -go_library( - name = "commandermock", - srcs = ["commandermock.go"], - importpath = "github.com/uber/tango/core/bazel/commandermock", - visibility = ["//visibility:public"], - deps = ["@org_uber_go_mock//gomock"], -) diff --git a/core/bazel/commandermock/commandermock.go b/core/bazel/commandermock/commandermock.go deleted file mode 100644 index f11f4c66..00000000 --- a/core/bazel/commandermock/commandermock.go +++ /dev/null @@ -1,99 +0,0 @@ -// Code generated by MockGen. DO NOT EDIT. -// Source: core/bazel/command.go -// -// Generated by this command: -// -// mockgen -package=commandermock -destination=core/bazel/commandermock/commandermock.go -source=core/bazel/command.go commander -// - -// Package commandermock is a generated GoMock package. -package commandermock - -import ( - io "io" - reflect "reflect" - - gomock "go.uber.org/mock/gomock" -) - -// Mockcommander is a mock of commander interface. -type Mockcommander struct { - ctrl *gomock.Controller - recorder *MockcommanderMockRecorder - isgomock struct{} -} - -// MockcommanderMockRecorder is the mock recorder for Mockcommander. -type MockcommanderMockRecorder struct { - mock *Mockcommander -} - -// NewMockcommander creates a new mock instance. -func NewMockcommander(ctrl *gomock.Controller) *Mockcommander { - mock := &Mockcommander{ctrl: ctrl} - mock.recorder = &MockcommanderMockRecorder{mock} - return mock -} - -// EXPECT returns an object that allows the caller to indicate expected use. -func (m *Mockcommander) EXPECT() *MockcommanderMockRecorder { - return m.recorder -} - -// Start mocks base method. -func (m *Mockcommander) Start() error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Start") - ret0, _ := ret[0].(error) - return ret0 -} - -// Start indicates an expected call of Start. -func (mr *MockcommanderMockRecorder) Start() *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Start", reflect.TypeOf((*Mockcommander)(nil).Start)) -} - -// StderrPipe mocks base method. -func (m *Mockcommander) StderrPipe() (io.ReadCloser, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "StderrPipe") - ret0, _ := ret[0].(io.ReadCloser) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// StderrPipe indicates an expected call of StderrPipe. -func (mr *MockcommanderMockRecorder) StderrPipe() *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StderrPipe", reflect.TypeOf((*Mockcommander)(nil).StderrPipe)) -} - -// StdoutPipe mocks base method. -func (m *Mockcommander) StdoutPipe() (io.ReadCloser, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "StdoutPipe") - ret0, _ := ret[0].(io.ReadCloser) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// StdoutPipe indicates an expected call of StdoutPipe. -func (mr *MockcommanderMockRecorder) StdoutPipe() *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StdoutPipe", reflect.TypeOf((*Mockcommander)(nil).StdoutPipe)) -} - -// Wait mocks base method. -func (m *Mockcommander) Wait() error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Wait") - ret0, _ := ret[0].(error) - return ret0 -} - -// Wait indicates an expected call of Wait. -func (mr *MockcommanderMockRecorder) Wait() *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Wait", reflect.TypeOf((*Mockcommander)(nil).Wait)) -} diff --git a/core/bazel/query.go b/core/bazel/query.go index 37f10a73..4e8594ce 100644 --- a/core/bazel/query.go +++ b/core/bazel/query.go @@ -15,9 +15,9 @@ package bazel import ( - "bytes" "compress/gzip" "context" + "errors" "fmt" "io" "os" @@ -25,11 +25,10 @@ import ( buildpb "github.com/bazelbuild/buildtools/build_proto" "go.uber.org/zap" - "golang.org/x/sync/errgroup" "google.golang.org/protobuf/proto" ) -func (b *BazelClient) setupCommand(ctx context.Context, query string, startupOptions []string, additionalArgs ...string) commander { +func (b *BazelClient) setupCommand(ctx context.Context, query string, startupOptions []string, additionalArgs ...string) Commander { // Build command: bazel query --output=streamed_proto args := make([]string, 0, len(startupOptions)+1+len(additionalArgs)+2) args = append(args, startupOptions...) @@ -50,55 +49,65 @@ func (b *BazelClient) ExecuteQuery(ctx context.Context, req *QueryRequest) (*Que } func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, startupOptions []string, additionalArgs ...string) (*buildpb.QueryResult, error) { - var ( - stdoutBuf, stderrBuf bytes.Buffer - queryResults *buildpb.QueryResult - ) cmdCtx, cancel := context.WithTimeout(ctx, b.queryTimeout) defer cancel() - // setup bazel query command - cmd := b.setupCommand(cmdCtx, query, startupOptions, additionalArgs...) - // Get pipes for stdout and stderr BEFORE starting the process - stdout, err := cmd.StdoutPipe() - if err != nil { + if err := cmdCtx.Err(); err != nil { return nil, err } - stderr, err := cmd.StderrPipe() + cmd := b.setupCommand(cmdCtx, query, startupOptions, additionalArgs...) + + // Keep command output backed by files rather than StdoutPipe, StderrPipe, + // or custom pipe writers. A Bazel descendant can inherit a pipe's write end + // and prevent readers from seeing EOF after Bazel exits. Calling Wait early + // to activate its pipe cleanup is not an alternative: Wait may close the + // read ends before they are fully drained, truncating output. Direct files + // avoid both lifecycle problems and are rewound for parsing below. + stdout, err := os.CreateTemp(b.tempDir, "tango-bazel-stdout-*") if err != nil { - return nil, err + return nil, fmt.Errorf("create stdout file: %w", err) } + defer os.Remove(stdout.Name()) + defer stdout.Close() + // In streamLogs mode, stderr goes straight to os.Stderr so operators see // bazel progress live; otherwise it's captured for inclusion in failure // errors (see wrapQueryFailure). - stderrSink := io.Writer(&stderrBuf) - if b.streamLogs { - stderrSink = os.Stderr + stderrSink := io.Writer(os.Stderr) + stderrPath := "" + var stderr *os.File + if !b.streamLogs { + stderr, err = os.CreateTemp(b.tempDir, "tango-bazel-stderr-*") + if err != nil { + return nil, fmt.Errorf("create stderr file: %w", err) + } + defer os.Remove(stderr.Name()) + defer stderr.Close() + stderrSink = stderr + stderrPath = stderr.Name() } - // 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 - if err = cmd.Start(); err != nil { - return nil, err + // CommandContext and WaitDelay still govern process shutdown; using files + // changes only output transport and does not weaken cancellation. + waitErr := cmd.Run(stdout, stderrSink) + if ctxErr := cmdCtx.Err(); waitErr != nil && ctxErr != nil { + // A process terminated by SIGTERM or SIGKILL reports an ExitError, which + // does not wrap the context error. Preserve both causes so callers can + // reliably identify cancellation, and avoid parsing output after it. + return nil, b.wrapQueryFailure("bazel query canceled", errors.Join(ctxErr, waitErr), "") + } + + if _, err := stdout.Seek(0, io.SeekStart); err != nil { + return nil, fmt.Errorf("rewind bazel query output: %w", err) } - // stream and parse targets - g.Go(func() error { - var err error - queryResults, err = streamAndParseTargets(gCtx, stdout, &stdoutBuf) - return err - }) - // stream stderr - g.Go(func() error { - return streamOutput(gCtx, stderr, stderrSink) - }) - waitErr := cmd.Wait() - streamErr := g.Wait() + queryResults, streamErr := getQueryResult(cmdCtx, stdout) if waitErr != nil { - return queryResults, b.wrapQueryFailure("bazel query failed", waitErr, &stderrBuf) + return queryResults, b.wrapQueryFailure("bazel query failed", waitErr, stderrPath) } if streamErr != nil { - return nil, b.wrapQueryFailure("stream processing failed", streamErr, &stderrBuf) + if cmdCtx.Err() != nil { + stderrPath = "" + } + return nil, b.wrapQueryFailure("stream processing failed", streamErr, stderrPath) } b.logger.Debugw("Parsed targets from bazel query", zap.Int("target_count", len(queryResults.Target))) return queryResults, nil @@ -108,10 +117,14 @@ func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, st // was captured (streamLogs off), its contents are appended so the failure is // self-contained. When streamLogs is on the operator has already seen stderr // live, so it's omitted. -func (b *BazelClient) wrapQueryFailure(msg string, cause error, stderrBuf *bytes.Buffer) error { +func (b *BazelClient) wrapQueryFailure(msg string, cause error, stderrPath string) error { tail := "" - if !b.streamLogs { - tail = "\nstderr:\n" + stderrBuf.String() + if stderrPath != "" { + stderr, err := os.ReadFile(stderrPath) + if err != nil { + return fmt.Errorf("%s: %w (read stderr: %v)", msg, cause, err) + } + tail = "\nstderr:\n" + string(stderr) } return fmt.Errorf("%s: %w%s", msg, cause, tail) } diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index c506214f..de96a06e 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -18,109 +18,57 @@ import ( "bytes" "context" "errors" + "fmt" "io" + "os" "strings" "testing" "time" buildpb "github.com/bazelbuild/buildtools/build_proto" + "github.com/golang/mock/gomock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/uber/tango/core/bazel/commandermock" "go.uber.org/goleak" - "go.uber.org/mock/gomock" "go.uber.org/zap" "google.golang.org/protobuf/encoding/protodelim" ) func TestExecuteQuery_Success(t *testing.T) { - defer goleak.VerifyNone(t) ctrl := gomock.NewController(t) - mockCmd := commandermock.NewMockcommander(ctrl) - var ( - ruleName, ruleClass = "//pkg:target", "go_library" - ) - target := &buildpb.Target{ - Type: buildpb.Target_RULE.Enum(), - Rule: &buildpb.Rule{ - Name: &ruleName, - RuleClass: &ruleClass, - }, - } - var protoData bytes.Buffer - _, err := protodelim.MarshalTo(&protoData, target) - require.NoError(t, err) - gomock.InOrder( - mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(&protoData), nil), - mockCmd.EXPECT().StderrPipe().Return(io.NopCloser(strings.NewReader("")), nil), - mockCmd.EXPECT().Start().Return(nil), - mockCmd.EXPECT().Wait().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 - }, + mockCmd := NewMockCommander(ctrl) + protoData := marshalTarget(t, "//pkg:target") + mockCmd.EXPECT().Run(gomock.Any(), gomock.Any()).DoAndReturn(func(stdout, _ io.Writer) error { + _, err := stdout.Write(protoData) + return err }) + client := newTestClient(t, func(context.Context, string, ...string) Commander { return mockCmd }) resp, err := client.ExecuteQuery(context.Background(), &QueryRequest{Query: "//..."}) + require.NoError(t, err) - require.NotNil(t, resp) - require.NotNil(t, resp.Result) - require.Equal(t, 1, len(resp.Result.Target)) - assert.Equal(t, &ruleName, resp.Result.Target[0].Rule.Name) - assert.Equal(t, &ruleClass, resp.Result.Target[0].Rule.RuleClass) + require.Len(t, resp.Result.Target, 1) + assert.Equal(t, "//pkg:target", resp.Result.Target[0].Rule.GetName()) } func TestExecuteQuery_WithStartupOptions(t *testing.T) { - defer goleak.VerifyNone(t) ctrl := gomock.NewController(t) - mockCmd := commandermock.NewMockcommander(ctrl) - var ( - ruleName, ruleClass = "//pkg:target", "go_library" - ) - target := &buildpb.Target{ - Type: buildpb.Target_RULE.Enum(), - Rule: &buildpb.Rule{ - Name: &ruleName, - RuleClass: &ruleClass, - }, - } - var protoData bytes.Buffer - _, err := protodelim.MarshalTo(&protoData, target) - require.NoError(t, err) + mockCmd := NewMockCommander(ctrl) + mockCmd.EXPECT().Run(gomock.Any(), gomock.Any()).Return(nil) var capturedArgs []string - gomock.InOrder( - mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(&protoData), nil), - mockCmd.EXPECT().StderrPipe().Return(io.NopCloser(strings.NewReader("")), nil), - mockCmd.EXPECT().Start().Return(nil), - mockCmd.EXPECT().Wait().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 { - capturedArgs = arg - return mockCmd - }, + client := newTestClient(t, func(_ context.Context, _ string, args ...string) Commander { + capturedArgs = args + return mockCmd }) - require.NoError(t, err) - resp, err := client.ExecuteQuery(context.Background(), &QueryRequest{ + _, err := client.ExecuteQuery(context.Background(), &QueryRequest{ Query: "//...", StartupOptions: []string{"--bazelrc=/custom/.bazelrc", "--output_base=/tmp/bazel"}, AdditionalArgs: []string{"--keep_going"}, }) - require.NoError(t, err) - require.NotNil(t, resp) - // Verify command structure: bazel query --output=streamed_proto - require.Equal(t, []string{ + require.NoError(t, err) + assert.Equal(t, []string{ "--bazelrc=/custom/.bazelrc", "--output_base=/tmp/bazel", "query", @@ -133,143 +81,143 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { defer goleak.VerifyNone(t) ctrl := gomock.NewController(t) - mockCmd := commandermock.NewMockcommander(ctrl) + mockCmd := NewMockCommander(ctrl) + var cmdCtx context.Context + mockCmd.EXPECT().Run(gomock.Any(), gomock.Any()).DoAndReturn(func(io.Writer, io.Writer) error { + <-cmdCtx.Done() + return errors.New("signal: terminated") + }) - prStdout, pwStdout := io.Pipe() - prStderr, pwStderr := io.Pipe() + client := newTestClient(t, func(ctx context.Context, _ string, _ ...string) Commander { + cmdCtx = ctx + return mockCmd + }) + client.queryTimeout = 10 * time.Millisecond + result, err := client.executeQueryInternal(context.Background(), "//...", nil) + + require.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.Nil(t, result) +} + +func TestExecuteQueryInternal_PreCanceledContext(t *testing.T) { + ctrl := gomock.NewController(t) + mockCmd := NewMockCommander(ctrl) + client := newTestClient(t, func(context.Context, string, ...string) Commander { return mockCmd }) + ctx, cancel := context.WithCancel(context.Background()) + cancel() - // Set up the mock expectations in the exact order they will be called. - 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 { - // Simulate process ending after timeout - return context.DeadlineExceeded - }), - ) + result, err := client.executeQueryInternal(ctx, "//...", nil) - client, err := NewBazelClient(context.Background(), Params{ - BazelCommand: "bazel", - WorkspacePath: "/tmp/test", - Logger: zap.NewNop().Sugar(), - EnvVarsMap: map[string]string{}, - QueryTimeout: 10 * time.Millisecond, // Short timeout for test + require.ErrorIs(t, err, context.Canceled) + assert.Nil(t, result) +} - ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander { - // Simulate process behavior: when context is cancelled, close pipes - go func() { - <-ctx.Done() - // Close pipes to unblock readers - pwStdout.Close() - pwStderr.Close() - }() - return mockCmd - }, +func TestExecuteQueryInternal_CancelDuringParsing(t *testing.T) { + defer goleak.VerifyNone(t) + ctrl := gomock.NewController(t) + mockCmd := NewMockCommander(ctrl) + ctx, cancel := context.WithCancel(context.Background()) + mockCmd.EXPECT().Run(gomock.Any(), gomock.Any()).DoAndReturn(func(stdout, _ io.Writer) error { + _, err := stdout.Write(marshalTarget(t, "//pkg:target")) + cancel() + return err }) - require.NoError(t, err) + + client := newTestClient(t, func(context.Context, string, ...string) Commander { return mockCmd }) + result, err := client.executeQueryInternal(ctx, "//...", nil) + + require.ErrorIs(t, err, context.Canceled) + assert.Nil(t, result) +} + +func TestExecuteQueryInternal_CapturesCompleteOutput(t *testing.T) { + ctrl := gomock.NewController(t) + mockCmd := NewMockCommander(ctrl) + const targetCount = 100 + const stderrTail = "FINAL STDERR LINE" + mockCmd.EXPECT().Run(gomock.Any(), gomock.Any()).DoAndReturn(func(stdout, stderr io.Writer) error { + for i := 0; i < targetCount; i++ { + if _, err := stdout.Write(marshalTarget(t, fmt.Sprintf("//pkg:target%d", i))); err != nil { + return err + } + } + _, err := io.WriteString(stderr, strings.Repeat("bazel progress line\n", 200)+stderrTail) + if err != nil { + return err + } + return errors.New("exit status 7") + }) + + client := newTestClient(t, func(context.Context, string, ...string) Commander { return mockCmd }) result, err := client.executeQueryInternal(context.Background(), "//...", nil) - require.Nil(t, result) + require.Error(t, err) - // Should get timeout or deadline exceeded error - assert.Contains(t, err.Error(), "deadline exceeded") + require.Len(t, result.Target, targetCount) + assert.Contains(t, err.Error(), stderrTail) } -func TestExecuteQueryInternal_Failures(t *testing.T) { - tests := []struct { - name string - setupMock func(*commandermock.Mockcommander) - expectedError string - expectNilResult bool - }{ - { - name: "stdout pipe failure", - setupMock: func(m *commandermock.Mockcommander) { - m.EXPECT().StdoutPipe().Return(nil, errors.New("stdout pipe failed")) - }, - expectedError: "stdout pipe failed", - expectNilResult: true, - }, - { - name: "stderr pipe failure", - setupMock: func(m *commandermock.Mockcommander) { - m.EXPECT().StdoutPipe().Return(io.NopCloser(strings.NewReader("")), nil) - m.EXPECT().StderrPipe().Return(nil, errors.New("stderr pipe failed")) - }, - expectedError: "stderr pipe failed", - expectNilResult: true, - }, - { - name: "command start failure", - setupMock: func(m *commandermock.Mockcommander) { - m.EXPECT().StdoutPipe().Return(io.NopCloser(strings.NewReader("")), nil) - m.EXPECT().StderrPipe().Return(io.NopCloser(strings.NewReader("")), nil) - m.EXPECT().Start().Return(errors.New("failed to start process")) - }, - expectedError: "failed to start process", - expectNilResult: true, - }, - { - name: "command wait failure", - setupMock: func(m *commandermock.Mockcommander) { - m.EXPECT().StdoutPipe().Return(io.NopCloser(strings.NewReader("")), nil) - m.EXPECT().StderrPipe().Return(io.NopCloser(strings.NewReader("")), nil) - m.EXPECT().Start().Return(nil) - m.EXPECT().Wait().Return(errors.New("command wait failed")) - }, - expectedError: "command wait failed", - expectNilResult: false, - }, - } +func TestExecuteQueryInternal_ParseFailure(t *testing.T) { + ctrl := gomock.NewController(t) + mockCmd := NewMockCommander(ctrl) + mockCmd.EXPECT().Run(gomock.Any(), gomock.Any()).DoAndReturn(func(stdout, _ io.Writer) error { + _, err := io.WriteString(stdout, "not a streamed proto") + return err + }) - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - defer goleak.VerifyNone(t) - ctrl := gomock.NewController(t) - mockCmd := commandermock.NewMockcommander(ctrl) - tt.setupMock(mockCmd) + client := newTestClient(t, func(context.Context, string, ...string) Commander { return mockCmd }) + result, err := client.executeQueryInternal(context.Background(), "//...", 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.Error(t, err) - if tt.expectNilResult { - require.Nil(t, result) - } else { - require.NotNil(t, result) - } - assert.Contains(t, err.Error(), tt.expectedError) - }) - } + require.Error(t, err) + assert.Nil(t, result) } -func TestExecuteQuery_ErrorCase(t *testing.T) { - defer goleak.VerifyNone(t) +func TestExecuteQueryInternal_StreamLogs(t *testing.T) { ctrl := gomock.NewController(t) - mockCmd := commandermock.NewMockcommander(ctrl) + mockCmd := NewMockCommander(ctrl) + mockCmd.EXPECT().Run(gomock.Any(), os.Stderr).Return(errors.New("exit status 1")) - mockCmd.EXPECT().StdoutPipe().Return(nil, errors.New("stdout pipe failed")) + client := newTestClient(t, func(context.Context, string, ...string) Commander { return mockCmd }) + client.streamLogs = true + _, err := client.executeQueryInternal(context.Background(), "//...", nil) + require.Error(t, err) +} - 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 - }, - }) +func TestExecuteQuery_ErrorCase(t *testing.T) { + ctrl := gomock.NewController(t) + mockCmd := NewMockCommander(ctrl) + mockCmd.EXPECT().Run(gomock.Any(), gomock.Any()).Return(errors.New("command failed")) + client := newTestClient(t, func(context.Context, string, ...string) Commander { return mockCmd }) resp, err := client.ExecuteQuery(context.Background(), &QueryRequest{Query: "//..."}) + require.Error(t, err) - require.Nil(t, resp) - assert.Contains(t, err.Error(), "stdout pipe failed") + assert.Nil(t, resp) + assert.Contains(t, err.Error(), "command failed") +} + +func marshalTarget(t *testing.T, name string) []byte { + t.Helper() + ruleClass := "go_library" + target := &buildpb.Target{ + Type: buildpb.Target_RULE.Enum(), + Rule: &buildpb.Rule{Name: &name, RuleClass: &ruleClass}, + } + var out bytes.Buffer + _, err := protodelim.MarshalTo(&out, target) + require.NoError(t, err) + return out.Bytes() +} + +func newTestClient(t *testing.T, execCmd func(context.Context, string, ...string) Commander) *BazelClient { + t.Helper() + client, err := NewBazelClient(context.Background(), Params{ + BazelCommand: "bazel", + WorkspacePath: "/tmp/test", + EnvVarsMap: map[string]string{}, + Logger: zap.NewNop().Sugar(), + ExecCommandContext: execCmd, + }) + require.NoError(t, err) + return client } diff --git a/core/bazel/stream.go b/core/bazel/stream.go index 2e78fd7b..751c1eb8 100644 --- a/core/bazel/stream.go +++ b/core/bazel/stream.go @@ -23,53 +23,17 @@ import ( "google.golang.org/protobuf/encoding/protodelim" ) -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 - } -} - // cancelCheckInterval is how often we poll ctx.Err() inside per-target hot loops. // Picked to keep overhead negligible while still surfacing cancellation in <100ms // for typical target rates. const cancelCheckInterval = 1024 // getQueryResult reads a QueryResult containing targets from the stream and returns it. -func getQueryResult(ctx context.Context, src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) { +func getQueryResult(ctx context.Context, src io.Reader) (*buildpb.QueryResult, error) { result := &buildpb.QueryResult{ Target: make([]*buildpb.Target, 0), } - tr := io.TeeReader(src, dst) - br := bufio.NewReader(tr) + br := bufio.NewReader(src) unmarshalOpts := protodelim.UnmarshalOptions{ MaxSize: 64 * 1024 * 1024, // 64MB limit } @@ -92,6 +56,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 }