From 43e3b5576fabb715b6e305a84adcd9be41e81d58 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Tue, 14 Jul 2026 13:48:43 -0700 Subject: [PATCH 1/8] config: consolidate streaming chunk-size fields into max_message_bytes Replaces service.streaming.{max_num_targets,max_num_changed_targets, max_num_metadata_entries} with a single service.max_message_bytes byte budget, matching config/README.md. The three separate counts all bounded the same underlying thing (gRPC message size) and it wasn't obvious to an operator which one to tune. config.ChunkSizesForByteBudget converts the budget into the three internal per-entry-type limits once, at wiring time, using fixed size assumptions (a target ~40KB worst case, a changed target ~2x that, a metadata entry ~85 bytes) -- not by measuring real messages per request, so there's no added per-request cost. Default of 4,250,000 bytes matches today's metadata chunk size exactly (50,000 x ~85 bytes). The byte-budget conversion lives in the config package (config-specific concern), not core/common, which keeps only its own small, independent fallback constants for chunkTargets/ChunkMetadata's <=0 safety net. --- config/BUILD.bazel | 6 ++- config/chunk_sizes.go | 45 +++++++++++++++++++ config/chunk_sizes_test.go | 69 +++++++++++++++++++++++++++++ config/config_test.go | 34 ++++++++++++++ config/service_config.go | 22 +++------ controller/BUILD.bazel | 1 + controller/controller.go | 24 +++------- controller/testhelper_test.go | 9 ++-- core/common/utils.go | 17 +++---- example/cmd/query-bench/BUILD.bazel | 1 + example/cmd/query-bench/main.go | 4 +- example/main.go | 8 ++-- orchestrator/native_orchestrator.go | 6 +-- orchestrator/testdata/config.yaml | 5 +-- 14 files changed, 188 insertions(+), 63 deletions(-) create mode 100644 config/chunk_sizes.go create mode 100644 config/chunk_sizes_test.go diff --git a/config/BUILD.bazel b/config/BUILD.bazel index bf0b6118..f9fd63e5 100644 --- a/config/BUILD.bazel +++ b/config/BUILD.bazel @@ -3,6 +3,7 @@ load("@rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "config", srcs = [ + "chunk_sizes.go", "config.go", "repository_config.go", "service_config.go", @@ -15,7 +16,10 @@ go_library( go_test( name = "config_test", - srcs = ["config_test.go"], + srcs = [ + "chunk_sizes_test.go", + "config_test.go", + ], embed = [":config"], deps = [ "@com_github_stretchr_testify//assert", diff --git a/config/chunk_sizes.go b/config/chunk_sizes.go new file mode 100644 index 00000000..db6e5825 --- /dev/null +++ b/config/chunk_sizes.go @@ -0,0 +1,45 @@ +// 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 config + +const ( + // DefaultMaxMessageBytes is the default value for ServiceConfig.MaxMessageBytes. + // Matches today's metadata chunk size exactly (50,000 entries × ~85 bytes). + DefaultMaxMessageBytes = 4_250_000 + + // bytesPerTargetEstimate is a worst-case size assumption for a single OptimizedTarget + // (target with ~10K direct deps × 4 bytes), used to convert a byte budget into a chunk size. + bytesPerTargetEstimate = 40_000 + + // bytesPerMetadataEntryEstimate is a worst-case size assumption for a single metadata + // map entry (~60-char avg target name + proto overhead), used the same way. + bytesPerMetadataEntryEstimate = 85 +) + +// ChunkSizesForByteBudget converts a MaxMessageBytes budget into the max number of +// targets, changed targets, and metadata entries per gRPC stream message, so responses +// stay under the 64MB default gRPC per-message limit. The conversion happens once — at +// controller/orchestrator wiring time, not per-request — using fixed size assumptions +// rather than measuring real messages. A ChangedTarget carries both an old and new +// target, so it gets half the plain-target budget. maxMessageBytes <= 0 falls back to +// DefaultMaxMessageBytes. +func ChunkSizesForByteBudget(maxMessageBytes int) (targetChunkSize, changedTargetChunkSize, metadataMapChunkSize int) { + if maxMessageBytes <= 0 { + maxMessageBytes = DefaultMaxMessageBytes + } + return max(1, maxMessageBytes/bytesPerTargetEstimate), + max(1, maxMessageBytes/(2*bytesPerTargetEstimate)), + max(1, maxMessageBytes/bytesPerMetadataEntryEstimate) +} diff --git a/config/chunk_sizes_test.go b/config/chunk_sizes_test.go new file mode 100644 index 00000000..c08af30f --- /dev/null +++ b/config/chunk_sizes_test.go @@ -0,0 +1,69 @@ +// 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 config + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestChunkSizesForByteBudget(t *testing.T) { + tests := []struct { + name string + giveMaxMessageBytes int + wantTargetChunkSize int + wantChangedTargetChunkSize int + wantMetadataMapChunkSize int + }{ + { + name: "non-positive falls back to default budget", + giveMaxMessageBytes: 0, + wantTargetChunkSize: 106, + wantChangedTargetChunkSize: 53, + wantMetadataMapChunkSize: 50_000, + }, + { + name: "negative falls back to default budget", + giveMaxMessageBytes: -1, + wantTargetChunkSize: 106, + wantChangedTargetChunkSize: 53, + wantMetadataMapChunkSize: 50_000, + }, + { + name: "custom budget scales proportionally", + giveMaxMessageBytes: 8_500_000, + wantTargetChunkSize: 212, + wantChangedTargetChunkSize: 106, + wantMetadataMapChunkSize: 100_000, + }, + { + name: "tiny budget never returns zero", + giveMaxMessageBytes: 1, + wantTargetChunkSize: 1, + wantChangedTargetChunkSize: 1, + wantMetadataMapChunkSize: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + targetChunkSize, changedTargetChunkSize, metadataMapChunkSize := ChunkSizesForByteBudget(tt.giveMaxMessageBytes) + assert.Equal(t, tt.wantTargetChunkSize, targetChunkSize) + assert.Equal(t, tt.wantChangedTargetChunkSize, changedTargetChunkSize) + assert.Equal(t, tt.wantMetadataMapChunkSize, metadataMapChunkSize) + }) + } +} diff --git a/config/config_test.go b/config/config_test.go index ce7ff74f..4288ac9a 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -70,6 +70,40 @@ repository: } } +func TestParse_ServiceDefaults(t *testing.T) { + tests := []struct { + name string + give string + wantMaxMessageBytes int + }{ + { + name: "max_message_bytes unset", + give: _baseServiceYAML + ` +repository: + - remote: "r1" +`, + wantMaxMessageBytes: 0, + }, + { + name: "max_message_bytes explicit value preserved", + give: _baseServiceYAML + ` + max_message_bytes: 1000000 +repository: + - remote: "r1" +`, + wantMaxMessageBytes: 1000000, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg, err := Parse(writeConfig(t, tt.give)) + require.NoError(t, err) + assert.Equal(t, tt.wantMaxMessageBytes, cfg.Service.MaxMessageBytes) + }) + } +} + func TestParse_RepositoryDefaults(t *testing.T) { tests := []struct { name string diff --git a/config/service_config.go b/config/service_config.go index 4c0aa09e..c04b3fdb 100644 --- a/config/service_config.go +++ b/config/service_config.go @@ -24,21 +24,9 @@ type ServiceConfig struct { // and worker checkouts. Required. Layout: // for // origin clones and /.workers//worker-{1..N}/ for // worker checkouts. - WorkspacesRootPath string `yaml:"workspaces_root_path"` - Streaming ChunkConfig `yaml:"streaming"` // streaming chunk sizes; zero values fall back to package defaults -} - -// ChunkConfig controls the number of entries per gRPC stream message. -// All fields are optional; a zero value means "use the package default". -// Tune these when a monorepo's per-target size causes messages to approach -// gRPC's default 4MB per-message limit. -type ChunkConfig struct { - // MaxNumTargets is the max number of OptimizedTarget entries per stream message. - MaxNumTargets int `yaml:"max_num_targets"` - // MaxNumChangedTargets is the max number of ChangedTarget entries per stream message. - // ChangedTarget carries both old and new targets (~2× the size of a regular target). - MaxNumChangedTargets int `yaml:"max_num_changed_targets"` - // MaxNumMetadataEntries is the max number of entries per metadata map chunk. - // Applies to target_id_mapping and attribute_string_value_mapping. - MaxNumMetadataEntries int `yaml:"max_num_metadata_entries"` + WorkspacesRootPath string `yaml:"workspaces_root_path"` + // MaxMessageBytes caps the size of each gRPC stream message (targets, changed + // targets, and metadata mappings). Zero falls back to common.DefaultMaxMessageBytes. + // See common.ChunkSizesForByteBudget for how this is converted into per-entry-type limits. + MaxMessageBytes int `yaml:"max_message_bytes"` } diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 9ec9c14d..62afc2f2 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -42,6 +42,7 @@ go_test( ], embed = [":controller"], deps = [ + "//config", "//core/common", "//core/storage", "//core/storage/storagemock", diff --git a/controller/controller.go b/controller/controller.go index 30523806..2c4feaef 100644 --- a/controller/controller.go +++ b/controller/controller.go @@ -20,7 +20,6 @@ import ( "github.com/uber-go/tally" "github.com/uber/tango/config" - "github.com/uber/tango/core/common" "github.com/uber/tango/core/storage" "github.com/uber/tango/orchestrator" pb "github.com/uber/tango/tangopb" @@ -31,11 +30,11 @@ import ( // Params are the parameters for the controller. type Params struct { fx.In - Logger *zap.Logger - Storage storage.Storage - Orchestrator orchestrator.Orchestrator - Scope tally.Scope `optional:"true"` - ChunkConfig config.ChunkConfig `optional:"true"` + Logger *zap.Logger + Storage storage.Storage + Orchestrator orchestrator.Orchestrator + Scope tally.Scope `optional:"true"` + MaxMessageBytes int `optional:"true"` } // _totalDurationBuckets covers 0–15 minutes in 10-second linear intervals. @@ -64,18 +63,7 @@ func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { if scope == nil { scope = tally.NoopScope } - targetChunkSize := p.ChunkConfig.MaxNumTargets - if targetChunkSize <= 0 { - targetChunkSize = common.DefaultTargetChunkSize - } - changedTargetChunkSize := p.ChunkConfig.MaxNumChangedTargets - if changedTargetChunkSize <= 0 { - changedTargetChunkSize = common.DefaultChangedTargetChunkSize - } - metadataMapChunkSize := p.ChunkConfig.MaxNumMetadataEntries - if metadataMapChunkSize <= 0 { - metadataMapChunkSize = common.DefaultMetadataMapChunkSize - } + targetChunkSize, changedTargetChunkSize, metadataMapChunkSize := config.ChunkSizesForByteBudget(p.MaxMessageBytes) return &controller{ logger: p.Logger, storage: p.Storage, diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go index 56ff4396..2d51d57e 100644 --- a/controller/testhelper_test.go +++ b/controller/testhelper_test.go @@ -18,17 +18,18 @@ import ( "context" "github.com/uber-go/tally" - "github.com/uber/tango/core/common" + "github.com/uber/tango/config" "go.uber.org/zap" ) func newTestController(logger *zap.Logger) *controller { + targetChunkSize, changedTargetChunkSize, metadataMapChunkSize := config.ChunkSizesForByteBudget(config.DefaultMaxMessageBytes) return &controller{ logger: logger, scope: tally.NoopScope, - targetChunkSize: common.DefaultTargetChunkSize, - changedTargetChunkSize: common.DefaultChangedTargetChunkSize, - metadataMapChunkSize: common.DefaultMetadataMapChunkSize, + targetChunkSize: targetChunkSize, + changedTargetChunkSize: changedTargetChunkSize, + metadataMapChunkSize: metadataMapChunkSize, totalDurationBuckets: _totalDurationBuckets, appCtx: context.Background(), } diff --git a/core/common/utils.go b/core/common/utils.go index e0f47d47..7a17c371 100644 --- a/core/common/utils.go +++ b/core/common/utils.go @@ -26,20 +26,17 @@ import ( ) const ( - // DefaultTargetChunkSize is the default number of OptimizedTarget entries per stream message. - // Sized conservatively: at ~40KB/target worst-case (target with ~10K direct deps × 4 bytes), - // 250 targets ≈ 10MB — well under the 64MB default gRPC per-message limit. + // DefaultTargetChunkSize is the fallback number of OptimizedTarget entries per stream + // message, used only if chunkTargets is called with a non-positive chunk size. DefaultTargetChunkSize = 250 - // DefaultChangedTargetChunkSize is the default number of ChangedTarget entries per stream message. - // A ChangedTarget carries both old_target and new_target (2× an OptimizedTarget), so we use - // half the regular chunk size to stay within the same byte budget. + // DefaultChangedTargetChunkSize is the fallback number of ChangedTarget entries per + // stream message. A ChangedTarget carries both old_target and new_target (2× an + // OptimizedTarget), so it gets half the plain-target fallback. DefaultChangedTargetChunkSize = 125 - // DefaultMetadataMapChunkSize is the max entries per metadata message chunk. - // target_id_mapping and attribute_string_value_mapping scale with repo size and can exceed - // the 64MB gRPC message limit for large monorepos, so they are split across multiple messages. - // At ~85 bytes/entry (60-char avg target name + proto overhead), 50 000 entries ≈ 4.25MB per chunk. + // DefaultMetadataMapChunkSize is the fallback max entries per metadata message chunk, + // used only if ChunkMetadata is called with a non-positive chunk size. DefaultMetadataMapChunkSize = 50_000 ) diff --git a/example/cmd/query-bench/BUILD.bazel b/example/cmd/query-bench/BUILD.bazel index ba71ce1c..6c798b8e 100644 --- a/example/cmd/query-bench/BUILD.bazel +++ b/example/cmd/query-bench/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/tango/example/cmd/query-bench", visibility = ["//visibility:private"], deps = [ + "//config", "//core/bazel", "//core/common", "//core/targethasher", diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 28e4e52b..0885a617 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -31,6 +31,7 @@ import ( "time" "github.com/gogo/protobuf/jsonpb" + "github.com/uber/tango/config" "github.com/uber/tango/core/bazel" "github.com/uber/tango/core/common" "github.com/uber/tango/core/targethasher" @@ -112,9 +113,10 @@ func run() error { totalDuration += elapsed fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) start = time.Now() + targetChunkSize, _, metadataMapChunkSize := config.ChunkSizesForByteBudget(config.DefaultMaxMessageBytes) response, err := common.ResultToGetTargetGraphResponse( ctx, targethasherResult, - common.DefaultTargetChunkSize, common.DefaultMetadataMapChunkSize, + targetChunkSize, metadataMapChunkSize, ) if err != nil { return fmt.Errorf("run %d: converting to GetTargetGraphResponse: %w", i+1, err) diff --git a/example/main.go b/example/main.go index 83043a61..66c51edc 100644 --- a/example/main.go +++ b/example/main.go @@ -100,10 +100,10 @@ func run() error { // Controller (YARPC server implementation). appCtx is forwarded so the // controller's background goroutines are tied to process lifetime. ctrl := controller.NewController(appCtx, controller.Params{ - Logger: zl, - Storage: store, - Orchestrator: orch, - ChunkConfig: cfg.Service.Streaming, + Logger: zl, + Storage: store, + Orchestrator: orch, + MaxMessageBytes: cfg.Service.MaxMessageBytes, }) // YARPC transports and dispatcher diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index 3f9ab40a..28628521 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -208,10 +208,8 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - responses, err := common.ResultToGetTargetGraphResponse(ctx, result, - b.config.Service.Streaming.MaxNumTargets, - b.config.Service.Streaming.MaxNumMetadataEntries, - ) + targetChunkSize, _, metadataMapChunkSize := config.ChunkSizesForByteBudget(b.config.Service.MaxMessageBytes) + responses, err := common.ResultToGetTargetGraphResponse(ctx, result, targetChunkSize, metadataMapChunkSize) if err != nil { return nil, fmt.Errorf("convert target graph to response: %w", err) } diff --git a/orchestrator/testdata/config.yaml b/orchestrator/testdata/config.yaml index 93f44392..2ebbc164 100644 --- a/orchestrator/testdata/config.yaml +++ b/orchestrator/testdata/config.yaml @@ -6,7 +6,4 @@ repository: service: max_worker_pool_size: 3 workspaces_root_path: "/tmp/tango-repo-manager" - streaming: - max_num_targets: 250 - max_num_changed_targets: 125 - max_num_metadata_entries: 50000 + max_message_bytes: 4250000 From b3fefe45eb45d82a3448a95a34d491f4a968a2d9 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Tue, 14 Jul 2026 15:15:31 -0700 Subject: [PATCH 2/8] config: measure real proto wire size instead of estimating chunk sizes Bound each streamed message by its actual serialized size (via each message type's generated Size() method) rather than converting max_message_bytes into a fixed per-entry-type chunk count using a guessed average entry size. Introduces internal/msgsplit for all chunking logic (BySize, ChunkMetadata, ResultToGetTargetGraphResponse, DefaultMaxMessageBytes), moving it out of core/common which was accumulating unrelated responsibilities. core/common retains only ToShortRemote and NameIDMapper. --- config/BUILD.bazel | 6 +- config/chunk_sizes.go | 45 ------ config/chunk_sizes_test.go | 69 -------- config/config.go | 8 +- config/config_test.go | 2 +- config/service_config.go | 6 +- controller/BUILD.bazel | 2 - controller/controller.go | 35 ++--- controller/getchangedtargets.go | 19 +-- controller/testhelper_test.go | 14 +- core/common/BUILD.bazel | 13 +- core/common/utils.go | 236 ---------------------------- core/common/utils_test.go | 95 ----------- example/cmd/query-bench/BUILD.bazel | 3 +- example/cmd/query-bench/main.go | 8 +- orchestrator/BUILD.bazel | 1 + orchestrator/native_orchestrator.go | 4 +- 17 files changed, 45 insertions(+), 521 deletions(-) delete mode 100644 config/chunk_sizes.go delete mode 100644 config/chunk_sizes_test.go diff --git a/config/BUILD.bazel b/config/BUILD.bazel index f9fd63e5..bf0b6118 100644 --- a/config/BUILD.bazel +++ b/config/BUILD.bazel @@ -3,7 +3,6 @@ load("@rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "config", srcs = [ - "chunk_sizes.go", "config.go", "repository_config.go", "service_config.go", @@ -16,10 +15,7 @@ go_library( go_test( name = "config_test", - srcs = [ - "chunk_sizes_test.go", - "config_test.go", - ], + srcs = ["config_test.go"], embed = [":config"], deps = [ "@com_github_stretchr_testify//assert", diff --git a/config/chunk_sizes.go b/config/chunk_sizes.go deleted file mode 100644 index db6e5825..00000000 --- a/config/chunk_sizes.go +++ /dev/null @@ -1,45 +0,0 @@ -// 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 config - -const ( - // DefaultMaxMessageBytes is the default value for ServiceConfig.MaxMessageBytes. - // Matches today's metadata chunk size exactly (50,000 entries × ~85 bytes). - DefaultMaxMessageBytes = 4_250_000 - - // bytesPerTargetEstimate is a worst-case size assumption for a single OptimizedTarget - // (target with ~10K direct deps × 4 bytes), used to convert a byte budget into a chunk size. - bytesPerTargetEstimate = 40_000 - - // bytesPerMetadataEntryEstimate is a worst-case size assumption for a single metadata - // map entry (~60-char avg target name + proto overhead), used the same way. - bytesPerMetadataEntryEstimate = 85 -) - -// ChunkSizesForByteBudget converts a MaxMessageBytes budget into the max number of -// targets, changed targets, and metadata entries per gRPC stream message, so responses -// stay under the 64MB default gRPC per-message limit. The conversion happens once — at -// controller/orchestrator wiring time, not per-request — using fixed size assumptions -// rather than measuring real messages. A ChangedTarget carries both an old and new -// target, so it gets half the plain-target budget. maxMessageBytes <= 0 falls back to -// DefaultMaxMessageBytes. -func ChunkSizesForByteBudget(maxMessageBytes int) (targetChunkSize, changedTargetChunkSize, metadataMapChunkSize int) { - if maxMessageBytes <= 0 { - maxMessageBytes = DefaultMaxMessageBytes - } - return max(1, maxMessageBytes/bytesPerTargetEstimate), - max(1, maxMessageBytes/(2*bytesPerTargetEstimate)), - max(1, maxMessageBytes/bytesPerMetadataEntryEstimate) -} diff --git a/config/chunk_sizes_test.go b/config/chunk_sizes_test.go deleted file mode 100644 index c08af30f..00000000 --- a/config/chunk_sizes_test.go +++ /dev/null @@ -1,69 +0,0 @@ -// 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 config - -import ( - "testing" - - "github.com/stretchr/testify/assert" -) - -func TestChunkSizesForByteBudget(t *testing.T) { - tests := []struct { - name string - giveMaxMessageBytes int - wantTargetChunkSize int - wantChangedTargetChunkSize int - wantMetadataMapChunkSize int - }{ - { - name: "non-positive falls back to default budget", - giveMaxMessageBytes: 0, - wantTargetChunkSize: 106, - wantChangedTargetChunkSize: 53, - wantMetadataMapChunkSize: 50_000, - }, - { - name: "negative falls back to default budget", - giveMaxMessageBytes: -1, - wantTargetChunkSize: 106, - wantChangedTargetChunkSize: 53, - wantMetadataMapChunkSize: 50_000, - }, - { - name: "custom budget scales proportionally", - giveMaxMessageBytes: 8_500_000, - wantTargetChunkSize: 212, - wantChangedTargetChunkSize: 106, - wantMetadataMapChunkSize: 100_000, - }, - { - name: "tiny budget never returns zero", - giveMaxMessageBytes: 1, - wantTargetChunkSize: 1, - wantChangedTargetChunkSize: 1, - wantMetadataMapChunkSize: 1, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - targetChunkSize, changedTargetChunkSize, metadataMapChunkSize := ChunkSizesForByteBudget(tt.giveMaxMessageBytes) - assert.Equal(t, tt.wantTargetChunkSize, targetChunkSize) - assert.Equal(t, tt.wantChangedTargetChunkSize, changedTargetChunkSize) - assert.Equal(t, tt.wantMetadataMapChunkSize, metadataMapChunkSize) - }) - } -} diff --git a/config/config.go b/config/config.go index debc9228..287cd6cd 100644 --- a/config/config.go +++ b/config/config.go @@ -21,7 +21,10 @@ import ( yaml "github.com/goccy/go-yaml" ) -const _defaultBazelQueryTimeoutSeconds = 600 +const ( + _defaultBazelQueryTimeoutSeconds = 600 + _defaultMaxMessageBytes = 4_250_000 +) var _bzlmodEnabledDefault = true @@ -67,6 +70,9 @@ func Parse(configFilePath string) (*Config, error) { if config.Service.MaxWorkerPoolSize <= 0 { return nil, fmt.Errorf("service.max_worker_pool_size must be > 0, got %d", config.Service.MaxWorkerPoolSize) } + if config.Service.MaxMessageBytes <= 0 { + config.Service.MaxMessageBytes = _defaultMaxMessageBytes + } config.repositoryByRemote = make(map[string]*RepositoryConfig, len(config.Repository)) for i := range config.Repository { remote := config.Repository[i].Remote diff --git a/config/config_test.go b/config/config_test.go index 4288ac9a..b14ed68f 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -82,7 +82,7 @@ func TestParse_ServiceDefaults(t *testing.T) { repository: - remote: "r1" `, - wantMaxMessageBytes: 0, + wantMaxMessageBytes: _defaultMaxMessageBytes, }, { name: "max_message_bytes explicit value preserved", diff --git a/config/service_config.go b/config/service_config.go index c04b3fdb..fa18bfc7 100644 --- a/config/service_config.go +++ b/config/service_config.go @@ -25,8 +25,8 @@ type ServiceConfig struct { // origin clones and /.workers//worker-{1..N}/ for // worker checkouts. WorkspacesRootPath string `yaml:"workspaces_root_path"` - // MaxMessageBytes caps the size of each gRPC stream message (targets, changed - // targets, and metadata mappings). Zero falls back to common.DefaultMaxMessageBytes. - // See common.ChunkSizesForByteBudget for how this is converted into per-entry-type limits. + // MaxMessageBytes caps the serialized size of each gRPC stream message (targets, + // changed targets, and metadata mappings), measured against each entry's actual + // wire size. Defaults to ~4.25MB when unset. MaxMessageBytes int `yaml:"max_message_bytes"` } diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 62afc2f2..534493d5 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -14,7 +14,6 @@ go_library( importpath = "github.com/uber/tango/controller", visibility = ["//visibility:public"], deps = [ - "//config", "//core/common", "//core/storage", "//entity", @@ -42,7 +41,6 @@ go_test( ], embed = [":controller"], deps = [ - "//config", "//core/common", "//core/storage", "//core/storage/storagemock", diff --git a/controller/controller.go b/controller/controller.go index 2c4feaef..ad4caa20 100644 --- a/controller/controller.go +++ b/controller/controller.go @@ -19,7 +19,6 @@ import ( "time" "github.com/uber-go/tally" - "github.com/uber/tango/config" "github.com/uber/tango/core/storage" "github.com/uber/tango/orchestrator" pb "github.com/uber/tango/tangopb" @@ -28,6 +27,9 @@ import ( ) // Params are the parameters for the controller. +// TODO: MaxMessageBytes should come from *config.Config via fx, not as a +// standalone Params field. Wire it through the fx module once config is +// provided as an fx dependency. type Params struct { fx.In Logger *zap.Logger @@ -41,14 +43,12 @@ type Params struct { var _totalDurationBuckets = tally.MustMakeLinearDurationBuckets(10*time.Second, 10*time.Second, 90) type controller struct { - logger *zap.Logger - storage storage.Storage - orchestrator orchestrator.Orchestrator - scope tally.Scope - targetChunkSize int - changedTargetChunkSize int - metadataMapChunkSize int - totalDurationBuckets tally.Buckets + logger *zap.Logger + storage storage.Storage + orchestrator orchestrator.Orchestrator + scope tally.Scope + maxMessageBytes int + totalDurationBuckets tally.Buckets // appCtx is the application lifetime; cancel it on process shutdown. // Used by linkRequestCtx and any fire-and-forget goroutines so they @@ -63,17 +63,14 @@ func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { if scope == nil { scope = tally.NoopScope } - targetChunkSize, changedTargetChunkSize, metadataMapChunkSize := config.ChunkSizesForByteBudget(p.MaxMessageBytes) return &controller{ - logger: p.Logger, - storage: p.Storage, - orchestrator: p.Orchestrator, - scope: scope.SubScope("controller"), - targetChunkSize: targetChunkSize, - changedTargetChunkSize: changedTargetChunkSize, - metadataMapChunkSize: metadataMapChunkSize, - totalDurationBuckets: _totalDurationBuckets, - appCtx: appCtx, + logger: p.Logger, + storage: p.Storage, + orchestrator: p.Orchestrator, + scope: scope.SubScope("controller"), + maxMessageBytes: p.MaxMessageBytes, + totalDurationBuckets: _totalDurationBuckets, + appCtx: appCtx, } } diff --git a/controller/getchangedtargets.go b/controller/getchangedtargets.go index 287db153..928533cc 100644 --- a/controller/getchangedtargets.go +++ b/controller/getchangedtargets.go @@ -447,33 +447,22 @@ func (c *controller) compareTargetGraphs(ctx context.Context, scope tally.Scope, // Emit changes in chunks to stay within gRPC per-message size limits, followed by chunked metadata. var results []*pb.GetChangedTargetsResponse - for i := 0; i < len(changed); i += c.changedTargetChunkSize { - end := i + c.changedTargetChunkSize - if end > len(changed) { - end = len(changed) - } + for _, chunk := range mapper.BySize(changed, c.maxMessageBytes) { results = append(results, &pb.GetChangedTargetsResponse{ Item: &pb.GetChangedTargetsResponse_ChangedTargets{ ChangedTargets: &pb.ChangedTargets{ - ChangedTargets: changed[i:end], + ChangedTargets: chunk, }, }, }) } - if len(results) == 0 { - results = append(results, &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{}, - }, - }) - } - for _, meta := range common.ChunkMetadata( + for _, meta := range mapper.ChunkMetadata( mappers.target.Invert(), mappers.ruleType.Invert(), mappers.tag.Invert(), mappers.attrName.Invert(), mappers.attrVal.Invert(), - c.metadataMapChunkSize, + c.maxMessageBytes, ) { results = append(results, &pb.GetChangedTargetsResponse{ Item: &pb.GetChangedTargetsResponse_Metadata{ diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go index 2d51d57e..1c470630 100644 --- a/controller/testhelper_test.go +++ b/controller/testhelper_test.go @@ -18,19 +18,15 @@ import ( "context" "github.com/uber-go/tally" - "github.com/uber/tango/config" "go.uber.org/zap" ) func newTestController(logger *zap.Logger) *controller { - targetChunkSize, changedTargetChunkSize, metadataMapChunkSize := config.ChunkSizesForByteBudget(config.DefaultMaxMessageBytes) return &controller{ - logger: logger, - scope: tally.NoopScope, - targetChunkSize: targetChunkSize, - changedTargetChunkSize: changedTargetChunkSize, - metadataMapChunkSize: metadataMapChunkSize, - totalDurationBuckets: _totalDurationBuckets, - appCtx: context.Background(), + logger: logger, + scope: tally.NoopScope, + maxMessageBytes: 4_250_000, + totalDurationBuckets: _totalDurationBuckets, + appCtx: context.Background(), } } diff --git a/core/common/BUILD.bazel b/core/common/BUILD.bazel index 11f4184d..518441f8 100644 --- a/core/common/BUILD.bazel +++ b/core/common/BUILD.bazel @@ -8,22 +8,11 @@ go_library( ], importpath = "github.com/uber/tango/core/common", visibility = ["//visibility:public"], - deps = [ - "//core/targethasher", - "//internal/mapper/idmapper", - "//tangopb", - "@com_github_bazelbuild_buildtools//build_proto", - ], ) go_test( name = "common_test", srcs = ["utils_test.go"], embed = [":common"], - deps = [ - "//core/targethasher", - "//tangopb", - "@com_github_stretchr_testify//assert", - "@com_github_stretchr_testify//require", - ], + deps = ["@com_github_stretchr_testify//assert"], ) diff --git a/core/common/utils.go b/core/common/utils.go index 7a17c371..d8062035 100644 --- a/core/common/utils.go +++ b/core/common/utils.go @@ -15,29 +15,7 @@ package common import ( - "context" - "encoding/hex" "strings" - - buildpb "github.com/bazelbuild/buildtools/build_proto" - "github.com/uber/tango/core/targethasher" - "github.com/uber/tango/internal/mapper/idmapper" - "github.com/uber/tango/tangopb" -) - -const ( - // DefaultTargetChunkSize is the fallback number of OptimizedTarget entries per stream - // message, used only if chunkTargets is called with a non-positive chunk size. - DefaultTargetChunkSize = 250 - - // DefaultChangedTargetChunkSize is the fallback number of ChangedTarget entries per - // stream message. A ChangedTarget carries both old_target and new_target (2× an - // OptimizedTarget), so it gets half the plain-target fallback. - DefaultChangedTargetChunkSize = 125 - - // DefaultMetadataMapChunkSize is the fallback max entries per metadata message chunk, - // used only if ChunkMetadata is called with a non-positive chunk size. - DefaultMetadataMapChunkSize = 50_000 ) // ToShortRemote returns the short remote name given a git ssh remote string. @@ -46,217 +24,3 @@ func ToShortRemote(remote string) string { strs := strings.Split(remote, ":") return strs[len(strs)-1] } - -// 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 = 4096 - -// ResultToGetTargetGraphResponse converts a Result to a GetTargetGraphResponse. -// targetChunkSize controls how many OptimizedTarget entries per stream message. -// metadataMapChunkSize controls how many entries per metadata map chunk. -// TODO: move this function to internal/mapper -func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result, targetChunkSize, metadataMapChunkSize int) ([]*tangopb.GetTargetGraphResponse, error) { - // Map target names to ids. This list is topologically sorted, so the ids are stable. - // IDs start at 1 — 0 is reserved as the proto3 "unset" sentinel so consumers using - // encoding/json (which honors `omitempty` on int32 fields) never silently lose a target. - targetNamesMapping := make(map[string]int32, len(result.TargetNames)) - for i, name := range result.TargetNames { - targetNamesMapping[name] = int32(i + 1) - } - - ruleTypeMapper := idmapper.NewMapper() - tagMapper := idmapper.NewMapper() - attrNameMapper := idmapper.NewMapper() - attrStrValMapper := idmapper.NewMapper() - - // Build the optimized targets slice - optimizedTargets := make([]*tangopb.OptimizedTarget, 0, len(result.Targets)) - - n := 0 - for _, t := range result.Targets { - if n%cancelCheckInterval == 0 { - if err := ctx.Err(); err != nil { - return nil, err - } - } - n++ - nameID := targetNamesMapping[t.Name] - - depIDs := make([]int32, 0, len(t.Deps)) - for _, depName := range t.Deps { - depID, ok := targetNamesMapping[depName] - if !ok { - continue - } - depIDs = append(depIDs, depID) - } - - ot := &tangopb.OptimizedTarget{ - Id: nameID, - Hash: hex.EncodeToString(t.Hash), - DirectDependencies: depIDs, - } - - // RuleType - if t.RuleType != "" { - id := ruleTypeMapper.ID(t.RuleType) - ot.RuleType = id - } - - // Tags - if len(t.Tags) > 0 { - tagIDs := make([]int32, 0, len(t.Tags)) - for _, tag := range t.Tags { - tagIDs = append(tagIDs, tagMapper.ID(tag)) - } - ot.Tags = tagIDs - } - ot.Root = t.Root - ot.External = t.External - if len(t.Attributes) > 0 { - attrs := make(map[int32]int32, len(t.Attributes)) - for _, attr := range t.Attributes { - // Only include STRING attributes with non-nil name and value to avoid nil dereferences. - if attr.GetType() == buildpb.Attribute_STRING && attr.Name != nil && attr.StringValue != nil { - nameID := attrNameMapper.ID(*attr.Name) - valID := attrStrValMapper.ID(*attr.StringValue) - attrs[nameID] = valID - } - } - ot.Attributes = attrs - } - - optimizedTargets = append(optimizedTargets, ot) - } - - // Invert mappings: string -> id => id -> string - targetIDToName := make(map[int32]string, len(targetNamesMapping)) - for s, id := range targetNamesMapping { - targetIDToName[id] = s - } - - ruleTypeIDToName := ruleTypeMapper.Invert() - tagIDToName := tagMapper.Invert() - attrNameIDToName := attrNameMapper.Invert() - attrStrValIDToVal := attrStrValMapper.Invert() - - // chunk targets into multiple messages for streaming - responses := chunkTargets(optimizedTargets, targetChunkSize) - for _, meta := range ChunkMetadata( - targetIDToName, - ruleTypeIDToName, - tagIDToName, - attrNameIDToName, - attrStrValIDToVal, - metadataMapChunkSize, - ) { - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Metadata{Metadata: meta}, - }) - } - - return responses, nil -} - -func chunkTargets(targets []*tangopb.OptimizedTarget, chunkSize int) []*tangopb.GetTargetGraphResponse { - if chunkSize <= 0 { - chunkSize = DefaultTargetChunkSize - } - - // at least one chunk - numChunks := max(1, (len(targets)+chunkSize-1)/chunkSize) - - responses := make([]*tangopb.GetTargetGraphResponse, 0, numChunks) - - for i := 0; i < len(targets); i += chunkSize { - end := i + chunkSize - if end > len(targets) { - end = len(targets) - } - - chunk := targets[i:end] - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Targets{ - Targets: &tangopb.OptimizedTargets{ - Targets: chunk, - }, - }, - }) - } - - // Handle empty targets case - if len(responses) == 0 { - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Targets{ - Targets: &tangopb.OptimizedTargets{ - Targets: []*tangopb.OptimizedTarget{}, - }, - }, - }) - } - - return responses -} - -// ChunkMetadata splits the metadata maps into multiple Metadata messages. -// target_id_mapping and attribute_string_value_mapping scale with repo size and can exceed the -// 64MB gRPC per-message limit for large monorepos; they are split across chunks of chunkSize entries. -// The small maps (rule_type, tag, attribute_name) always fit in one message and are sent in the first chunk. -func ChunkMetadata( - targetIDToName map[int32]string, - ruleTypeIDToName map[int32]string, - tagIDToName map[int32]string, - attrNameIDToName map[int32]string, - attrStrValIDToVal map[int32]string, - chunkSize int, -) []*tangopb.Metadata { - if chunkSize <= 0 { - chunkSize = DefaultMetadataMapChunkSize - } - - targetChunks := splitMap(targetIDToName, chunkSize) - attrValChunks := splitMap(attrStrValIDToVal, chunkSize) - - numChunks := max(1, max(len(targetChunks), len(attrValChunks))) - chunks := make([]*tangopb.Metadata, 0, numChunks) - - for i := range numChunks { - meta := &tangopb.Metadata{} - // Small maps are always small enough to fit in one message; include them in the first chunk. - if i == 0 { - meta.RuleTypeMapping = ruleTypeIDToName - meta.TagMapping = tagIDToName - meta.AttributeNameMapping = attrNameIDToName - } - if i < len(targetChunks) { - meta.TargetIdMapping = targetChunks[i] - } - if i < len(attrValChunks) { - meta.AttributeStringValueMapping = attrValChunks[i] - } - chunks = append(chunks, meta) - } - - return chunks -} - -// splitMap splits a map[int32]string into slices of at most size entries each. -func splitMap(m map[int32]string, size int) []map[int32]string { - if len(m) == 0 { - return nil - } - chunks := make([]map[int32]string, 0, (len(m)+size-1)/size) - current := make(map[int32]string, size) - for k, v := range m { - current[k] = v - if len(current) >= size { - chunks = append(chunks, current) - current = make(map[int32]string, size) - } - } - if len(current) > 0 { - chunks = append(chunks, current) - } - return chunks -} diff --git a/core/common/utils_test.go b/core/common/utils_test.go index 63e5eb24..ddcd9969 100644 --- a/core/common/utils_test.go +++ b/core/common/utils_test.go @@ -15,14 +15,9 @@ package common import ( - "context" - "fmt" "testing" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "github.com/uber/tango/core/targethasher" - pb "github.com/uber/tango/tangopb" ) func TestToShortRemote(t *testing.T) { @@ -55,93 +50,3 @@ func TestToShortRemote(t *testing.T) { }) } } - -func TestChunkTargets(t *testing.T) { - t.Parallel() - - // Create 25 targets, chunk by 10 → expect 3 chunks (10, 10, 5) - targets := make([]*pb.OptimizedTarget, 25) - for i := range targets { - targets[i] = &pb.OptimizedTarget{Id: int32(i)} - } - - responses := chunkTargets(targets, 10) - - require.Len(t, responses, 3) - - // Verify total count and order preserved - var total int - for _, resp := range responses { - item := resp.Item.(*pb.GetTargetGraphResponse_Targets) - for _, target := range item.Targets.Targets { - assert.Equal(t, int32(total), target.Id) - total++ - } - } - assert.Equal(t, 25, total) -} - -func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { - t.Parallel() - - numTargets := 50 - result := targethasher.Result{ - TargetNames: make([]string, numTargets), - Targets: make(map[string]*targethasher.Target, numTargets), - } - for i := 0; i < numTargets; i++ { - name := fmt.Sprintf("//pkg:target%d", i) - result.TargetNames[i] = name - result.Targets[name] = &targethasher.Target{Name: name, Hash: []byte{0}, RuleType: "go_library"} - } - - tests := []struct { - name string - targetChunkSize int - metadataMapChunkSize int - wantTargetChunks int - wantMetadataChunks int - }{ - { - name: "25 per chunk", - targetChunkSize: 25, - metadataMapChunkSize: 20, - wantTargetChunks: 2, - wantMetadataChunks: 3, - }, - { - name: "10 per chunk", - targetChunkSize: 10, - metadataMapChunkSize: 10, - wantTargetChunks: 5, - wantMetadataChunks: 5, - }, - { - name: "all in one chunk", - targetChunkSize: 100, - metadataMapChunkSize: 5_000, - wantTargetChunks: 1, - wantMetadataChunks: 1, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - responses, err := ResultToGetTargetGraphResponse(context.Background(), result, tt.targetChunkSize, tt.metadataMapChunkSize) - require.NoError(t, err) - - var targetChunks, metadataChunks int - for _, resp := range responses { - switch resp.Item.(type) { - case *pb.GetTargetGraphResponse_Targets: - targetChunks++ - case *pb.GetTargetGraphResponse_Metadata: - metadataChunks++ - } - } - assert.Equal(t, tt.wantTargetChunks, targetChunks) - assert.Equal(t, tt.wantMetadataChunks, metadataChunks) - }) - } -} diff --git a/example/cmd/query-bench/BUILD.bazel b/example/cmd/query-bench/BUILD.bazel index 6c798b8e..07ac5907 100644 --- a/example/cmd/query-bench/BUILD.bazel +++ b/example/cmd/query-bench/BUILD.bazel @@ -6,10 +6,9 @@ go_library( importpath = "github.com/uber/tango/example/cmd/query-bench", visibility = ["//visibility:private"], deps = [ - "//config", "//core/bazel", - "//core/common", "//core/targethasher", + "//internal/mapper", "@com_github_gogo_protobuf//jsonpb", "@org_uber_go_zap//:zap", ], diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 0885a617..37aa7cb9 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -31,10 +31,9 @@ import ( "time" "github.com/gogo/protobuf/jsonpb" - "github.com/uber/tango/config" "github.com/uber/tango/core/bazel" - "github.com/uber/tango/core/common" "github.com/uber/tango/core/targethasher" + "github.com/uber/tango/internal/mapper" "go.uber.org/zap" ) @@ -113,10 +112,9 @@ func run() error { totalDuration += elapsed fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) start = time.Now() - targetChunkSize, _, metadataMapChunkSize := config.ChunkSizesForByteBudget(config.DefaultMaxMessageBytes) - response, err := common.ResultToGetTargetGraphResponse( + response, err := mapper.ResultToGetTargetGraphResponse( ctx, targethasherResult, - targetChunkSize, metadataMapChunkSize, + 4_250_000, ) if err != nil { return fmt.Errorf("run %d: converting to GetTargetGraphResponse: %w", i+1, err) diff --git a/orchestrator/BUILD.bazel b/orchestrator/BUILD.bazel index fe456af3..065fc0d0 100644 --- a/orchestrator/BUILD.bazel +++ b/orchestrator/BUILD.bazel @@ -20,6 +20,7 @@ go_library( "//entity", "//graphrunner", "//internal/cachekey", + "//internal/mapper", "@com_github_uber_go_tally//:tally", "@org_uber_go_zap//:zap", ], diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index 28628521..a6448449 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -34,6 +34,7 @@ import ( "github.com/uber/tango/entity" "github.com/uber/tango/graphrunner" "github.com/uber/tango/internal/cachekey" + "github.com/uber/tango/internal/mapper" "go.uber.org/zap" ) @@ -208,8 +209,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - targetChunkSize, _, metadataMapChunkSize := config.ChunkSizesForByteBudget(b.config.Service.MaxMessageBytes) - responses, err := common.ResultToGetTargetGraphResponse(ctx, result, targetChunkSize, metadataMapChunkSize) + responses, err := mapper.ResultToGetTargetGraphResponse(ctx, result, b.config.Service.MaxMessageBytes) if err != nil { return nil, fmt.Errorf("convert target graph to response: %w", err) } From d3f0dee6169a0efe9311c2b82b3c2804ee9ebeff Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Tue, 14 Jul 2026 21:32:05 -0700 Subject: [PATCH 3/8] Update --- controller/controller.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/controller/controller.go b/controller/controller.go index ad4caa20..b2417d00 100644 --- a/controller/controller.go +++ b/controller/controller.go @@ -63,6 +63,12 @@ func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { if scope == nil { scope = tally.NoopScope } + // TODO: MaxMessageBytes should come from *config.Config via fx, not as a + // standalone Params field. Wire it through the fx module once config is + // provided as an fx dependency. + if p.MaxMessageBytes <= 0 { + p.MaxMessageBytes = 4_250_000 + } return &controller{ logger: p.Logger, storage: p.Storage, From cbdb27af42c7bbf25e9771a257cd077b96a6647d Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Wed, 15 Jul 2026 10:35:57 -0700 Subject: [PATCH 4/8] fix(mapper): guarantee at least one chunk and validate non-empty splits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BySize now always returns at least one chunk (empty input yields a single empty chunk) so callers always have a message to send on the stream. Non-first chunks that are empty return an error since that indicates a bug in the splitting logic. Same contract for ChunkMetadata. Removes the per-type wrapper functions (chunkTargets, ChunkChangedTargets) — callers inline the one-liner proto wrapping themselves, keeping BySize and ChunkMetadata as the two public APIs. Co-Authored-By: Claude Sonnet 5 --- controller/getchangedtargets.go | 18 ++++++---- controller/getchangedtargets_test.go | 50 +++++++++++++++++++++++----- 2 files changed, 54 insertions(+), 14 deletions(-) diff --git a/controller/getchangedtargets.go b/controller/getchangedtargets.go index 928533cc..64b9ecda 100644 --- a/controller/getchangedtargets.go +++ b/controller/getchangedtargets.go @@ -447,23 +447,29 @@ func (c *controller) compareTargetGraphs(ctx context.Context, scope tally.Scope, // Emit changes in chunks to stay within gRPC per-message size limits, followed by chunked metadata. var results []*pb.GetChangedTargetsResponse - for _, chunk := range mapper.BySize(changed, c.maxMessageBytes) { + changedChunks, err := mapper.BySize(changed, c.maxMessageBytes) + if err != nil { + return nil, err + } + for _, chunk := range changedChunks { results = append(results, &pb.GetChangedTargetsResponse{ Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{ - ChangedTargets: chunk, - }, + ChangedTargets: &pb.ChangedTargets{ChangedTargets: chunk}, }, }) } - for _, meta := range mapper.ChunkMetadata( + metaChunks, err := mapper.ChunkMetadata( mappers.target.Invert(), mappers.ruleType.Invert(), mappers.tag.Invert(), mappers.attrName.Invert(), mappers.attrVal.Invert(), c.maxMessageBytes, - ) { + ) + if err != nil { + return nil, err + } + for _, meta := range metaChunks { results = append(results, &pb.GetChangedTargetsResponse{ Item: &pb.GetChangedTargetsResponse_Metadata{ Metadata: meta, diff --git a/controller/getchangedtargets_test.go b/controller/getchangedtargets_test.go index 1c363e1f..2bb1b2df 100644 --- a/controller/getchangedtargets_test.go +++ b/controller/getchangedtargets_test.go @@ -296,18 +296,52 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { stream.EXPECT().Send(gomock.Any()).Return(errors.New("send error")) storagemock := storagemock.NewMockStorage(ctrl) - var buf bytes.Buffer - gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, + // Build two graphs with differing hashes so compareTargetGraphs produces a change. + var buf1 bytes.Buffer + w1 := gogio.NewDelimitedWriter(&buf1) + w1.WriteMsg(&pb.GetTargetGraphResponse{ + Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{ + Targets: []*pb.OptimizedTarget{{Id: 1, Hash: "old"}}, + }}, }) + w1.WriteMsg(&pb.GetTargetGraphResponse{ + Item: &pb.GetTargetGraphResponse_Metadata{Metadata: &pb.Metadata{ + TargetIdMapping: map[int32]string{1: "//pkg:t"}, + RuleTypeMapping: map[int32]string{}, + }}, + }) + graph1Bytes := buf1.Bytes() + + var buf2 bytes.Buffer + w2 := gogio.NewDelimitedWriter(&buf2) + w2.WriteMsg(&pb.GetTargetGraphResponse{ + Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{ + Targets: []*pb.OptimizedTarget{{Id: 1, Hash: "new"}}, + }}, + }) + w2.WriteMsg(&pb.GetTargetGraphResponse{ + Item: &pb.GetTargetGraphResponse_Metadata{Metadata: &pb.Metadata{ + TargetIdMapping: map[int32]string{1: "//pkg:t"}, + RuleTypeMapping: map[int32]string{}, + }}, + }) + graph2Bytes := buf2.Bytes() + storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req storage.DownloadRequest) (storage.DownloadResponse, error) { - if strings.Contains(req.Key, "compared-targets") { + switch { + case strings.Contains(req.Key, "compared-targets"): return storage.DownloadResponse{}, storage.NewNotFoundError(req.Key) + case strings.Contains(req.Key, "sha1"): + return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader([]byte("treehash1")))}, nil + case strings.Contains(req.Key, "sha2"): + return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader([]byte("treehash2")))}, nil + case strings.Contains(req.Key, "treehash1"): + return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader(graph1Bytes))}, nil + case strings.Contains(req.Key, "treehash2"): + return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader(graph2Bytes))}, nil + default: + return storage.DownloadResponse{}, fmt.Errorf("unexpected key: %s", req.Key) } - if strings.Contains(req.Key, "th") { - return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader(buf.Bytes()))}, nil - } - return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader([]byte("th")))}, nil }).AnyTimes() // Put is launched in a goroutine — use a channel to wait for it before the test ends. From 41d3271583982c0bebe76cfa402242c69ed8e3ac Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Wed, 15 Jul 2026 22:00:36 -0700 Subject: [PATCH 5/8] refactor(orchestrator): remove proto from target graph write path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add mapper.WriteTargetGraph which accepts entity.TargetGraph and handles proto conversion + storage write internally. The orchestrator now only touches entity types when writing computed graphs — proto conversion is confined to the mapper layer. Co-Authored-By: Claude Sonnet 5 --- internal/mapper/BUILD.bazel | 5 ++++- internal/mapper/graphstore.go | 19 +++++++++++++++++++ orchestrator/native_orchestrator.go | 6 +++--- 3 files changed, 26 insertions(+), 4 deletions(-) create mode 100644 internal/mapper/graphstore.go diff --git a/internal/mapper/BUILD.bazel b/internal/mapper/BUILD.bazel index 4a3ec64a..8e506749 100644 --- a/internal/mapper/BUILD.bazel +++ b/internal/mapper/BUILD.bazel @@ -4,15 +4,17 @@ go_library( name = "mapper", srcs = [ "build_description.go", + "chunk.go", + "graphstore.go", "target_graph.go", ], importpath = "github.com/uber/tango/internal/mapper", visibility = ["//visibility:public"], deps = [ + "//core/storage", "//core/targethasher", "//entity", "//internal/mapper/idmapper", - "//internal/streaming", "//tangopb", "@com_github_bazelbuild_buildtools//build_proto", ], @@ -22,6 +24,7 @@ go_test( name = "mapper_test", srcs = [ "build_description_test.go", + "chunk_test.go", "target_graph_test.go", ], embed = [":mapper"], diff --git a/internal/mapper/graphstore.go b/internal/mapper/graphstore.go new file mode 100644 index 00000000..ac6579c5 --- /dev/null +++ b/internal/mapper/graphstore.go @@ -0,0 +1,19 @@ +package mapper + +import ( + "context" + + "github.com/uber/tango/core/storage" + "github.com/uber/tango/entity" +) + +// WriteTargetGraph converts an entity.TargetGraph to chunked proto +// responses and writes them to storage. This keeps proto conversion +// out of the orchestrator — callers only work with entity types. +func WriteTargetGraph(ctx context.Context, st storage.Storage, key string, graph entity.TargetGraph, maxMessageBytes int) error { + responses, err := TargetGraphToProto(ctx, graph, maxMessageBytes) + if err != nil { + return err + } + return storage.WriteGraphStream(ctx, st, key, responses) +} diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index a6448449..dad62382 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -209,11 +209,11 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - responses, err := mapper.ResultToGetTargetGraphResponse(ctx, result, b.config.Service.MaxMessageBytes) + graph, err := mapper.ResultToTargetGraph(ctx, result) if err != nil { - return nil, fmt.Errorf("convert target graph to response: %w", err) + return nil, fmt.Errorf("convert target graph: %w", err) } - err = storage.WriteGraphStream(ctx, b.storage, treehashPath, responses) + err = mapper.WriteTargetGraph(ctx, b.storage, treehashPath, graph, b.config.Service.MaxMessageBytes) if err != nil { return nil, fmt.Errorf("write graph to storage at %s: %w", treehashPath, err) } From dd4e0a3bb5ec8f358f343a0db8fb1f58d78819ad Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 16 Jul 2026 13:28:41 -0700 Subject: [PATCH 6/8] refactor(example): use two-step ResultToTargetGraph + TargetGraphToProto Update query-bench to use the explicit two-step flow now that the ResultToGetTargetGraphResponse convenience function has been removed. Co-Authored-By: Claude Sonnet 5 --- example/cmd/query-bench/main.go | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 37aa7cb9..54d6c56c 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -112,15 +112,16 @@ func run() error { totalDuration += elapsed fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) start = time.Now() - response, err := mapper.ResultToGetTargetGraphResponse( - ctx, targethasherResult, - 4_250_000, - ) + graph, err := mapper.ResultToTargetGraph(ctx, targethasherResult) if err != nil { - return fmt.Errorf("run %d: converting to GetTargetGraphResponse: %w", i+1, err) + return fmt.Errorf("run %d: converting to target graph: %w", i+1, err) + } + response, err := mapper.TargetGraphToProto(ctx, graph, 4_250_000) + if err != nil { + return fmt.Errorf("run %d: converting to proto: %w", i+1, err) } elapsed = time.Since(start) - fmt.Printf("run %d: ResultToGetTargetGraphResponse: %v (%d responses)\n", i+1, elapsed.Round(time.Millisecond), len(response)) + fmt.Printf("run %d: TargetGraphToProto: %v (%d responses)\n", i+1, elapsed.Round(time.Millisecond), len(response)) m := jsonpb.Marshaler{Indent: " "} for _, r := range response { if err := m.Marshal(os.Stdout, r); err != nil { From f21c929bb01bc6fb68c22d30cdde0074709f46c6 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 16 Jul 2026 13:49:10 -0700 Subject: [PATCH 7/8] refactor: use entity-based GraphChunkReader throughout orchestrator and controller MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace storage.GraphReader (proto-typed) with mapper.GraphChunkReader (entity-typed) in the Orchestrator interface, native orchestrator, and controller. Proto conversion now happens at stream-send time in the controller via GraphChunkToProto, keeping the orchestrator proto-free. Add mapper.NewGraphChunkReader (wraps storage.NewGraphReader, converts proto → entity on read) and mapper.ReadAllGraphChunks for batch reads. Co-Authored-By: Claude Sonnet 5 --- controller/BUILD.bazel | 1 + controller/getchangedtargets.go | 4 +- controller/getchangedtargets_test.go | 30 +++-------- controller/gettargetgraph.go | 9 ++-- controller/gettargetgraph_test.go | 28 ++++------ controller/testhelper_test.go | 29 ++++++++++ internal/mapper/graphstore.go | 53 +++++++++++++++++++ orchestrator/native_orchestrator.go | 6 +-- orchestrator/native_orchestrator_test.go | 8 ++- orchestrator/orchestrator.go | 5 +- orchestrator/orchestratormock/BUILD.bazel | 2 +- .../orchestratormock/orchestratormock.go | 6 +-- 12 files changed, 121 insertions(+), 60 deletions(-) diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 534493d5..5f5c9f6f 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -45,6 +45,7 @@ go_test( "//core/storage", "//core/storage/storagemock", "//entity", + "//internal/mapper", "//orchestrator/orchestratormock", "//tangopb", "//tangopb/tangopbmock", diff --git a/controller/getchangedtargets.go b/controller/getchangedtargets.go index 64b9ecda..75664abe 100644 --- a/controller/getchangedtargets.go +++ b/controller/getchangedtargets.go @@ -262,7 +262,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, l } defer graphReader.Close() - // Read all chunks from the stream + // Read all chunks from the stream and convert to proto var chunks []*pb.GetTargetGraphResponse for { chunk, err := graphReader.Read() @@ -274,7 +274,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, l results <- graphResult{order: idx, err: err} return } - chunks = append(chunks, chunk) + chunks = append(chunks, mapper.GraphChunkToProto(&chunk)) } }(i) } diff --git a/controller/getchangedtargets_test.go b/controller/getchangedtargets_test.go index 2bb1b2df..18692c44 100644 --- a/controller/getchangedtargets_test.go +++ b/controller/getchangedtargets_test.go @@ -31,6 +31,7 @@ import ( "github.com/uber/tango/core/storage" storagemock "github.com/uber/tango/core/storage/storagemock" "github.com/uber/tango/entity" + "github.com/uber/tango/internal/mapper" orchestratormock "github.com/uber/tango/orchestrator/orchestratormock" pb "github.com/uber/tango/tangopb" tangomock "github.com/uber/tango/tangopb/tangopbmock" @@ -1414,21 +1415,6 @@ func TestSendTrimmedChangedTargets_RetainsDeletedAtMaxDistanceOne(t *testing.T) assert.True(t, gotDeleted, "DELETED entry at distance 0 must survive max_distance=1") } -func newMockGraphReader(ctrl *gomock.Controller, chunks ...*pb.GetTargetGraphResponse) *storagemock.MockGraphReader { - r := storagemock.NewMockGraphReader(ctrl) - idx := 0 - r.EXPECT().Read().DoAndReturn(func() (*pb.GetTargetGraphResponse, error) { - if idx >= len(chunks) { - return nil, io.EOF - } - c := chunks[idx] - idx++ - return c, nil - }).AnyTimes() - r.EXPECT().Close().Return(nil).AnyTimes() - return r -} - func changedTargetsRequest() *pb.GetChangedTargetsRequest { return &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -1546,8 +1532,8 @@ func TestFetchTargetGraphs(t *testing.T) { ctrl := gomock.NewController(t) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, _ entity.GetTargetGraphRequest) (storage.GraphReader, error) { - return newMockGraphReader(ctrl, chunk), nil + func(_ context.Context, _ entity.GetTargetGraphRequest) (*mapper.GraphChunkReader, error) { + return newGraphChunkReaderFromProto(t, ctrl, chunk), nil }).Times(2) c := newTestController(zaptest.NewLogger(t)) @@ -1564,11 +1550,11 @@ func TestFetchTargetGraphs(t *testing.T) { injected := errors.New("orchestrator boom") orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, p entity.GetTargetGraphRequest) (storage.GraphReader, error) { + func(_ context.Context, p entity.GetTargetGraphRequest) (*mapper.GraphChunkReader, error) { if p.Build.BaseSha == "sha1" { return nil, injected } - return newMockGraphReader(ctrl, chunk), nil + return newGraphChunkReaderFromProto(t, ctrl, chunk), nil }).Times(2) c := newTestController(zaptest.NewLogger(t)) @@ -1585,8 +1571,8 @@ func TestFetchTargetGraphs(t *testing.T) { ctrl := gomock.NewController(t) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, _ entity.GetTargetGraphRequest) (storage.GraphReader, error) { - return newMockGraphReader(ctrl), nil + func(_ context.Context, _ entity.GetTargetGraphRequest) (*mapper.GraphChunkReader, error) { + return newGraphChunkReaderFromProto(t, ctrl), nil }).Times(2) c := newTestController(zaptest.NewLogger(t)) @@ -1600,7 +1586,7 @@ func TestFetchTargetGraphs(t *testing.T) { ctrl := gomock.NewController(t) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, _ entity.GetTargetGraphRequest) (storage.GraphReader, error) { + func(_ context.Context, _ entity.GetTargetGraphRequest) (*mapper.GraphChunkReader, error) { panic("boom in orchestrator") }).Times(2) diff --git a/controller/gettargetgraph.go b/controller/gettargetgraph.go index 630e16d0..451b76b4 100644 --- a/controller/gettargetgraph.go +++ b/controller/gettargetgraph.go @@ -67,7 +67,7 @@ func (c *controller) GetTargetGraph(request *pb.GetTargetGraphRequest, stream pb sendStart := time.Now() outputConfig := request.GetOutputConfig() for { - graphStreamChunk, err := graphReader.Read() + chunk, err := graphReader.Read() if err == io.EOF { sendDuration := time.Since(sendStart) totalDuration := time.Since(start) @@ -82,7 +82,8 @@ func (c *controller) GetTargetGraph(request *pb.GetTargetGraphRequest, stream pb if err != nil { return common.WithReason(failureReasonGraphFetch, common.ErrorTypeInfra, err) } - toSend := applyOptimizedTargetsOutputConfigToChunk(graphStreamChunk, outputConfig) + protoResp := mapper.GraphChunkToProto(&chunk) + toSend := applyOptimizedTargetsOutputConfigToChunk(protoResp, outputConfig) err = stream.Send(toSend) if err != nil { return common.WithReason(failureReasonSend, common.ErrorTypeInfra, fmt.Errorf("send graph: %w", err)) @@ -96,7 +97,7 @@ func (c *controller) GetTargetGraph(request *pb.GetTargetGraphRequest, stream pb // entries store the full payload and stripping happens at send time, so // letting an orchestrator see it could poison the shared cache with // stripped graphs. -func (c *controller) getGraph(ctx context.Context, req entity.GetTargetGraphRequest) (storage.GraphReader, error) { +func (c *controller) getGraph(ctx context.Context, req entity.GetTargetGraphRequest) (*mapper.GraphChunkReader, error) { start := time.Now() logger := c.logger.With( zap.Any("build_description", req.Build), @@ -123,7 +124,7 @@ func (c *controller) getGraph(ctx context.Context, req entity.GetTargetGraphRequ treehashPath := cachekey.GetGraphByTreeHash(req.Build.Remote, string(treehashBytes), req.Build.Strategy, req.ExcludeFilesRegex) // Download the target graph based on treehash. storageStart := time.Now() - graphReader, err := storage.NewGraphReader(ctx, c.storage, treehashPath) + graphReader, err := mapper.NewGraphChunkReader(ctx, c.storage, treehashPath) if err != nil { if ctx.Err() != nil { return nil, common.WithReason(common.FailureReasonCancelled, common.ErrorTypeUser, ctx.Err()) diff --git a/controller/gettargetgraph_test.go b/controller/gettargetgraph_test.go index 4a1ff02e..402f7197 100644 --- a/controller/gettargetgraph_test.go +++ b/controller/gettargetgraph_test.go @@ -191,17 +191,12 @@ func TestGetTargetGraph_TreehashNotFound_NoError(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, storage.NewNotFoundError("x")) orchestrator := orchestratormock.NewMockOrchestrator(ctrl) - // Provide a fake GraphReader that yields one message then EOF - graphReader := storagemock.NewMockGraphReader(ctrl) - graphReader.EXPECT().Read().DoAndReturn(func() (*pb.GetTargetGraphResponse, error) { - return &pb.GetTargetGraphResponse{ + reader := newGraphChunkReaderFromProto(t, ctrl, + &pb.GetTargetGraphResponse{ Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }, nil - }).Times(1) - // Controller may call Read again to observe EOF - graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) - graphReader.EXPECT().Close().Return(nil) - orchestrator.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) + }, + ) + orchestrator.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(reader, nil) c := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, @@ -314,13 +309,12 @@ func TestGetTargetGraph_GraphNotFound_FallsThrough(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, storage.NewNotFoundError("graphs/abc")), ) orch := orchestratormock.NewMockOrchestrator(ctrl) - graphReader := storagemock.NewMockGraphReader(ctrl) - graphReader.EXPECT().Read().Return(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }, nil).Times(1) - graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) - graphReader.EXPECT().Close().Return(nil) - orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) + reader := newGraphChunkReaderFromProto(t, ctrl, + &pb.GetTargetGraphResponse{ + Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, + }, + ) + orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(reader, nil) c := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go index 1c470630..0109fa7d 100644 --- a/controller/testhelper_test.go +++ b/controller/testhelper_test.go @@ -15,9 +15,19 @@ package controller import ( + "bytes" "context" + "io" + "testing" + gogio "github.com/gogo/protobuf/io" + "github.com/stretchr/testify/require" "github.com/uber-go/tally" + "github.com/uber/tango/core/storage" + storagemock "github.com/uber/tango/core/storage/storagemock" + "github.com/uber/tango/internal/mapper" + pb "github.com/uber/tango/tangopb" + "go.uber.org/mock/gomock" "go.uber.org/zap" ) @@ -30,3 +40,22 @@ func newTestController(logger *zap.Logger) *controller { appCtx: context.Background(), } } + +// newGraphChunkReaderFromProto creates a *mapper.GraphChunkReader backed by +// in-memory proto responses, for use in controller tests that mock the +// orchestrator return value. +func newGraphChunkReaderFromProto(t *testing.T, ctrl *gomock.Controller, responses ...*pb.GetTargetGraphResponse) *mapper.GraphChunkReader { + t.Helper() + var buf bytes.Buffer + w := gogio.NewDelimitedWriter(&buf) + for _, r := range responses { + require.NoError(t, w.WriteMsg(r)) + } + st := storagemock.NewMockStorage(ctrl) + st.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ + ReadCloser: io.NopCloser(bytes.NewReader(buf.Bytes())), + }, nil) + reader, err := mapper.NewGraphChunkReader(context.Background(), st, "test-key") + require.NoError(t, err) + return reader +} diff --git a/internal/mapper/graphstore.go b/internal/mapper/graphstore.go index ac6579c5..ebe0ecab 100644 --- a/internal/mapper/graphstore.go +++ b/internal/mapper/graphstore.go @@ -2,6 +2,7 @@ package mapper import ( "context" + "io" "github.com/uber/tango/core/storage" "github.com/uber/tango/entity" @@ -17,3 +18,55 @@ func WriteTargetGraph(ctx context.Context, st storage.Storage, key string, graph } return storage.WriteGraphStream(ctx, st, key, responses) } + +// GraphChunkReader reads entity.GraphChunk values from storage. It wraps +// storage.GraphReader and converts proto → entity on the fly. +type GraphChunkReader struct { + inner storage.GraphReader +} + +// NewGraphChunkReader opens a stored target graph at key and returns a +// reader that yields entity.GraphChunk values. Proto deserialization +// happens internally — callers see only entity types. +func NewGraphChunkReader(ctx context.Context, st storage.Storage, key string) (*GraphChunkReader, error) { + gr, err := storage.NewGraphReader(ctx, st, key) + if err != nil { + return nil, err + } + return &GraphChunkReader{inner: gr}, nil +} + +// Read returns the next GraphChunk. Returns io.EOF at the end of the stream. +func (r *GraphChunkReader) Read() (entity.GraphChunk, error) { + resp, err := r.inner.Read() + if err != nil { + return entity.GraphChunk{}, err + } + return ProtoToGraphChunk(resp), nil +} + +// Close releases the underlying reader. +func (r *GraphChunkReader) Close() error { + return r.inner.Close() +} + +// ReadAllGraphChunks reads all GraphChunks from storage at key into a slice. +func ReadAllGraphChunks(ctx context.Context, st storage.Storage, key string) ([]entity.GraphChunk, error) { + r, err := NewGraphChunkReader(ctx, st, key) + if err != nil { + return nil, err + } + defer r.Close() + + var chunks []entity.GraphChunk + for { + chunk, err := r.Read() + if err == io.EOF { + return chunks, nil + } + if err != nil { + return nil, err + } + chunks = append(chunks, chunk) + } +} diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index dad62382..d3c7abf0 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -99,7 +99,7 @@ func NewNativeOrchestrator(appCtx context.Context, p Params) (Orchestrator, erro // GetTargetGraph is used to compute the target graph locally. // It leases a workspace, checks out the base revision, applies the change requests, and computes the target graph. -func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetTargetGraphRequest) (_ storage.GraphReader, retErr error) { +func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetTargetGraphRequest) (_ *mapper.GraphChunkReader, retErr error) { scope := b.scope.SubScope("get_target_graph") scope.Counter("calls").Inc(1) defer func() { @@ -172,7 +172,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT } treehashPath := cachekey.GetGraphByTreeHash(build.Remote, treehash, build.Strategy, req.ExcludeFilesRegex) if !req.BypassCache { - graphReader, err := storage.NewGraphReader(ctx, b.storage, treehashPath) + graphReader, err := mapper.NewGraphChunkReader(ctx, b.storage, treehashPath) if err == nil { logger.Infow("GetTargetGraph: Cache hit on treehash", zap.String("treehash", treehash)) return graphReader, nil @@ -223,7 +223,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT if err != nil { return nil, fmt.Errorf("store treehash mapping at %s: %w", treehashCachePath, err) } - graphReader, err := storage.NewGraphReader(ctx, b.storage, treehashPath) + graphReader, err := mapper.NewGraphChunkReader(ctx, b.storage, treehashPath) if err != nil { return nil, fmt.Errorf("create graph reader at %s: %w", treehashPath, err) } diff --git a/orchestrator/native_orchestrator_test.go b/orchestrator/native_orchestrator_test.go index e3650e4d..8b0819bb 100644 --- a/orchestrator/native_orchestrator_test.go +++ b/orchestrator/native_orchestrator_test.go @@ -79,13 +79,11 @@ func TestNative_GetTargetGraph_Success(t *testing.T) { require.NoError(t, err) require.NotNil(t, reader) defer reader.Close() - graph, rerr := reader.Read() + chunk, rerr := reader.Read() require.NoError(t, rerr) - require.NotNil(t, graph) - assert.NotNil(t, graph.GetTargets()) - graph, rerr = reader.Read() + assert.NotNil(t, chunk.Targets) + _, rerr = reader.Read() assert.Equal(t, io.EOF, rerr) - assert.Nil(t, graph) } func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) { diff --git a/orchestrator/orchestrator.go b/orchestrator/orchestrator.go index ace1b322..80628dda 100644 --- a/orchestrator/orchestrator.go +++ b/orchestrator/orchestrator.go @@ -17,12 +17,11 @@ package orchestrator import ( "context" - "github.com/uber/tango/core/storage" - "github.com/uber/tango/entity" + "github.com/uber/tango/internal/mapper" ) // Orchestrator defines high-level execution interface that "does everything" type Orchestrator interface { - GetTargetGraph(ctx context.Context, req entity.GetTargetGraphRequest) (storage.GraphReader, error) + GetTargetGraph(ctx context.Context, req entity.GetTargetGraphRequest) (*mapper.GraphChunkReader, error) } diff --git a/orchestrator/orchestratormock/BUILD.bazel b/orchestrator/orchestratormock/BUILD.bazel index 50e72aae..daf60af4 100644 --- a/orchestrator/orchestratormock/BUILD.bazel +++ b/orchestrator/orchestratormock/BUILD.bazel @@ -6,8 +6,8 @@ go_library( importpath = "github.com/uber/tango/orchestrator/orchestratormock", visibility = ["//visibility:public"], deps = [ - "//core/storage", "//entity", + "//internal/mapper", "@org_uber_go_mock//gomock", ], ) diff --git a/orchestrator/orchestratormock/orchestratormock.go b/orchestrator/orchestratormock/orchestratormock.go index 5955502f..ba1d818a 100644 --- a/orchestrator/orchestratormock/orchestratormock.go +++ b/orchestrator/orchestratormock/orchestratormock.go @@ -13,8 +13,8 @@ import ( context "context" reflect "reflect" - storage "github.com/uber/tango/core/storage" entity "github.com/uber/tango/entity" + mapper "github.com/uber/tango/internal/mapper" gomock "go.uber.org/mock/gomock" ) @@ -43,10 +43,10 @@ func (m *MockOrchestrator) EXPECT() *MockOrchestratorMockRecorder { } // GetTargetGraph mocks base method. -func (m *MockOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetTargetGraphRequest) (storage.GraphReader, error) { +func (m *MockOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetTargetGraphRequest) (*mapper.GraphChunkReader, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "GetTargetGraph", ctx, req) - ret0, _ := ret[0].(storage.GraphReader) + ret0, _ := ret[0].(*mapper.GraphChunkReader) ret1, _ := ret[1].(error) return ret0, ret1 } From e1d5cf608672ceb23d21988e82148ca9673e6410 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 16 Jul 2026 15:12:28 -0700 Subject: [PATCH 8/8] refactor(controller): use streaming.SplitBySize and SplitMetadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Update getchangedtargets to use the new streaming package for splitting changed targets and metadata, replacing the removed mapper.BySize and mapper.ChunkMetadata. Export EntityMetadataToProto for the entity → proto metadata conversion at the controller boundary. Co-Authored-By: Claude Sonnet 5 --- controller/BUILD.bazel | 1 + controller/getchangedtargets.go | 13 +++++++------ internal/mapper/BUILD.bazel | 3 +-- internal/mapper/target_graph.go | 4 ++-- 4 files changed, 11 insertions(+), 10 deletions(-) diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 5f5c9f6f..2a002930 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -20,6 +20,7 @@ go_library( "//internal/cachekey", "//internal/mapper", "//internal/mapper/idmapper", + "//internal/streaming", "//internal/targetdiff", "//orchestrator", "//tangopb", diff --git a/controller/getchangedtargets.go b/controller/getchangedtargets.go index 75664abe..e8581f17 100644 --- a/controller/getchangedtargets.go +++ b/controller/getchangedtargets.go @@ -28,6 +28,7 @@ import ( "github.com/uber/tango/internal/cachekey" "github.com/uber/tango/internal/mapper" "github.com/uber/tango/internal/mapper/idmapper" + "github.com/uber/tango/internal/streaming" "github.com/uber/tango/internal/targetdiff" pb "github.com/uber/tango/tangopb" "go.uber.org/zap" @@ -447,18 +448,18 @@ func (c *controller) compareTargetGraphs(ctx context.Context, scope tally.Scope, // Emit changes in chunks to stay within gRPC per-message size limits, followed by chunked metadata. var results []*pb.GetChangedTargetsResponse - changedChunks, err := mapper.BySize(changed, c.maxMessageBytes) + changedGroups, err := streaming.SplitBySize(changed, c.maxMessageBytes) if err != nil { return nil, err } - for _, chunk := range changedChunks { + for _, group := range changedGroups { results = append(results, &pb.GetChangedTargetsResponse{ Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{ChangedTargets: chunk}, + ChangedTargets: &pb.ChangedTargets{ChangedTargets: group}, }, }) } - metaChunks, err := mapper.ChunkMetadata( + metaGroups, err := streaming.SplitMetadata( mappers.target.Invert(), mappers.ruleType.Invert(), mappers.tag.Invert(), @@ -469,10 +470,10 @@ func (c *controller) compareTargetGraphs(ctx context.Context, scope tally.Scope, if err != nil { return nil, err } - for _, meta := range metaChunks { + for _, meta := range metaGroups { results = append(results, &pb.GetChangedTargetsResponse{ Item: &pb.GetChangedTargetsResponse_Metadata{ - Metadata: meta, + Metadata: mapper.EntityMetadataToProto(meta), }, }) } diff --git a/internal/mapper/BUILD.bazel b/internal/mapper/BUILD.bazel index 8e506749..12ab69d3 100644 --- a/internal/mapper/BUILD.bazel +++ b/internal/mapper/BUILD.bazel @@ -4,7 +4,6 @@ go_library( name = "mapper", srcs = [ "build_description.go", - "chunk.go", "graphstore.go", "target_graph.go", ], @@ -15,6 +14,7 @@ go_library( "//core/targethasher", "//entity", "//internal/mapper/idmapper", + "//internal/streaming", "//tangopb", "@com_github_bazelbuild_buildtools//build_proto", ], @@ -24,7 +24,6 @@ go_test( name = "mapper_test", srcs = [ "build_description_test.go", - "chunk_test.go", "target_graph_test.go", ], embed = [":mapper"], diff --git a/internal/mapper/target_graph.go b/internal/mapper/target_graph.go index 8ca4d459..5912fe8c 100644 --- a/internal/mapper/target_graph.go +++ b/internal/mapper/target_graph.go @@ -199,7 +199,7 @@ func GraphChunkToProto(chunk *entity.GraphChunk) *tangopb.GetTargetGraphResponse if chunk.Metadata != nil { return &tangopb.GetTargetGraphResponse{ Item: &tangopb.GetTargetGraphResponse_Metadata{ - Metadata: entityMetadataToProto(chunk.Metadata), + Metadata: EntityMetadataToProto(chunk.Metadata), }, } } @@ -257,7 +257,7 @@ func protoTargetToIDTarget(t *tangopb.OptimizedTarget) entity.IDTarget { } } -func entityMetadataToProto(m *entity.GraphMetadata) *tangopb.Metadata { +func EntityMetadataToProto(m *entity.GraphMetadata) *tangopb.Metadata { return &tangopb.Metadata{ TargetIdMapping: m.TargetIDMapping, RuleTypeMapping: m.RuleTypeMapping,