From 8ddef96e9f6dfcfd573f2bd6030206b461ecb601 Mon Sep 17 00:00:00 2001 From: yushan Date: Tue, 24 Feb 2026 22:28:59 +0000 Subject: [PATCH 1/9] [tango] Update bazel query command to do plain proto dump and unmarshal --- core/bazel/BUILD.bazel | 3 - core/bazel/query.go | 64 ++++++++------------ core/bazel/query_test.go | 59 ++++++++---------- core/bazel/stream.go | 68 --------------------- core/bazelrunner/native.go | 1 + example/cmd/query-bench/BUILD.bazel | 18 ++++++ example/cmd/query-bench/main.go | 94 +++++++++++++++++++++++++++++ 7 files changed, 164 insertions(+), 143 deletions(-) delete mode 100644 core/bazel/stream.go create mode 100644 example/cmd/query-bench/BUILD.bazel create mode 100644 example/cmd/query-bench/main.go diff --git a/core/bazel/BUILD.bazel b/core/bazel/BUILD.bazel index a9d2922f..06aeaf7d 100644 --- a/core/bazel/BUILD.bazel +++ b/core/bazel/BUILD.bazel @@ -6,15 +6,12 @@ go_library( "bazel.go", "command.go", "query.go", - "stream.go", ], importpath = "github.com/uber/tango/core/bazel", visibility = ["//visibility:public"], deps = [ "@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", ], ) diff --git a/core/bazel/query.go b/core/bazel/query.go index 9efc27db..2edf1ffc 100644 --- a/core/bazel/query.go +++ b/core/bazel/query.go @@ -1,7 +1,6 @@ package bazel import ( - "bytes" "compress/gzip" "context" "io" @@ -10,17 +9,16 @@ 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 { - // Build command: bazel query --output=streamed_proto + // Build command: bazel query --output=proto args := make([]string, 0, len(startupOptions)+1+len(additionalArgs)+2) args = append(args, startupOptions...) args = append(args, "query") args = append(args, additionalArgs...) - args = append(args, "--output=streamed_proto") + args = append(args, "--output=proto") args = append(args, query) b.logger.Info("Querying Bazel", zap.String("workspacePath", b.workspacePath), zap.String("query", query)) return b.execCommandContext(ctx, b.bazelCommand, args...) @@ -35,56 +33,44 @@ 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 + + // Get stdout pipe stdout, err := cmd.StdoutPipe() if err != nil { return nil, err } - stderr, err := cmd.StderrPipe() - if err != nil { + + // Start the process + if err := cmd.Start(); err != nil { return nil, err } - // 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 { + // Read all output at once (batch mode) - much faster than incremental parsing + out, err := io.ReadAll(stdout) + if err != nil { return nil, 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, &stderrBuf) - }) - - streamErr := g.Wait() - waitErr := cmd.Wait() - // The command itself failed. - if waitErr != nil { - b.logger.Error("Bazel query failed failed: %v", waitErr) - return queryResults, waitErr + + // Wait for command to complete + if err := cmd.Wait(); err != nil { + b.logger.Error("Bazel query failed: %v", err) + return nil, err } - // The command succeeded, but there was an error in the stream processing. - if streamErr != nil { - b.logger.Error("Error in stream processing: %v", streamErr) - return nil, streamErr + + // Single unmarshal operation - much faster than per-target streaming + var qr buildpb.QueryResult + if err := proto.Unmarshal(out, &qr); err != nil { + b.logger.Error("Failed to unmarshal query result: %v", err) + return nil, err } - b.logger.Debugf("Parsed %d targets from bazel query", len(queryResults.Target)) - return queryResults, nil + + b.logger.Debugf("Parsed %d targets from bazel query", len(qr.Target)) + return &qr, nil } // FromFile reads a proto file generated by bazel query. diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index bdab8145..27be952e 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -16,7 +16,7 @@ import ( "go.uber.org/goleak" "go.uber.org/mock/gomock" "go.uber.org/zap" - "google.golang.org/protobuf/encoding/protodelim" + "google.golang.org/protobuf/proto" ) func TestExecuteQuery_Success(t *testing.T) { @@ -33,12 +33,16 @@ func TestExecuteQuery_Success(t *testing.T) { RuleClass: &ruleClass, }, } - var protoData bytes.Buffer - _, err := protodelim.MarshalTo(&protoData, target) + + // Create a QueryResult with the target (batch proto format, not protodelim streaming) + queryResult := &buildpb.QueryResult{ + Target: []*buildpb.Target{target}, + } + protoData, err := proto.Marshal(queryResult) 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().StdoutPipe().Return(io.NopCloser(bytes.NewReader(protoData)), nil), mockCmd.EXPECT().Start().Return(nil), mockCmd.EXPECT().Wait().Return(nil), ) @@ -75,14 +79,17 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { RuleClass: &ruleClass, }, } - var protoData bytes.Buffer - _, err := protodelim.MarshalTo(&protoData, target) + + // Create a QueryResult with the target (batch proto format) + queryResult := &buildpb.QueryResult{ + Target: []*buildpb.Target{target}, + } + protoData, err := proto.Marshal(queryResult) require.NoError(t, err) var capturedArgs []string gomock.InOrder( - mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(&protoData), nil), - mockCmd.EXPECT().StderrPipe().Return(io.NopCloser(strings.NewReader("")), nil), + mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(bytes.NewReader(protoData)), nil), mockCmd.EXPECT().Start().Return(nil), mockCmd.EXPECT().Wait().Return(nil), ) @@ -105,31 +112,30 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { require.NoError(t, err) require.NotNil(t, resp) - // Verify command structure: bazel query --output=streamed_proto + // Verify command structure: bazel query --output=proto require.Equal(t, []string{ "--bazelrc=/custom/.bazelrc", "--output_base=/tmp/bazel", "query", "--keep_going", - "--output=streamed_proto", + "--output=proto", "//...", }, capturedArgs) } -func TestexecuteQueryInternal_ContextTimeout(t *testing.T) { +func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { defer goleak.VerifyNone(t) ctrl := gomock.NewController(t) mockCmd := commandermock.NewMockcommander(ctrl) prStdout, pwStdout := io.Pipe() - prStderr, pwStderr := io.Pipe() // Set up the mock expectations in the exact order they will be called. + // Note: When io.ReadAll gets an error (like context deadline), we return immediately + // without calling Wait(), so we don't expect Wait() to be called here. gomock.InOrder( mockCmd.EXPECT().StdoutPipe().Return(prStdout, nil), - mockCmd.EXPECT().StderrPipe().Return(prStderr, nil), mockCmd.EXPECT().Start().Return(nil), - mockCmd.EXPECT().Wait().Return(context.DeadlineExceeded), ) client, err := NewBazelClient(Params{ @@ -146,12 +152,10 @@ func TestexecuteQueryInternal_ContextTimeout(t *testing.T) { go func() { <-ctx.Done() // Wait for the timeout to fire - // "Killing" the process: close the pipes. - // This unblocks the Read() calls in your - // streamAndParseTargets and streamOutput goroutines. - // We close with the context's error so g.Wait() sees it. + // "Killing" the process: close the pipe. + // This unblocks the Read() call in io.ReadAll. + // We close with the context's error so the read sees it. pwStdout.CloseWithError(ctx.Err()) - pwStderr.CloseWithError(ctx.Err()) }() return mockCmd }, @@ -162,7 +166,7 @@ func TestexecuteQueryInternal_ContextTimeout(t *testing.T) { assert.Contains(t, err.Error(), "context deadline exceeded") } -func TestexecuteQueryInternal_Failures(t *testing.T) { +func TestExecuteQueryInternal_Failures(t *testing.T) { tests := []struct { name string setupMock func(*commandermock.Mockcommander) @@ -177,20 +181,10 @@ func TestexecuteQueryInternal_Failures(t *testing.T) { 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", @@ -200,12 +194,11 @@ func TestexecuteQueryInternal_Failures(t *testing.T) { 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, + expectNilResult: true, // Changed from false - implementation returns nil on Wait() error }, } diff --git a/core/bazel/stream.go b/core/bazel/stream.go deleted file mode 100644 index 176bf24c..00000000 --- a/core/bazel/stream.go +++ /dev/null @@ -1,68 +0,0 @@ -package bazel - -import ( - "bufio" - "context" - "io" - - buildpb "github.com/bazelbuild/buildtools/build_proto" - "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(src, dst) - done <- result{queryResult: queryResult, err: err} - }() - - select { - case <-ctx.Done(): - return nil, ctx.Err() - case res := <-done: - return res.queryResult, res.err - } -} - -// getQueryResult reads a QueryResult containing targets from the stream and returns it. -func getQueryResult(src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) { - result := &buildpb.QueryResult{ - Target: make([]*buildpb.Target, 0), - } - tr := io.TeeReader(src, dst) - br := bufio.NewReader(tr) - - for { - var target buildpb.Target - err := protodelim.UnmarshalFrom(br, &target) - if err == io.EOF { - break - } - if err != nil { - return result, err - } - result.Target = append(result.Target, &target) - } - - return result, nil -} diff --git a/core/bazelrunner/native.go b/core/bazelrunner/native.go index 78fa6c3e..7711eda2 100644 --- a/core/bazelrunner/native.go +++ b/core/bazelrunner/native.go @@ -43,6 +43,7 @@ func (g *nativeGraphRunner) Compute(ctx context.Context, ws workspace.Workspace) // --noproto: parameters exclude fields from the output that are not used for hashing anyways, making // proto blob smaller and serialization/deserialization faster // TODO: pass in --enable_workspace or --enable_bzlmod based on the config + AdditionalArgs: []string{"--order_output=no", "--proto:locations", "--noproto:default_values"}, }) if err != nil { diff --git a/example/cmd/query-bench/BUILD.bazel b/example/cmd/query-bench/BUILD.bazel new file mode 100644 index 00000000..11d37eee --- /dev/null +++ b/example/cmd/query-bench/BUILD.bazel @@ -0,0 +1,18 @@ +load("@rules_go//go:def.bzl", "go_binary", "go_library") + +go_library( + name = "query-bench_lib", + srcs = ["main.go"], + importpath = "github.com/uber/tango/example/cmd/query-bench", + visibility = ["//visibility:private"], + deps = [ + "//core/bazel", + "@org_uber_go_zap//:zap", + ], +) + +go_binary( + name = "query-bench", + embed = [":query-bench_lib"], + visibility = ["//visibility:public"], +) diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go new file mode 100644 index 00000000..ba5ed8a2 --- /dev/null +++ b/example/cmd/query-bench/main.go @@ -0,0 +1,94 @@ +// query-bench is a local benchmarking tool for the bazel query execution path. +// It runs the same query that nativeGraphRunner uses and reports timing and target counts. +// +// Usage: +// +// bazel run //cmd/query-bench -- --workspace /path/to/repo +// bazel run //cmd/query-bench -- --workspace /path/to/repo --bazel bazelisk --runs 3 +// bazel run //cmd/query-bench -- --workspace /path/to/repo --exclude-external +// bazel run //cmd/query-bench -- --workspace /path/to/repo --query '//...:all-targets' +package main + +import ( + "context" + "flag" + "fmt" + "os" + "time" + + "github.com/uber/tango/core/bazel" + "go.uber.org/zap" +) + +func main() { + if err := run(); err != nil { + fmt.Fprintf(os.Stderr, "error: %v\n", err) + os.Exit(1) + } +} + +func run() error { + bazelCmd := flag.String("bazel", "", "bazel binary to invoke (default: auto-detect bazel/bazelisk on PATH)") + workspace := flag.String("workspace", ".", "workspace root to run bazel query in") + query := flag.String("query", "", "bazel query expression (default: the standard nativeGraphRunner query)") + excludeExternal := flag.Bool("exclude-external", false, "use deps(//...:all-targets) instead of including //external:all-targets") + runs := flag.Int("runs", 1, "number of times to run the query (for benchmarking)") + timeout := flag.Duration("timeout", 30*time.Minute, "per-run timeout") + flag.Parse() + + q := *query + if q == "" { + if *excludeExternal { + q = "deps(//...:all-targets)" + } else { + q = "//external:all-targets + deps(//...:all-targets)" + } + } + + logger, err := zap.NewDevelopment() + if err != nil { + return fmt.Errorf("creating logger: %w", err) + } + defer logger.Sync() + + client, err := bazel.NewBazelClient(bazel.Params{ + BazelCommand: *bazelCmd, + WorkspacePath: *workspace, + Logger: logger.Sugar(), + QueryTimeout: *timeout, + }) + if err != nil { + return fmt.Errorf("creating bazel client: %w", err) + } + + req := &bazel.QueryRequest{ + Query: q, + AdditionalArgs: []string{"--order_output=no", "--proto:locations", "--noproto:default_values"}, + } + + fmt.Printf("workspace: %s\n", *workspace) + fmt.Printf("query: %s\n", q) + fmt.Printf("runs: %d\n\n", *runs) + + var totalDuration time.Duration + for i := range *runs { + ctx, cancel := context.WithTimeout(context.Background(), *timeout) + start := time.Now() + resp, err := client.ExecuteQuery(ctx, req) + elapsed := time.Since(start) + cancel() + + if err != nil { + return fmt.Errorf("run %d: query failed: %w", i+1, err) + } + + totalDuration += elapsed + fmt.Printf("run %d: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(resp.Result.Target)) + } + + if *runs > 1 { + fmt.Printf("\naverage: %v\n", (totalDuration / time.Duration(*runs)).Round(time.Millisecond)) + } + + return nil +} From 6c99321ca317ea28df8cc39a903773434d69fb17 Mon Sep 17 00:00:00 2001 From: yushan Date: Tue, 24 Feb 2026 22:42:09 +0000 Subject: [PATCH 2/9] Update tango benchmark --- core/bazel/BUILD.bazel | 2 +- example/cmd/query-bench/BUILD.bazel | 1 + example/cmd/query-bench/main.go | 13 ++++++++++--- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/core/bazel/BUILD.bazel b/core/bazel/BUILD.bazel index 06aeaf7d..6b253a0d 100644 --- a/core/bazel/BUILD.bazel +++ b/core/bazel/BUILD.bazel @@ -28,7 +28,7 @@ go_test( "@com_github_bazelbuild_buildtools//build_proto", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", - "@org_golang_google_protobuf//encoding/protodelim", + "@org_golang_google_protobuf//proto", "@org_uber_go_goleak//:goleak", "@org_uber_go_mock//gomock", "@org_uber_go_zap//:zap", diff --git a/example/cmd/query-bench/BUILD.bazel b/example/cmd/query-bench/BUILD.bazel index 11d37eee..06a7986d 100644 --- a/example/cmd/query-bench/BUILD.bazel +++ b/example/cmd/query-bench/BUILD.bazel @@ -7,6 +7,7 @@ go_library( visibility = ["//visibility:private"], deps = [ "//core/bazel", + "//core/targethasher", "@org_uber_go_zap//:zap", ], ) diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index ba5ed8a2..f5cf6e0d 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -1,4 +1,4 @@ -// query-bench is a local benchmarking tool for the bazel query execution path. +// query-bench is a local benchmarking tool for the bazel query execution path and targethasher. // It runs the same query that nativeGraphRunner uses and reports timing and target counts. // // Usage: @@ -18,6 +18,7 @@ import ( "github.com/uber/tango/core/bazel" "go.uber.org/zap" + "github.com/uber/tango/core/targethasher" ) func main() { @@ -82,13 +83,19 @@ func run() error { return fmt.Errorf("run %d: query failed: %w", i+1, err) } + fmt.Printf("run %d: bazel query: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(resp.Result.Target)) + start = time.Now() + targethasherResult, err := targethasher.FromProto(context.Background(), resp.Result, *workspace, targethasher.HashConfig{}) + if err != nil { + return fmt.Errorf("converting result to targethasher.Result: %w", err) + } + elapsed = time.Since(start) + fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) totalDuration += elapsed - fmt.Printf("run %d: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(resp.Result.Target)) } if *runs > 1 { fmt.Printf("\naverage: %v\n", (totalDuration / time.Duration(*runs)).Round(time.Millisecond)) } - return nil } From 668555af7af0c9dd909c596134cb3adac9abff8a Mon Sep 17 00:00:00 2001 From: yushan Date: Tue, 24 Feb 2026 22:45:52 +0000 Subject: [PATCH 3/9] Update --- example/cmd/query-bench/main.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index f5cf6e0d..022d0adc 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -83,15 +83,16 @@ func run() error { return fmt.Errorf("run %d: query failed: %w", i+1, err) } + totalDuration += elapsed fmt.Printf("run %d: bazel query: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(resp.Result.Target)) start = time.Now() - targethasherResult, err := targethasher.FromProto(context.Background(), resp.Result, *workspace, targethasher.HashConfig{}) + targethasherResult, err := targethasher.FromProto(ctx, resp.Result, *workspace, targethasher.HashConfig{}) if err != nil { return fmt.Errorf("converting result to targethasher.Result: %w", err) } elapsed = time.Since(start) - fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) totalDuration += elapsed + fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) } if *runs > 1 { From fa9e2e57f90e77ca1bba0594bb0f97d9a2e061cc Mon Sep 17 00:00:00 2001 From: yushan Date: Wed, 25 Feb 2026 22:11:43 +0000 Subject: [PATCH 4/9] Update with fix --- core/bazel/BUILD.bazel | 3 + core/bazel/command.go | 4 +- core/bazel/commandermock/commandermock.go | 6 ++ core/bazel/query.go | 76 ++++++++++++++++------- core/bazel/query_test.go | 76 ++++++++++++++++------- core/bazel/stream.go | 72 +++++++++++++++++++++ example/cmd/query-bench/main.go | 4 +- 7 files changed, 192 insertions(+), 49 deletions(-) create mode 100644 core/bazel/stream.go diff --git a/core/bazel/BUILD.bazel b/core/bazel/BUILD.bazel index 6b253a0d..39cf50d4 100644 --- a/core/bazel/BUILD.bazel +++ b/core/bazel/BUILD.bazel @@ -6,12 +6,15 @@ go_library( "bazel.go", "command.go", "query.go", + "stream.go", ], importpath = "github.com/uber/tango/core/bazel", visibility = ["//visibility:public"], deps = [ "@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", ], ) diff --git a/core/bazel/command.go b/core/bazel/command.go index 839fd388..abf09ad7 100644 --- a/core/bazel/command.go +++ b/core/bazel/command.go @@ -1,6 +1,8 @@ package bazel -import "io" +import ( + "io" +) type commander interface { StdoutPipe() (io.ReadCloser, error) diff --git a/core/bazel/commandermock/commandermock.go b/core/bazel/commandermock/commandermock.go index 13aba3c6..f11f4c66 100644 --- a/core/bazel/commandermock/commandermock.go +++ b/core/bazel/commandermock/commandermock.go @@ -1,5 +1,10 @@ // 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 @@ -15,6 +20,7 @@ import ( type Mockcommander struct { ctrl *gomock.Controller recorder *MockcommanderMockRecorder + isgomock struct{} } // MockcommanderMockRecorder is the mock recorder for Mockcommander. diff --git a/core/bazel/query.go b/core/bazel/query.go index 2edf1ffc..3d7c18c6 100644 --- a/core/bazel/query.go +++ b/core/bazel/query.go @@ -1,25 +1,31 @@ package bazel import ( + "bytes" "compress/gzip" "context" + "fmt" "io" "os" "strings" 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 { - // Build command: bazel query --output=proto - args := make([]string, 0, len(startupOptions)+1+len(additionalArgs)+2) + // Build command: bazel query --order_output=no --output=streamed_proto + args := make([]string, 0, len(startupOptions)+1+len(additionalArgs)+3) args = append(args, startupOptions...) args = append(args, "query") - args = append(args, additionalArgs...) - args = append(args, "--output=proto") args = append(args, query) + args = append(args, "--order_output=no") + args = append(args, "--output=streamed_proto") + args = append(args, additionalArgs...) b.logger.Info("Querying Bazel", zap.String("workspacePath", b.workspacePath), zap.String("query", query)) return b.execCommandContext(ctx, b.bazelCommand, args...) } @@ -32,47 +38,69 @@ func (b *BazelClient) ExecuteQuery(ctx context.Context, req *QueryRequest) (*Que return &QueryResponse{Result: result}, nil } + func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, startupOptions []string, additionalArgs ...string) (*buildpb.QueryResult, error) { + var ( + 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 stdout pipe + // Get pipes for stdout and stderr BEFORE starting the process stdout, err := cmd.StdoutPipe() if err != nil { return nil, err } - - // Start the process - if err := cmd.Start(); err != nil { + stderr, err := cmd.StderrPipe() + if err != nil { return nil, err } + // orchestrate `allOfFailFast` + // create a `g` group and a new `gCtx` derived from our timeout `ctx`. + g, gCtx := errgroup.WithContext(cmdCtx) - // Read all output at once (batch mode) - much faster than incremental parsing - out, err := io.ReadAll(stdout) - if err != nil { + // Start the process + if err = cmd.Start(); err != nil { return nil, err } - // Wait for command to complete - if err := cmd.Wait(); err != nil { - b.logger.Error("Bazel query failed: %v", err) - return nil, err + // Goroutine 1: stream and parse targets from stdout + g.Go(func() error { + queryResults, err = streamAndParseTargets(gCtx, stdout) + return err + }) + + // Goroutine 2: stream stderr + g.Go(func() error { + return streamOutput(gCtx, stderr, &stderrBuf) + }) + + // Wait for both to complete to prevent zombie processes and goroutine leaks + streamErr := g.Wait() + waitErr := cmd.Wait() + + // Get stderr for error messages + stderrStr := stderrBuf.String() + + // Handle errors with full context (both have completed at this point) + if waitErr != nil { + b.logger.Error("Bazel query process failed", zap.Error(waitErr), zap.String("stderr", stderrStr)) + return nil, fmt.Errorf("bazel query process failed: %w\nstderr:\n%s", waitErr, stderrStr) } - // Single unmarshal operation - much faster than per-target streaming - var qr buildpb.QueryResult - if err := proto.Unmarshal(out, &qr); err != nil { - b.logger.Error("Failed to unmarshal query result: %v", err) - return nil, err + if streamErr != nil { + b.logger.Error("Stream processing failed", zap.Error(streamErr), zap.String("stderr", stderrStr)) + return nil, fmt.Errorf("error parsing bazel query output: %w\nstderr:\n%s", streamErr, stderrStr) } - b.logger.Debugf("Parsed %d targets from bazel query", len(qr.Target)) - return &qr, nil + // Both succeeded + b.logger.Debugf("Parsed %d targets from bazel query", len(queryResults.Target)) + return queryResults, nil } + // FromFile reads a proto file generated by bazel query. func FromFile(path string) (*buildpb.QueryResult, error) { var f io.ReadCloser diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index 27be952e..7b1d848e 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -34,15 +34,14 @@ func TestExecuteQuery_Success(t *testing.T) { }, } - // Create a QueryResult with the target (batch proto format, not protodelim streaming) - queryResult := &buildpb.QueryResult{ - Target: []*buildpb.Target{target}, - } - protoData, err := proto.Marshal(queryResult) + // Create delimited proto data (streaming format) + var buf bytes.Buffer + err := writeDelimitedTarget(&buf, target) require.NoError(t, err) gomock.InOrder( - mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(bytes.NewReader(protoData)), nil), + mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(&buf), nil), + mockCmd.EXPECT().StderrPipe().Return(io.NopCloser(strings.NewReader("")), nil), mockCmd.EXPECT().Start().Return(nil), mockCmd.EXPECT().Wait().Return(nil), ) @@ -65,6 +64,24 @@ func TestExecuteQuery_Success(t *testing.T) { assert.Equal(t, &ruleClass, resp.Result.Target[0].Rule.RuleClass) } +// Helper to write delimited proto messages +func writeDelimitedTarget(w io.Writer, target *buildpb.Target) error { + data, err := proto.Marshal(target) + if err != nil { + return err + } + // Write varint length prefix + length := len(data) + for length >= 0x80 { + w.Write([]byte{byte(length) | 0x80}) + length >>= 7 + } + w.Write([]byte{byte(length)}) + // Write message + _, err = w.Write(data) + return err +} + func TestExecuteQuery_WithStartupOptions(t *testing.T) { defer goleak.VerifyNone(t) ctrl := gomock.NewController(t) @@ -80,16 +97,15 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { }, } - // Create a QueryResult with the target (batch proto format) - queryResult := &buildpb.QueryResult{ - Target: []*buildpb.Target{target}, - } - protoData, err := proto.Marshal(queryResult) + // Create delimited proto data + var buf bytes.Buffer + err := writeDelimitedTarget(&buf, target) require.NoError(t, err) var capturedArgs []string gomock.InOrder( - mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(bytes.NewReader(protoData)), nil), + mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(&buf), nil), + mockCmd.EXPECT().StderrPipe().Return(io.NopCloser(strings.NewReader("")), nil), mockCmd.EXPECT().Start().Return(nil), mockCmd.EXPECT().Wait().Return(nil), ) @@ -112,14 +128,15 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { require.NoError(t, err) require.NotNil(t, resp) - // Verify command structure: bazel query --output=proto + // Verify command structure: bazel query --order_output=no --output=streamed_proto require.Equal(t, []string{ "--bazelrc=/custom/.bazelrc", "--output_base=/tmp/bazel", "query", - "--keep_going", - "--output=proto", "//...", + "--order_output=no", + "--output=streamed_proto", + "--keep_going", }, capturedArgs) } @@ -129,13 +146,17 @@ func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { mockCmd := commandermock.NewMockcommander(ctrl) prStdout, pwStdout := io.Pipe() + prStderr, pwStderr := io.Pipe() // Set up the mock expectations in the exact order they will be called. - // Note: When io.ReadAll gets an error (like context deadline), we return immediately - // without calling Wait(), so we don't expect Wait() to be called here. 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 { + // Wait will be called after timeout + return context.DeadlineExceeded + }), ) client, err := NewBazelClient(Params{ @@ -152,10 +173,9 @@ func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { go func() { <-ctx.Done() // Wait for the timeout to fire - // "Killing" the process: close the pipe. - // This unblocks the Read() call in io.ReadAll. - // We close with the context's error so the read sees it. + // "Killing" the process: close the pipes. pwStdout.CloseWithError(ctx.Err()) + pwStderr.CloseWithError(ctx.Err()) }() return mockCmd }, @@ -163,7 +183,7 @@ func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { require.NoError(t, err) result, err := client.executeQueryInternal(context.Background(), "//...", nil) require.Nil(t, result) - assert.Contains(t, err.Error(), "context deadline exceeded") + assert.Error(t, err) } func TestExecuteQueryInternal_Failures(t *testing.T) { @@ -181,10 +201,20 @@ func TestExecuteQueryInternal_Failures(t *testing.T) { 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", @@ -194,11 +224,12 @@ func TestExecuteQueryInternal_Failures(t *testing.T) { 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: true, // Changed from false - implementation returns nil on Wait() error + expectNilResult: true, }, } @@ -247,6 +278,7 @@ func TestExecuteQuery_ErrorCase(t *testing.T) { return mockCmd }, }) + require.NoError(t, err) resp, err := client.ExecuteQuery(context.Background(), &QueryRequest{Query: "//..."}) require.Error(t, err) diff --git a/core/bazel/stream.go b/core/bazel/stream.go new file mode 100644 index 00000000..e45e1400 --- /dev/null +++ b/core/bazel/stream.go @@ -0,0 +1,72 @@ +package bazel + +import ( + "bufio" + "context" + "fmt" + "io" + + buildpb "github.com/bazelbuild/buildtools/build_proto" + "google.golang.org/protobuf/encoding/protodelim" +) + +// streamOutput copies data from src to dst, checking context periodically +func streamOutput(ctx context.Context, src io.Reader, dst io.Writer) error { + buf := make([]byte, 32*1024) + for { + n, readErr := src.Read(buf) + if n > 0 { + if _, writeErr := dst.Write(buf[:n]); writeErr != nil { + // Write failed but MUST keep reading to drain pipe + for { + if _, err := src.Read(buf); err != nil { + break + } + } + return writeErr + } + } + if readErr == io.EOF { + return nil + } + if readErr != nil { + return readErr + } + } +} + +// streamAndParseTargets reads delimited Target protos from src +func streamAndParseTargets(ctx context.Context, src io.Reader) (*buildpb.QueryResult, error) { + result := &buildpb.QueryResult{ + Target: make([]*buildpb.Target, 0), + } + + br := bufio.NewReader(src) + unmarshalOpts := protodelim.UnmarshalOptions{ + MaxSize: 64 * 1024 * 1024, // 64MB limit + } + + var parseErr error + for { + var target buildpb.Target + err := unmarshalOpts.UnmarshalFrom(br, &target) + if err == io.EOF { + break + } + if err != nil { + // Save first error but MUST continue to EOF to drain pipe + if parseErr == nil { + parseErr = fmt.Errorf("failed to unmarshal target: %w", err) + } + // Continue reading - critical to prevent Bazel from blocking on write + continue + } + + // Only collect targets if no error yet + if parseErr == nil { + result.Target = append(result.Target, &target) + } + } + + return result, parseErr +} diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 022d0adc..ae0f38bf 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -17,8 +17,8 @@ import ( "time" "github.com/uber/tango/core/bazel" - "go.uber.org/zap" "github.com/uber/tango/core/targethasher" + "go.uber.org/zap" ) func main() { @@ -77,7 +77,7 @@ func run() error { start := time.Now() resp, err := client.ExecuteQuery(ctx, req) elapsed := time.Since(start) - cancel() + defer cancel() if err != nil { return fmt.Errorf("run %d: query failed: %w", i+1, err) From 2235c13821249156357a15f6c7550b56d0712b2e Mon Sep 17 00:00:00 2001 From: yushan Date: Wed, 25 Feb 2026 22:41:40 +0000 Subject: [PATCH 5/9] Update implementation --- core/bazel/query.go | 26 +++++--------------------- core/bazel/query_test.go | 19 +++++++++---------- core/bazel/stream.go | 5 ++--- 3 files changed, 16 insertions(+), 34 deletions(-) diff --git a/core/bazel/query.go b/core/bazel/query.go index 3d7c18c6..35fc7c4a 100644 --- a/core/bazel/query.go +++ b/core/bazel/query.go @@ -11,7 +11,6 @@ import ( buildpb "github.com/bazelbuild/buildtools/build_proto" "go.uber.org/zap" - "golang.org/x/sync/errgroup" "google.golang.org/protobuf/proto" ) @@ -48,7 +47,6 @@ func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, st 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 { return nil, err @@ -57,44 +55,30 @@ func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, st if err != nil { return nil, err } - // orchestrate `allOfFailFast` - // create a `g` group and a new `gCtx` derived from our timeout `ctx`. - g, gCtx := errgroup.WithContext(cmdCtx) // Start the process if err = cmd.Start(); err != nil { return nil, err } - // Goroutine 1: stream and parse targets from stdout - g.Go(func() error { - queryResults, err = streamAndParseTargets(gCtx, stdout) - return err - }) + queryResults, err = streamAndParseTargets(stdout) - // Goroutine 2: stream stderr - g.Go(func() error { - return streamOutput(gCtx, stderr, &stderrBuf) - }) + err = streamOutput(stderr, &stderrBuf) + if err != nil { + return nil, err + } // Wait for both to complete to prevent zombie processes and goroutine leaks - streamErr := g.Wait() waitErr := cmd.Wait() // Get stderr for error messages stderrStr := stderrBuf.String() - // Handle errors with full context (both have completed at this point) if waitErr != nil { b.logger.Error("Bazel query process failed", zap.Error(waitErr), zap.String("stderr", stderrStr)) return nil, fmt.Errorf("bazel query process failed: %w\nstderr:\n%s", waitErr, stderrStr) } - if streamErr != nil { - b.logger.Error("Stream processing failed", zap.Error(streamErr), zap.String("stderr", stderrStr)) - return nil, fmt.Errorf("error parsing bazel query output: %w\nstderr:\n%s", streamErr, stderrStr) - } - // Both succeeded b.logger.Debugf("Parsed %d targets from bazel query", len(queryResults.Target)) return queryResults, nil diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index 7b1d848e..151af32f 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -164,18 +164,15 @@ func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { WorkspacePath: "/tmp/test", Logger: zap.NewNop().Sugar(), EnvVarsMap: map[string]string{}, - QueryTimeout: 1 * time.Nanosecond, // Induce timeout immediately + QueryTimeout: 10 * time.Millisecond, // Short timeout for test ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander { - // This goroutine simulates the OS/exec.Cmd behavior: - // When the context is canceled, the process is "killed", - // which closes its stdout/stderr pipes. + // Simulate process behavior: when context is cancelled, close pipes go func() { - <-ctx.Done() // Wait for the timeout to fire - - // "Killing" the process: close the pipes. - pwStdout.CloseWithError(ctx.Err()) - pwStderr.CloseWithError(ctx.Err()) + <-ctx.Done() + // Close pipes to unblock readers + pwStdout.Close() + pwStderr.Close() }() return mockCmd }, @@ -183,7 +180,9 @@ func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { require.NoError(t, err) result, err := client.executeQueryInternal(context.Background(), "//...", nil) require.Nil(t, result) - assert.Error(t, err) + require.Error(t, err) + // Should get timeout or deadline exceeded error + assert.Contains(t, err.Error(), "deadline exceeded") } func TestExecuteQueryInternal_Failures(t *testing.T) { diff --git a/core/bazel/stream.go b/core/bazel/stream.go index e45e1400..4db15d80 100644 --- a/core/bazel/stream.go +++ b/core/bazel/stream.go @@ -2,7 +2,6 @@ package bazel import ( "bufio" - "context" "fmt" "io" @@ -11,7 +10,7 @@ import ( ) // streamOutput copies data from src to dst, checking context periodically -func streamOutput(ctx context.Context, src io.Reader, dst io.Writer) error { +func streamOutput(src io.Reader, dst io.Writer) error { buf := make([]byte, 32*1024) for { n, readErr := src.Read(buf) @@ -36,7 +35,7 @@ func streamOutput(ctx context.Context, src io.Reader, dst io.Writer) error { } // streamAndParseTargets reads delimited Target protos from src -func streamAndParseTargets(ctx context.Context, src io.Reader) (*buildpb.QueryResult, error) { +func streamAndParseTargets(src io.Reader) (*buildpb.QueryResult, error) { result := &buildpb.QueryResult{ Target: make([]*buildpb.Target, 0), } From 0d4709c42bfa272b9808e752456e1e36036f04a6 Mon Sep 17 00:00:00 2001 From: yushan Date: Wed, 25 Feb 2026 23:09:59 +0000 Subject: [PATCH 6/9] update streaming --- core/bazel/query.go | 59 ++++++++++++++-------------- core/bazel/query_test.go | 49 ++++++------------------ core/bazel/stream.go | 83 +++++++++++++++++++++------------------- 3 files changed, 86 insertions(+), 105 deletions(-) diff --git a/core/bazel/query.go b/core/bazel/query.go index 35fc7c4a..39524960 100644 --- a/core/bazel/query.go +++ b/core/bazel/query.go @@ -4,27 +4,24 @@ import ( "bytes" "compress/gzip" "context" - "fmt" "io" "os" "strings" 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 { - // Build command: bazel query --order_output=no --output=streamed_proto - args := make([]string, 0, len(startupOptions)+1+len(additionalArgs)+3) + // Build command: bazel query --output=streamed_proto + args := make([]string, 0, len(startupOptions)+1+len(additionalArgs)+2) args = append(args, startupOptions...) args = append(args, "query") - args = append(args, query) - args = append(args, "--order_output=no") - args = append(args, "--output=streamed_proto") args = append(args, additionalArgs...) + args = append(args, "--output=streamed_proto") + args = append(args, query) b.logger.Info("Querying Bazel", zap.String("workspacePath", b.workspacePath), zap.String("query", query)) return b.execCommandContext(ctx, b.bazelCommand, args...) } @@ -37,16 +34,16 @@ func (b *BazelClient) ExecuteQuery(ctx context.Context, req *QueryRequest) (*Que return &QueryResponse{Result: result}, nil } - func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, startupOptions []string, additionalArgs ...string) (*buildpb.QueryResult, error) { var ( - stderrBuf bytes.Buffer - queryResults *buildpb.QueryResult + 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 { return nil, err @@ -55,36 +52,40 @@ func (b *BazelClient) executeQueryInternal(ctx context.Context, query string, st if err != nil { return nil, err } + // 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 } - - queryResults, err = streamAndParseTargets(stdout) - - err = streamOutput(stderr, &stderrBuf) - if err != nil { - return nil, err - } - - // Wait for both to complete to prevent zombie processes and goroutine leaks + // 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, &stderrBuf) + }) waitErr := cmd.Wait() - - // Get stderr for error messages - stderrStr := stderrBuf.String() - // Handle errors with full context (both have completed at this point) + streamErr := g.Wait() + // The command itself failed. if waitErr != nil { - b.logger.Error("Bazel query process failed", zap.Error(waitErr), zap.String("stderr", stderrStr)) - return nil, fmt.Errorf("bazel query process failed: %w\nstderr:\n%s", waitErr, stderrStr) + b.logger.Error("Bazel query failed failed: %v", waitErr) + return queryResults, waitErr + } + // The command succeeded, but there was an error in the stream processing. + if streamErr != nil { + b.logger.Error("Error in stream processing: %v", streamErr) + return nil, streamErr } - - // Both succeeded b.logger.Debugf("Parsed %d targets from bazel query", len(queryResults.Target)) return queryResults, nil } - // FromFile reads a proto file generated by bazel query. func FromFile(path string) (*buildpb.QueryResult, error) { var f io.ReadCloser diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index 151af32f..449c66b8 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -16,7 +16,7 @@ import ( "go.uber.org/goleak" "go.uber.org/mock/gomock" "go.uber.org/zap" - "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/encoding/protodelim" ) func TestExecuteQuery_Success(t *testing.T) { @@ -33,14 +33,11 @@ func TestExecuteQuery_Success(t *testing.T) { RuleClass: &ruleClass, }, } - - // Create delimited proto data (streaming format) - var buf bytes.Buffer - err := writeDelimitedTarget(&buf, target) + var protoData bytes.Buffer + _, err := protodelim.MarshalTo(&protoData, target) require.NoError(t, err) - gomock.InOrder( - mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(&buf), nil), + 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), @@ -64,24 +61,6 @@ func TestExecuteQuery_Success(t *testing.T) { assert.Equal(t, &ruleClass, resp.Result.Target[0].Rule.RuleClass) } -// Helper to write delimited proto messages -func writeDelimitedTarget(w io.Writer, target *buildpb.Target) error { - data, err := proto.Marshal(target) - if err != nil { - return err - } - // Write varint length prefix - length := len(data) - for length >= 0x80 { - w.Write([]byte{byte(length) | 0x80}) - length >>= 7 - } - w.Write([]byte{byte(length)}) - // Write message - _, err = w.Write(data) - return err -} - func TestExecuteQuery_WithStartupOptions(t *testing.T) { defer goleak.VerifyNone(t) ctrl := gomock.NewController(t) @@ -96,15 +75,13 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { RuleClass: &ruleClass, }, } - - // Create delimited proto data - var buf bytes.Buffer - err := writeDelimitedTarget(&buf, target) + var protoData bytes.Buffer + _, err := protodelim.MarshalTo(&protoData, target) require.NoError(t, err) var capturedArgs []string gomock.InOrder( - mockCmd.EXPECT().StdoutPipe().Return(io.NopCloser(&buf), nil), + 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), @@ -128,15 +105,14 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { require.NoError(t, err) require.NotNil(t, resp) - // Verify command structure: bazel query --order_output=no --output=streamed_proto + // Verify command structure: bazel query --output=streamed_proto require.Equal(t, []string{ "--bazelrc=/custom/.bazelrc", "--output_base=/tmp/bazel", "query", - "//...", - "--order_output=no", - "--output=streamed_proto", "--keep_going", + "--output=streamed_proto", + "//...", }, capturedArgs) } @@ -154,7 +130,7 @@ func TestExecuteQueryInternal_ContextTimeout(t *testing.T) { mockCmd.EXPECT().StderrPipe().Return(prStderr, nil), mockCmd.EXPECT().Start().Return(nil), mockCmd.EXPECT().Wait().DoAndReturn(func() error { - // Wait will be called after timeout + // Simulate process ending after timeout return context.DeadlineExceeded }), ) @@ -228,7 +204,7 @@ func TestExecuteQueryInternal_Failures(t *testing.T) { m.EXPECT().Wait().Return(errors.New("command wait failed")) }, expectedError: "command wait failed", - expectNilResult: true, + expectNilResult: false, }, } @@ -277,7 +253,6 @@ func TestExecuteQuery_ErrorCase(t *testing.T) { return mockCmd }, }) - require.NoError(t, err) resp, err := client.ExecuteQuery(context.Background(), &QueryRequest{Query: "//..."}) require.Error(t, err) diff --git a/core/bazel/stream.go b/core/bazel/stream.go index 4db15d80..81c0bcbb 100644 --- a/core/bazel/stream.go +++ b/core/bazel/stream.go @@ -2,69 +2,74 @@ package bazel import ( "bufio" - "fmt" "io" - + "context" buildpb "github.com/bazelbuild/buildtools/build_proto" "google.golang.org/protobuf/encoding/protodelim" ) -// streamOutput copies data from src to dst, checking context periodically -func streamOutput(src io.Reader, dst io.Writer) error { - buf := make([]byte, 32*1024) - for { - n, readErr := src.Read(buf) - if n > 0 { - if _, writeErr := dst.Write(buf[:n]); writeErr != nil { - // Write failed but MUST keep reading to drain pipe - for { - if _, err := src.Read(buf); err != nil { - break - } - } - return writeErr - } - } - if readErr == io.EOF { - return nil - } - if readErr != nil { - return readErr - } +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 } } -// streamAndParseTargets reads delimited Target protos from src -func streamAndParseTargets(src io.Reader) (*buildpb.QueryResult, error) { - result := &buildpb.QueryResult{ - Target: make([]*buildpb.Target, 0), +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) - br := bufio.NewReader(src) - unmarshalOpts := protodelim.UnmarshalOptions{ - MaxSize: 64 * 1024 * 1024, // 64MB limit + go func() { + queryResult, err := getQueryResult(src, dst) + done <- result{queryResult: queryResult, err: err} + }() + + select { + case <-ctx.Done(): + return nil, ctx.Err() + case res := <-done: + return res.queryResult, res.err } +} + + +// getQueryResult reads a QueryResult containing targets from the stream and returns it. +func getQueryResult(src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) { + result := &buildpb.QueryResult{ + Target: make([]*buildpb.Target, 0), + } + tr := io.TeeReader(src, dst) + br := bufio.NewReader(tr) + // unmarshalOpts := protodelim.UnmarshalOptions{ + // MaxSize: 64 * 1024 * 1024, // 64MB limit + // } var parseErr error for { var target buildpb.Target - err := unmarshalOpts.UnmarshalFrom(br, &target) + err := protodelim.UnmarshalFrom(br, &target) if err == io.EOF { break } if err != nil { - // Save first error but MUST continue to EOF to drain pipe if parseErr == nil { - parseErr = fmt.Errorf("failed to unmarshal target: %w", err) + parseErr = err } // Continue reading - critical to prevent Bazel from blocking on write continue } - - // Only collect targets if no error yet - if parseErr == nil { - result.Target = append(result.Target, &target) - } + result.Target = append(result.Target, &target) } return result, parseErr From a4eb6b0491ec240dc7b6641850521f2e55a49dc4 Mon Sep 17 00:00:00 2001 From: yushan Date: Wed, 25 Feb 2026 23:10:28 +0000 Subject: [PATCH 7/9] update --- core/bazel/stream.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/core/bazel/stream.go b/core/bazel/stream.go index 81c0bcbb..03e9e673 100644 --- a/core/bazel/stream.go +++ b/core/bazel/stream.go @@ -52,13 +52,13 @@ func getQueryResult(src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) } tr := io.TeeReader(src, dst) br := bufio.NewReader(tr) - // unmarshalOpts := protodelim.UnmarshalOptions{ - // MaxSize: 64 * 1024 * 1024, // 64MB limit - // } + unmarshalOpts := protodelim.UnmarshalOptions{ + MaxSize: 64 * 1024 * 1024, // 64MB limit + } var parseErr error for { var target buildpb.Target - err := protodelim.UnmarshalFrom(br, &target) + err := unmarshalOpts.UnmarshalFrom(br, &target) if err == io.EOF { break } From b89b1d292c7e208b122395933e74b634f7ea1c56 Mon Sep 17 00:00:00 2001 From: yushan Date: Wed, 25 Feb 2026 23:34:03 +0000 Subject: [PATCH 8/9] Update --- core/bazel/stream.go | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/core/bazel/stream.go b/core/bazel/stream.go index 03e9e673..6c53e06d 100644 --- a/core/bazel/stream.go +++ b/core/bazel/stream.go @@ -31,7 +31,7 @@ func streamAndParseTargets(ctx context.Context, src io.Reader, dst io.Writer) (* done := make(chan result, 1) go func() { - queryResult, err := getQueryResult(src, dst) + queryResult, err := getQueryResult(ctx,src, dst) done <- result{queryResult: queryResult, err: err} }() @@ -46,7 +46,7 @@ func streamAndParseTargets(ctx context.Context, src io.Reader, dst io.Writer) (* // getQueryResult reads a QueryResult containing targets from the stream and returns it. -func getQueryResult(src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) { +func getQueryResult(ctx context.Context, src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) { result := &buildpb.QueryResult{ Target: make([]*buildpb.Target, 0), } @@ -63,9 +63,7 @@ func getQueryResult(src io.Reader, dst io.Writer) (*buildpb.QueryResult, error) break } if err != nil { - if parseErr == nil { - parseErr = err - } + parseErr = err // Continue reading - critical to prevent Bazel from blocking on write continue } From c56b0e8351845419148f94ddd9784bb602ac5207 Mon Sep 17 00:00:00 2001 From: yushan Date: Wed, 25 Feb 2026 23:34:54 +0000 Subject: [PATCH 9/9] Update --- core/bazel/BUILD.bazel | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/bazel/BUILD.bazel b/core/bazel/BUILD.bazel index 39cf50d4..a9d2922f 100644 --- a/core/bazel/BUILD.bazel +++ b/core/bazel/BUILD.bazel @@ -31,7 +31,7 @@ go_test( "@com_github_bazelbuild_buildtools//build_proto", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", - "@org_golang_google_protobuf//proto", + "@org_golang_google_protobuf//encoding/protodelim", "@org_uber_go_goleak//:goleak", "@org_uber_go_mock//gomock", "@org_uber_go_zap//:zap",