config: bound stream chunks by real message size, not estimated

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

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: _defaultMaxMessageBytes,
+		},
+		{
+			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

config/service_config.go

@@ -24,21 +24,9 @@ type ServiceConfig struct {
 	// and worker checkouts. Required. Layout: <workspaces_root_path>/<repo>/ for
 	// origin clones and <workspaces_root_path>/.workers/<repo>/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 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"`
 }

controller/BUILD.bazel

@@ -14,13 +14,13 @@ go_library(
     importpath = "github.com/uber/tango/controller",
     visibility = ["//visibility:public"],
     deps = [
-        "//config",
         "//core/common",
         "//core/storage",
         "//entity",
         "//internal/cachekey",
         "//internal/mapper",
         "//internal/mapper/idmapper",
+        "//internal/streaming",
         "//internal/targetdiff",
         "//orchestrator",
         "//tangopb",
@@ -46,6 +46,7 @@ go_test(
         "//core/storage",
         "//core/storage/storagemock",
         "//entity",
+        "//internal/mapper",
         "//orchestrator/orchestratormock",
         "//tangopb",
         "//tangopb/tangopbmock",

controller/controller.go

@@ -19,8 +19,6 @@ import (
 	"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"
@@ -29,27 +27,28 @@ 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
-	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.
 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
@@ -64,28 +63,20 @@ 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
+	// 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,
-		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,
 	}
 }
 

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"
@@ -262,7 +263,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 +275,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)
 	}
@@ -447,37 +448,32 @@ 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)
-		}
-		results = append(results, &pb.GetChangedTargetsResponse{
-			Item: &pb.GetChangedTargetsResponse_ChangedTargets{
-				ChangedTargets: &pb.ChangedTargets{
-					ChangedTargets: changed[i:end],
-				},
-			},
-		})
+	changedGroups, err := streaming.SplitBySize(changed, c.maxMessageBytes)
+	if err != nil {
+		return nil, err
 	}
-	if len(results) == 0 {
+	for _, group := range changedGroups {
 		results = append(results, &pb.GetChangedTargetsResponse{
 			Item: &pb.GetChangedTargetsResponse_ChangedTargets{
-				ChangedTargets: &pb.ChangedTargets{},
+				ChangedTargets: &pb.ChangedTargets{ChangedTargets: group},
 			},
 		})
 	}
-	for _, meta := range common.ChunkMetadata(
+	metaGroups, err := streaming.SplitMetadata(
 		mappers.target.Invert(),
 		mappers.ruleType.Invert(),
 		mappers.tag.Invert(),
 		mappers.attrName.Invert(),
 		mappers.attrVal.Invert(),
-		c.metadataMapChunkSize,
-	) {
+		c.maxMessageBytes,
+	)
+	if err != nil {
+		return nil, err
+	}
+	for _, meta := range metaGroups {
 		results = append(results, &pb.GetChangedTargetsResponse{
 			Item: &pb.GetChangedTargetsResponse_Metadata{
-				Metadata: meta,
+				Metadata: mapper.EntityMetadataToProto(meta),
 			},
 		})
 	}