diff --git a/config/BUILD.bazel b/config/BUILD.bazel index 0c95044..84c243f 100644 --- a/config/BUILD.bazel +++ b/config/BUILD.bazel @@ -1,4 +1,4 @@ -load("@rules_go//go:def.bzl", "go_library") +load("@rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "config", @@ -12,3 +12,10 @@ go_library( visibility = ["//visibility:public"], deps = ["@com_github_goccy_go_yaml//:go-yaml"], ) + +go_test( + name = "config_test", + srcs = ["config_test.go"], + embed = [":config"], + deps = ["@com_github_stretchr_testify//require"], +) diff --git a/config/config.go b/config/config.go index 77ca493..41decb7 100644 --- a/config/config.go +++ b/config/config.go @@ -70,12 +70,24 @@ func Parse(configFilePath string) (*Config, error) { if config.Service.WorkerPoolSize <= 0 { return nil, fmt.Errorf("service.worker_pool_size must be > 0, got %d", config.Service.WorkerPoolSize) } + if config.Service.Chunking.TargetChunkSize <= 0 { + return nil, fmt.Errorf("service.chunking.target_chunk_size must be > 0, got %d", config.Service.Chunking.TargetChunkSize) + } + if config.Service.Chunking.ChangedTargetChunkSize <= 0 { + return nil, fmt.Errorf("service.chunking.changed_target_chunk_size must be > 0, got %d", config.Service.Chunking.ChangedTargetChunkSize) + } + if config.Service.Chunking.MetadataMapChunkSize <= 0 { + return nil, fmt.Errorf("service.chunking.metadata_map_chunk_size must be > 0, got %d", config.Service.Chunking.MetadataMapChunkSize) + } config.repositoryByRemote = make(map[string]*RepositoryConfig, len(config.Repository)) for i := range config.Repository { remote := config.Repository[i].Remote if remote == "" { return nil, fmt.Errorf("repository[%d].remote must not be empty", i) } + if config.Repository[i].QueryTimeoutSeconds <= 0 { + return nil, fmt.Errorf("repository[%d].query_timeout_seconds must be > 0", i) + } if _, exists := config.repositoryByRemote[remote]; exists { return nil, fmt.Errorf("duplicate repository remote %q", remote) } diff --git a/config/config_test.go b/config/config_test.go new file mode 100644 index 0000000..eff9be6 --- /dev/null +++ b/config/config_test.go @@ -0,0 +1,186 @@ +// 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 ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestParse_Validation(t *testing.T) { + t.Parallel() + + validConfig := ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +storage: + type: memory +service: + worker_pool_size: 5 + repo_manager_clone_path: /tmp/tango + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +` + + tests := []struct { + name string + yaml string + wantErr bool + }{ + { + name: "valid config", + yaml: validConfig, + }, + { + name: "worker_pool_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "target_chunk_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "changed_target_chunk_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "metadata_map_chunk_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 +`, + wantErr: true, + }, + { + name: "repository remote not set", + yaml: ` +repository: + - query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "query_timeout_seconds not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "duplicate repository remote", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "worker_root_path without repo_manager_clone_path", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + worker_root_path: /tmp/workers + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "config.yaml") + require.NoError(t, os.WriteFile(path, []byte(tt.yaml), 0o644)) + + cfg, err := Parse(path) + if tt.wantErr { + require.Error(t, err, "expected error when %s", tt.name) + } else { + require.NoError(t, err) + require.NotNil(t, cfg) + } + }) + } +} diff --git a/config/repository_config.go b/config/repository_config.go index 4052ec4..8e23f0a 100644 --- a/config/repository_config.go +++ b/config/repository_config.go @@ -22,7 +22,7 @@ type RepositoryConfig struct { ExcludeExternalTargets bool `yaml:"exclude_external_targets"` BzlmodEnabled bool `yaml:"bzlmod_enabled"` BazelCommand string `yaml:"bazel_command"` - QueryTimeout int64 `yaml:"query_timeout"` // in seconds + QueryTimeoutSeconds int64 `yaml:"query_timeout_seconds"` BazelExtraArgs []string `yaml:"bazel_extra_args"` StreamBazelLogs bool `yaml:"stream_bazel_logs"` } diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 581ba3b..2194182 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -33,10 +33,10 @@ go_test( "getchangedtargets_test.go", "gettargetgraph_test.go", "output_filter_test.go", - "testhelper_test.go", ], embed = [":controller"], deps = [ + "//config", "//core/common", "//core/storage", "//core/storage/storagemock", @@ -46,7 +46,6 @@ go_test( "@com_github_gogo_protobuf//io", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", - "@com_github_uber_go_tally//:tally", "@org_uber_go_mock//gomock", "@org_uber_go_zap//:zap", "@org_uber_go_zap//zaptest", diff --git a/controller/controller.go b/controller/controller.go index 34ca33c..92089c9 100644 --- a/controller/controller.go +++ b/controller/controller.go @@ -16,11 +16,11 @@ package controller import ( "context" + "fmt" "time" "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" @@ -59,34 +59,31 @@ type controller struct { // NewController creates a new controller. appCtx is cancelled on process // shutdown to abort background work. -func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { +func NewController(appCtx context.Context, p Params) (pb.TangoYARPCServer, error) { scope := p.Scope if scope == nil { scope = tally.NoopScope } - targetChunkSize := p.ChunkConfig.TargetChunkSize - if targetChunkSize <= 0 { - targetChunkSize = common.DefaultTargetChunkSize + if p.ChunkConfig.TargetChunkSize <= 0 { + return nil, fmt.Errorf("target_chunk_size must be > 0, got %d", p.ChunkConfig.TargetChunkSize) } - changedTargetChunkSize := p.ChunkConfig.ChangedTargetChunkSize - if changedTargetChunkSize <= 0 { - changedTargetChunkSize = common.DefaultChangedTargetChunkSize + if p.ChunkConfig.ChangedTargetChunkSize <= 0 { + return nil, fmt.Errorf("changed_target_chunk_size must be > 0, got %d", p.ChunkConfig.ChangedTargetChunkSize) } - metadataMapChunkSize := p.ChunkConfig.MetadataMapChunkSize - if metadataMapChunkSize <= 0 { - metadataMapChunkSize = common.DefaultMetadataMapChunkSize + if p.ChunkConfig.MetadataMapChunkSize <= 0 { + return nil, fmt.Errorf("metadata_map_chunk_size must be > 0, got %d", p.ChunkConfig.MetadataMapChunkSize) } return &controller{ logger: p.Logger, storage: p.Storage, orchestrator: p.Orchestrator, scope: scope.SubScope("controller"), - targetChunkSize: targetChunkSize, - changedTargetChunkSize: changedTargetChunkSize, - metadataMapChunkSize: metadataMapChunkSize, + targetChunkSize: p.ChunkConfig.TargetChunkSize, + changedTargetChunkSize: p.ChunkConfig.ChangedTargetChunkSize, + metadataMapChunkSize: p.ChunkConfig.MetadataMapChunkSize, totalDurationBuckets: _totalDurationBuckets, appCtx: appCtx, - } + }, nil } // linkRequestCtx returns a context derived from reqCtx that is also cancelled diff --git a/controller/controller_test.go b/controller/controller_test.go index b891f87..c591866 100644 --- a/controller/controller_test.go +++ b/controller/controller_test.go @@ -19,11 +19,26 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/tango/config" orchestratormock "github.com/uber/tango/orchestrator/orchestratormock" "go.uber.org/mock/gomock" "go.uber.org/zap" ) +var _testChunkConfig = config.ChunkConfig{ + TargetChunkSize: 250, + ChangedTargetChunkSize: 125, + MetadataMapChunkSize: 50_000, +} + +func newTestController(t *testing.T, appCtx context.Context, p Params) *controller { + t.Helper() + ctrl, err := NewController(appCtx, p) + require.NoError(t, err) + return ctrl.(*controller) +} + // TestNewController_StoresAppContext verifies the caller-supplied context is // retained and is the one observed by background goroutines. func TestNewController_StoresAppContext(t *testing.T) { @@ -31,10 +46,11 @@ func TestNewController_StoresAppContext(t *testing.T) { appCtx, cancel := context.WithCancel(context.Background()) defer cancel() - c := NewController(appCtx, Params{ + c := newTestController(t, appCtx, Params{ Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), - }).(*controller) + ChunkConfig: _testChunkConfig, + }) assert.Same(t, appCtx, c.appCtx) assert.NoError(t, c.appCtx.Err()) @@ -101,3 +117,54 @@ func TestLinkRequestCtx_CancelReleasesAfterFunc(t *testing.T) { cancelApp() assert.Equal(t, firstErr, linked.Err(), "linked.Err() must not change after cancelLink") } + +func TestNewController_Validation(t *testing.T) { + t.Parallel() + + valid := config.ChunkConfig{ + TargetChunkSize: 250, + ChangedTargetChunkSize: 125, + MetadataMapChunkSize: 50_000, + } + + tests := []struct { + name string + chunk config.ChunkConfig + wantErr bool + }{ + { + name: "valid", + chunk: valid, + }, + { + name: "target_chunk_size not set", + chunk: config.ChunkConfig{ChangedTargetChunkSize: 125, MetadataMapChunkSize: 50_000}, + wantErr: true, + }, + { + name: "changed_target_chunk_size not set", + chunk: config.ChunkConfig{TargetChunkSize: 250, MetadataMapChunkSize: 50_000}, + wantErr: true, + }, + { + name: "metadata_map_chunk_size not set", + chunk: config.ChunkConfig{TargetChunkSize: 250, ChangedTargetChunkSize: 125}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, err := NewController(context.Background(), Params{ + Logger: zap.NewNop(), + ChunkConfig: tt.chunk, + }) + if tt.wantErr { + require.Error(t, err, "expected error when %s", tt.name) + } else { + require.NoError(t, err) + } + }) + } +} diff --git a/controller/getchangedtargets_test.go b/controller/getchangedtargets_test.go index 1e854fc..02dda25 100644 --- a/controller/getchangedtargets_test.go +++ b/controller/getchangedtargets_test.go @@ -133,7 +133,7 @@ func TestValidateGetChangedTargetsRequest(t *testing.T) { } func TestCompareTargetGraphs(t *testing.T) { - c := newTestController(zap.NewNop()) + c := newTestController(t, context.Background(), Params{Logger: zap.NewNop(), ChunkConfig: _testChunkConfig}) firstGraph := &pb.GetTargetGraphResponse{ Item: &pb.GetTargetGraphResponse_Metadata{ @@ -155,7 +155,11 @@ func TestGetChangedTargets_ValidationError(t *testing.T) { ctrl := gomock.NewController(t) stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) - c := NewController(context.Background(), Params{Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl)}) + c := newTestController(t, context.Background(), Params{ + Logger: zap.NewNop(), + Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, + }) err := c.GetChangedTargets(nil, stream) assert.EqualError(t, err, "request cannot be nil") @@ -196,10 +200,11 @@ func TestGetChangedTargets_CacheHit(t *testing.T) { stream.EXPECT().Send(gomock.Any()).Return(nil).Times(2) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) request := &pb.GetChangedTargetsRequest{ @@ -227,10 +232,11 @@ func TestGetChangedTargets_TreehashReadError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{}, injected).Times(2) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zap.NewNop(), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) request := &pb.GetChangedTargetsRequest{ @@ -316,10 +322,11 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { return nil }) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) request := &pb.GetChangedTargetsRequest{ @@ -423,10 +430,11 @@ func TestGetChangedTargets_streamChunks(t *testing.T) { return nil }) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) request := &pb.GetChangedTargetsRequest{ @@ -523,10 +531,11 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { appCtx, cancelApp := context.WithCancel(context.Background()) defer cancelApp() - c := NewController(appCtx, Params{ + c := newTestController(t, appCtx, Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) request := &pb.GetChangedTargetsRequest{ @@ -576,7 +585,7 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { } func TestCompareTargetGraphs_NewTarget_CanonicalIDs(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) first := []*pb.GetTargetGraphResponse{ { @@ -629,7 +638,7 @@ func TestCompareTargetGraphs_NewTarget_CanonicalIDs(t *testing.T) { } func TestCompareTargetGraphs_SourceFileDirectAndPropagation(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: source file A (id 1, hash h1), lib L (id 2, hash h1, dep -> A) first := []*pb.GetTargetGraphResponse{ @@ -713,7 +722,7 @@ func TestCompareTargetGraphs_SourceFileDirectAndPropagation(t *testing.T) { } func TestCompareTargetGraphs_ChangedRuleUnreachableFromAnySeed(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (id 1, rule), no deps first := []*pb.GetTargetGraphResponse{ @@ -769,7 +778,7 @@ func TestCompareTargetGraphs_ChangedRuleUnreachableFromAnySeed(t *testing.T) { } func TestCompareTargetGraphs_ChangedWhenDependenciesChanged(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (id 1, rule) with deps on A first := []*pb.GetTargetGraphResponse{ @@ -845,7 +854,7 @@ func TestCompareTargetGraphs_ChangedWhenDependenciesChanged(t *testing.T) { } func TestCompareTargetGraphs_ChangedWhenAttributesChanged(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T with attribute "key1" -> "value1" first := []*pb.GetTargetGraphResponse{ @@ -918,7 +927,7 @@ func TestCompareTargetGraphs_ChangedWhenAttributesChanged(t *testing.T) { } func TestCompareTargetGraphs_ChangedWhenNewAttributeAdded(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T with one attribute first := []*pb.GetTargetGraphResponse{ @@ -1159,10 +1168,11 @@ func TestGetChangedTargets_CacheHitWithDistanceFilter(t *testing.T) { return nil }).Times(2) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) request := &pb.GetChangedTargetsRequest{ @@ -1256,7 +1266,7 @@ func TestComputeDistances_NilMetadata(t *testing.T) { } func TestCompareTargetGraphs_HashOnlyChangePropagatesViaBFS(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (rule) with deps on source file A (id 10) and attributes first := []*pb.GetTargetGraphResponse{ @@ -1349,7 +1359,7 @@ func TestCompareTargetGraphs_HashOnlyChangePropagatesViaBFS(t *testing.T) { } func TestCompareTargetGraphs_SiblingRuleNotPromotedToSeed(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Source file A (id 1) owned by rule L (id 2). // Rule T (id 3) depends on L (sibling rule), NOT directly on A. @@ -1428,7 +1438,7 @@ func TestCompareTargetGraphs_SiblingRuleNotPromotedToSeed(t *testing.T) { } func TestCompareTargetGraphs_DeletedTargetEmitted(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (rule) exists; New: T is gone. first := []*pb.GetTargetGraphResponse{ diff --git a/controller/gettargetgraph_test.go b/controller/gettargetgraph_test.go index e62856f..ed55bd8 100644 --- a/controller/gettargetgraph_test.go +++ b/controller/gettargetgraph_test.go @@ -46,9 +46,10 @@ func TestGetTargetGraph_CacheMiss_NoSend(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte{})}, nil), ) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) req := &pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ @@ -71,9 +72,10 @@ func TestGetTargetGraph_StorageError_Propagates(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) storagemock := storagemock.NewMockStorage(ctrl) storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, expected) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: storagemock, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: storagemock, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ @@ -97,9 +99,10 @@ func TestGetTargetGraph_DecodeError_ReturnsError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("bad-bytes"))}, nil), ) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: storagemock, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: storagemock, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ @@ -130,9 +133,10 @@ func TestGetTargetGraph_SendsWhenItemPresent(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-xyz"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ @@ -152,9 +156,10 @@ func TestGetTargetGraph_BuildDescriptionMissingRequiredFields_ReturnsError(t *te stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ @@ -173,9 +178,10 @@ func TestGetTargetGraph_MissingBuildDescription_ReturnsError(t *testing.T) { stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{}, stream) assert.Error(t, err) @@ -202,10 +208,11 @@ func TestGetTargetGraph_TreehashNotFound_NoError(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orchestrator.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orchestrator, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, @@ -220,9 +227,10 @@ func TestGetTargetGraph_TreehashReadError(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, @@ -240,9 +248,10 @@ func TestGetTargetGraph_GraphFetchError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("graph error")), ) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, @@ -264,9 +273,10 @@ func TestGetTargetGraph_GraphReadError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil), ) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, @@ -292,9 +302,10 @@ func TestGetTargetGraph_StreamSendError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: storagemock, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: storagemock, + ChunkConfig: _testChunkConfig, }) err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, @@ -321,10 +332,11 @@ func TestGetTargetGraph_GraphNotFound_FallsThrough(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, @@ -343,9 +355,10 @@ func TestGetTargetGraph_GraphReadCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("context canceled")), ) - c := NewController(context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c := newTestController(t, context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, @@ -367,10 +380,11 @@ func TestGetTargetGraph_OrchestratorCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, &storage.NotFoundError{Path: "x"}) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(nil, errors.New("context canceled")) - c := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, + ChunkConfig: _testChunkConfig, }) err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go deleted file mode 100644 index 56ff439..0000000 --- a/controller/testhelper_test.go +++ /dev/null @@ -1,35 +0,0 @@ -// Copyright (c) 2026 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 controller - -import ( - "context" - - "github.com/uber-go/tally" - "github.com/uber/tango/core/common" - "go.uber.org/zap" -) - -func newTestController(logger *zap.Logger) *controller { - return &controller{ - logger: logger, - scope: tally.NoopScope, - targetChunkSize: common.DefaultTargetChunkSize, - changedTargetChunkSize: common.DefaultChangedTargetChunkSize, - metadataMapChunkSize: common.DefaultMetadataMapChunkSize, - totalDurationBuckets: _totalDurationBuckets, - appCtx: context.Background(), - } -} diff --git a/core/bazel/bazel.go b/core/bazel/bazel.go index 217283a..033930e 100644 --- a/core/bazel/bazel.go +++ b/core/bazel/bazel.go @@ -30,11 +30,6 @@ import ( "go.uber.org/zap" ) -const ( - // default query timeout if not provided in config - _queryTimeout = 15 * time.Minute -) - type QueryRequest struct { Query string StartupOptions []string @@ -84,8 +79,8 @@ func NewBazelClient(ctx context.Context, p Params) (*BazelClient, error) { } } timeout := p.QueryTimeout - if timeout == 0 { - timeout = _queryTimeout + if timeout <= 0 { + return nil, fmt.Errorf("query timeout must be > 0, got %s", timeout) } bazelCommand, err := detectBazelExecutable(ctx, p.BazelCommand) if err != nil { diff --git a/core/bazel/bazel_test.go b/core/bazel/bazel_test.go index c76342a..c5cf88b 100644 --- a/core/bazel/bazel_test.go +++ b/core/bazel/bazel_test.go @@ -18,6 +18,7 @@ import ( "context" "os/exec" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -36,6 +37,7 @@ func TestNewBazelClient(t *testing.T) { WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{"FOO": "bar"}, Logger: zap.NewNop().Sugar(), + QueryTimeout: 5 * time.Minute, ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander { return nil }, @@ -48,6 +50,7 @@ func TestNewBazelClient(t *testing.T) { WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{"FOO": "bar"}, Logger: zap.NewNop().Sugar(), + QueryTimeout: 5 * time.Minute, }, }, } @@ -72,6 +75,7 @@ func TestNewBazelClient_WithNilExecCommand(t *testing.T) { WorkspacePath: "/workspace", EnvVarsMap: map[string]string{"KEY": "value"}, Logger: zap.NewNop().Sugar(), + QueryTimeout: 5 * time.Minute, }) require.NoError(t, err) require.NotNil(t, client) @@ -84,3 +88,45 @@ func TestNewBazelClient_WithNilExecCommand(t *testing.T) { assert.Equal(t, "/workspace", execCmd.Dir) assert.Contains(t, execCmd.Env, "KEY=value") } + +func TestNewBazelClient_Validation(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + timeout time.Duration + wantErr bool + }{ + { + name: "valid timeout", + timeout: 5 * time.Minute, + }, + { + name: "query timeout not set", + timeout: 0, + wantErr: true, + }, + { + name: "negative query timeout", + timeout: -1 * time.Second, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, err := NewBazelClient(context.Background(), Params{ + BazelCommand: "bazel", + WorkspacePath: "/tmp/test", + Logger: zap.NewNop().Sugar(), + QueryTimeout: tt.timeout, + }) + if tt.wantErr { + require.Error(t, err, "expected error when %s", tt.name) + } else { + require.NoError(t, err) + } + }) + } +} diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index c506214..09337a8 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -57,6 +57,7 @@ func TestExecuteQuery_Success(t *testing.T) { mockCmd.EXPECT().Wait().Return(nil), ) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -101,6 +102,7 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { mockCmd.EXPECT().Wait().Return(nil), ) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -230,6 +232,7 @@ func TestExecuteQueryInternal_Failures(t *testing.T) { tt.setupMock(mockCmd) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -259,6 +262,7 @@ func TestExecuteQuery_ErrorCase(t *testing.T) { mockCmd.EXPECT().StdoutPipe().Return(nil, errors.New("stdout pipe failed")) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, diff --git a/core/common/utils.go b/core/common/utils.go index 90df802..e46a170 100644 --- a/core/common/utils.go +++ b/core/common/utils.go @@ -28,24 +28,6 @@ import ( "github.com/uber/tango/tangopb" ) -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 = 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 = 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 = 50_000 -) - // ToShortRemote returns the short remote name given a git ssh remote string. // For example, "git@github:uber/tango" will return "uber/tango". func ToShortRemote(remote string) string { @@ -137,8 +119,10 @@ func HashRequestOptions(opts *tangopb.RequestOptions) string { // for typical target rates. const cancelCheckInterval = 4096 -// ResultToGetTargetGraphResponse converts a Result to a GetTargetGraphResponse -func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result) ([]*tangopb.GetTargetGraphResponse, error) { +// 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. +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. @@ -224,14 +208,14 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res attrStrValIDToVal := attrStrValMapper.Invert() // chunk targets into multiple messages for streaming - responses := chunkTargets(optimizedTargets, DefaultTargetChunkSize) + responses := chunkTargets(optimizedTargets, targetChunkSize) for _, meta := range ChunkMetadata( targetIDToName, ruleTypeIDToName, tagIDToName, attrNameIDToName, attrStrValIDToVal, - DefaultMetadataMapChunkSize, + metadataMapChunkSize, ) { responses = append(responses, &tangopb.GetTargetGraphResponse{ Item: &tangopb.GetTargetGraphResponse_Metadata{Metadata: meta}, @@ -242,10 +226,6 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res } 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) @@ -293,10 +273,6 @@ func ChunkMetadata( attrStrValIDToVal map[int32]string, chunkSize int, ) []*tangopb.Metadata { - if chunkSize <= 0 { - chunkSize = DefaultMetadataMapChunkSize - } - targetChunks := splitMap(targetIDToName, chunkSize) attrValChunks := splitMap(attrStrValIDToVal, chunkSize) diff --git a/core/common/utils_test.go b/core/common/utils_test.go index 4172da0..ea1ecd9 100644 --- a/core/common/utils_test.go +++ b/core/common/utils_test.go @@ -186,7 +186,7 @@ func TestChunkTargets(t *testing.T) { func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { t.Parallel() - // 500 targets with DefaultTargetChunkSize=250 → 2 target chunks + 1 metadata = 3 responses + // 500 targets with targetChunkSize=250 → 2 target chunks + 1 metadata = 3 responses numTargets := 500 result := targethasher.Result{ TargetNames: make([]string, numTargets), @@ -198,10 +198,10 @@ func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { result.Targets[name] = &targethasher.Target{Name: name, Hash: []byte{0}, RuleType: "go_library"} } - responses, err := ResultToGetTargetGraphResponse(context.Background(), result) + responses, err := ResultToGetTargetGraphResponse(context.Background(), result, 250, 50_000) require.NoError(t, err) - // 2 target chunks + 1 metadata chunk (500 targets well under DefaultMetadataMapChunkSize) + // 2 target chunks + 1 metadata chunk (500 targets well under metadataMapChunkSize) require.Len(t, responses, 3) for _, resp := range responses[:2] { diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 3071b01..8efaa25 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -112,7 +112,7 @@ 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 := common.ResultToGetTargetGraphResponse(ctx, targethasherResult) // also print the time to convert to GetTargetGraphResponse format + response, err := common.ResultToGetTargetGraphResponse(ctx, targethasherResult, 250, 50_000) 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 05a3036..c5c900d 100644 --- a/example/main.go +++ b/example/main.go @@ -105,11 +105,15 @@ 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{ + ctrl, err := controller.NewController(appCtx, controller.Params{ Logger: zl, Storage: store, Orchestrator: orch, + ChunkConfig: cfg.Service.Chunking, }) + if err != nil { + return fmt.Errorf("failed to create controller: %w", err) + } // YARPC transports and dispatcher grpcTransport := yarpcgrpc.NewTransport() diff --git a/example/tango-config.yaml b/example/tango-config.yaml index ec43e6a..8608c65 100644 --- a/example/tango-config.yaml +++ b/example/tango-config.yaml @@ -15,7 +15,7 @@ repository: - "^@@?bazel_tools/" exclude_external_targets: true bzlmod_enabled: true - query_timeout: 300 + query_timeout_seconds: 300 # Service configuration service: @@ -24,3 +24,7 @@ service: repo_manager_clone_path: "/tmp/tango-repo-manager" # root for worker checkouts; defaults to repo_manager_clone_path/.workers worker_root_path: "/tmp/tango/workers" + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index f46c2ab..0f99272 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -187,7 +187,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, param GetTarget WorkspacePath: ws.Path(), Logger: b.logger, BazelCommand: repoCfg.BazelCommand, - QueryTimeout: time.Duration(repoCfg.QueryTimeout) * time.Second, + QueryTimeout: time.Duration(repoCfg.QueryTimeoutSeconds) * time.Second, StreamLogs: repoCfg.StreamBazelLogs, }) if err != nil { @@ -206,7 +206,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, param GetTarget if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - responses, err := common.ResultToGetTargetGraphResponse(ctx, result) + responses, err := common.ResultToGetTargetGraphResponse(ctx, result, b.config.Service.Chunking.TargetChunkSize, b.config.Service.Chunking.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 f420f74..2fc7793 100644 --- a/orchestrator/testdata/config.yaml +++ b/orchestrator/testdata/config.yaml @@ -7,7 +7,11 @@ repository: - "*.gen.go" exclude_external_targets: true bzlmod_enabled: true - query_timeout: 15 + query_timeout_seconds: 15 service: worker_pool_size: 3 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000