From b29fe48e4978fd30539f14dce0a48ee05434467f Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Mon, 20 Jul 2026 11:19:20 -0700 Subject: [PATCH 1/4] feat(streaming): add streaming split utilities for gRPC message sizing Add internal/streaming package with: - SplitBySize: generic function to split items into groups by wire size - SplitMetadata: splits metadata maps into multiple messages - SplitTargetGraph: splits targets and metadata into wire-safe response chunks Co-Authored-By: Claude Sonnet 5 --- internal/streaming/BUILD.bazel | 24 ++++ internal/streaming/streaming.go | 181 +++++++++++++++++++++++++++ internal/streaming/streaming_test.go | 159 +++++++++++++++++++++++ 3 files changed, 364 insertions(+) create mode 100644 internal/streaming/BUILD.bazel create mode 100644 internal/streaming/streaming.go create mode 100644 internal/streaming/streaming_test.go diff --git a/internal/streaming/BUILD.bazel b/internal/streaming/BUILD.bazel new file mode 100644 index 00000000..d6ea663a --- /dev/null +++ b/internal/streaming/BUILD.bazel @@ -0,0 +1,24 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "streaming", + srcs = ["streaming.go"], + importpath = "github.com/uber/tango/internal/streaming", + visibility = ["//:__subpackages__"], + deps = [ + "//entity", + "//tangopb", + ], +) + +go_test( + name = "streaming_test", + srcs = ["streaming_test.go"], + embed = [":streaming"], + deps = [ + "//entity", + "//tangopb", + "@com_github_stretchr_testify//assert", + "@com_github_stretchr_testify//require", + ], +) diff --git a/internal/streaming/streaming.go b/internal/streaming/streaming.go new file mode 100644 index 00000000..4dc5c49a --- /dev/null +++ b/internal/streaming/streaming.go @@ -0,0 +1,181 @@ +package streaming + +import ( + "fmt" + + "github.com/uber/tango/entity" + "github.com/uber/tango/tangopb" +) + +// Sizer is satisfied by any type that reports its serialized byte length. +type Sizer interface { + Size() int +} + +// SplitBySize splits items into consecutive runs whose cumulative Size() +// stays at or under maxBytes. A single item larger than the budget ships +// alone since it can't be split further. Always returns at least one +// group: empty input yields a single empty group so callers always have a +// message to send on the stream. Returns an error if a multi-group split +// produces any empty group after the first. +func SplitBySize[T Sizer](items []T, maxBytes int) ([][]T, error) { + if len(items) == 0 { + return [][]T{nil}, nil + } + groups := make([][]T, 0, 1) + var current []T + currentBytes := 0 + for _, item := range items { + itemBytes := item.Size() + if len(current) > 0 && currentBytes+itemBytes > maxBytes { + groups = append(groups, current) + current = nil + currentBytes = 0 + } + current = append(current, item) + currentBytes += itemBytes + } + groups = append(groups, current) + for i := 1; i < len(groups); i++ { + if len(groups[i]) == 0 { + return nil, fmt.Errorf("internal error: group %d of %d is empty", i, len(groups)) + } + } + return groups, nil +} + +// SplitMetadata splits the metadata maps into multiple Metadata +// messages so each stays at or under maxBytes. The two large maps (target +// names, attribute string values) are split independently by measured +// entry wire size; consumers merge all metadata before use. The small +// maps (rule_type, tag, attribute_name) are sent in the first message. +// Always returns at least one message. Returns an error if a non-first +// message is completely empty. +func SplitMetadata( + targetIDToName map[int32]string, + ruleTypeIDToName map[int32]string, + tagIDToName map[int32]string, + attrNameIDToName map[int32]string, + attrStrValIDToVal map[int32]string, + maxBytes int, +) ([]*entity.Metadata, error) { + targetGroups := splitMapByBytes(targetIDToName, maxBytes) + attrValGroups := splitMapByBytes(attrStrValIDToVal, maxBytes) + + metas := make([]*entity.Metadata, 0, max(1, len(targetGroups)+len(attrValGroups))) + for _, g := range targetGroups { + metas = append(metas, &entity.Metadata{TargetIDMapping: g}) + } + for _, g := range attrValGroups { + metas = append(metas, &entity.Metadata{AttributeStringValueMapping: g}) + } + if len(metas) == 0 { + metas = append(metas, &entity.Metadata{}) + } + metas[0].RuleTypeMapping = ruleTypeIDToName + metas[0].TagMapping = tagIDToName + metas[0].AttributeNameMapping = attrNameIDToName + + for i := 1; i < len(metas); i++ { + m := metas[i] + if len(m.TargetIDMapping) == 0 && + len(m.RuleTypeMapping) == 0 && + len(m.TagMapping) == 0 && + len(m.AttributeNameMapping) == 0 && + len(m.AttributeStringValueMapping) == 0 { + return nil, fmt.Errorf("internal error: metadata group %d of %d is empty", i, len(metas)) + } + } + + return metas, nil +} + +func splitMapByBytes(m map[int32]string, maxBytes int) []map[int32]string { + if len(m) == 0 { + return nil + } + var groups []map[int32]string + current := make(map[int32]string) + currentBytes := 0 + for k, v := range m { + entryBytes := mapEntryWireSize(k, v) + if len(current) > 0 && currentBytes+entryBytes > maxBytes { + groups = append(groups, current) + current = make(map[int32]string) + currentBytes = 0 + } + current[k] = v + currentBytes += entryBytes + } + if len(current) > 0 { + groups = append(groups, current) + } + return groups +} + +func mapEntryWireSize(k int32, v string) int { + mapEntrySize := 1 + varintSize(uint64(k)) + 1 + len(v) + varintSize(uint64(len(v))) + return mapEntrySize + 1 + varintSize(uint64(mapEntrySize)) +} + +func varintSize(x uint64) int { + n := 1 + for x >= 0x80 { + x >>= 7 + n++ + } + return n +} + +// SplitTargetGraph splits targets and metadata into wire-safe +// entity.GetTargetGraphResponse chunks bounded by maxMessageBytes. +func SplitTargetGraph(targets []entity.OptimizedTarget, meta *entity.Metadata, maxMessageBytes int) ([]entity.GetTargetGraphResponse, error) { + protoTargets := make([]*tangopb.OptimizedTarget, len(targets)) + for i := range targets { + protoTargets[i] = optimizedTargetToProto(&targets[i]) + } + + targetGroups, err := SplitBySize(protoTargets, maxMessageBytes) + if err != nil { + return nil, err + } + + var chunks []entity.GetTargetGraphResponse + idx := 0 + for _, g := range targetGroups { + chunks = append(chunks, entity.GetTargetGraphResponse{ + Targets: targets[idx : idx+len(g)], + }) + idx += len(g) + } + + metaGroups, err := SplitMetadata( + meta.TargetIDMapping, + meta.RuleTypeMapping, + meta.TagMapping, + meta.AttributeNameMapping, + meta.AttributeStringValueMapping, + maxMessageBytes, + ) + if err != nil { + return nil, err + } + for _, m := range metaGroups { + chunks = append(chunks, entity.GetTargetGraphResponse{Metadata: m}) + } + + return chunks, nil +} + +func optimizedTargetToProto(t *entity.OptimizedTarget) *tangopb.OptimizedTarget { + return &tangopb.OptimizedTarget{ + Id: t.ID, + Hash: t.Hash, + DirectDependencies: t.DirectDependencies, + RuleType: t.RuleType, + Tags: t.Tags, + Root: t.Root, + External: t.External, + Attributes: t.Attributes, + } +} diff --git a/internal/streaming/streaming_test.go b/internal/streaming/streaming_test.go new file mode 100644 index 00000000..d56cbfec --- /dev/null +++ b/internal/streaming/streaming_test.go @@ -0,0 +1,159 @@ +package streaming + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/tango/entity" + pb "github.com/uber/tango/tangopb" +) + +func TestSplitBySize(t *testing.T) { + t.Parallel() + + targets := make([]*pb.OptimizedTarget, 25) + for i := range targets { + targets[i] = &pb.OptimizedTarget{Id: int32(i + 1)} + } + maxBytes := targets[0].Size() * 10 + + groups, err := SplitBySize(targets, maxBytes) + require.NoError(t, err) + require.Len(t, groups, 3) + assert.Len(t, groups[0], 10) + assert.Len(t, groups[1], 10) + assert.Len(t, groups[2], 5) + + var total int + for _, g := range groups { + for _, target := range g { + assert.Equal(t, int32(total+1), target.Id) + total++ + } + } + assert.Equal(t, 25, total) +} + +func TestSplitBySize_SingleOversizedItemShipsAlone(t *testing.T) { + t.Parallel() + + oversized := &pb.OptimizedTarget{Id: 1, Hash: strings.Repeat("a", 1000)} + small := &pb.OptimizedTarget{Id: 2} + maxBytes := small.Size() + + groups, err := SplitBySize([]*pb.OptimizedTarget{oversized, small}, maxBytes) + require.NoError(t, err) + require.Len(t, groups, 2) + assert.Equal(t, []*pb.OptimizedTarget{oversized}, groups[0]) + assert.Equal(t, []*pb.OptimizedTarget{small}, groups[1]) +} + +func TestSplitBySize_EmptyInputReturnsOneEmptyGroup(t *testing.T) { + t.Parallel() + + groups, err := SplitBySize([]*pb.OptimizedTarget{}, 100) + require.NoError(t, err) + require.Len(t, groups, 1) + assert.Empty(t, groups[0]) +} + +func TestMapEntryWireSize_MatchesGeneratedSize(t *testing.T) { + t.Parallel() + + m := &pb.Metadata{TargetIdMapping: map[int32]string{7: "hello world"}} + assert.Equal(t, m.Size(), mapEntryWireSize(7, "hello world")) + + m2 := &pb.Metadata{TargetIdMapping: map[int32]string{1234567: strings.Repeat("x", 300)}} + assert.Equal(t, m2.Size(), mapEntryWireSize(1234567, strings.Repeat("x", 300))) +} + +func TestSplitMetadata_SplitsTargetMapByBytes(t *testing.T) { + t.Parallel() + + targetMap := map[int32]string{1: "a", 2: "b", 3: "c", 4: "d"} + ruleType := map[int32]string{1: "go_library"} + entryBytes := mapEntryWireSize(1, "a") + + metas, err := SplitMetadata(targetMap, ruleType, nil, nil, nil, entryBytes*2) + require.NoError(t, err) + require.Len(t, metas, 2) + + merged := map[int32]string{} + for i, meta := range metas { + for k, v := range meta.TargetIDMapping { + merged[k] = v + } + if i == 0 { + assert.Equal(t, ruleType, meta.RuleTypeMapping) + } else { + assert.Empty(t, meta.RuleTypeMapping) + } + } + assert.Equal(t, targetMap, merged) +} + +func TestSplitMetadata_AllEmptyMapsReturnsOneEmptyMessage(t *testing.T) { + t.Parallel() + + metas, err := SplitMetadata(nil, nil, nil, nil, nil, 100) + require.NoError(t, err) + require.Len(t, metas, 1) + assert.Empty(t, metas[0].TargetIDMapping) + assert.Empty(t, metas[0].RuleTypeMapping) +} + +func TestSplitMetadata_EmptyBigMapsStillCarrySmallMaps(t *testing.T) { + t.Parallel() + + ruleType := map[int32]string{1: "go_library"} + metas, err := SplitMetadata(nil, ruleType, nil, nil, nil, 100) + require.NoError(t, err) + require.Len(t, metas, 1) + assert.Equal(t, ruleType, metas[0].RuleTypeMapping) +} + +func TestSplitTargetGraph(t *testing.T) { + t.Parallel() + + numTargets := 50 + targets := make([]entity.OptimizedTarget, numTargets) + for i := range targets { + targets[i] = entity.OptimizedTarget{ID: int32(i + 1), Hash: "ab", RuleType: 1} + } + meta := &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//pkg:a"}, + RuleTypeMapping: map[int32]string{1: "go_library"}, + } + + protoSize := optimizedTargetToProto(&targets[0]).Size() + + tests := []struct { + name string + maxMessageBytes int + wantTargetChunks int + }{ + {name: "25 per chunk", maxMessageBytes: protoSize * 25, wantTargetChunks: 2}, + {name: "10 per chunk", maxMessageBytes: protoSize * 10, wantTargetChunks: 5}, + {name: "all in one", maxMessageBytes: protoSize * 100, wantTargetChunks: 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + chunks, err := SplitTargetGraph(targets, meta, tt.maxMessageBytes) + require.NoError(t, err) + + var targetChunks, totalTargets int + for _, c := range chunks { + if len(c.Targets) > 0 || c.Metadata == nil { + targetChunks++ + totalTargets += len(c.Targets) + } + } + assert.Equal(t, tt.wantTargetChunks, targetChunks) + assert.Equal(t, numTargets, totalTargets) + }) + } +} From eacdca891801ca77a9fd598a8013a3777d74ff85 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Mon, 20 Jul 2026 11:40:51 -0700 Subject: [PATCH 2/4] refactor(streaming): remove proto dependency from SplitTargetGraph Add Size() method directly on entity.OptimizedTarget using packed protobuf wire size calculation, so SplitTargetGraph works with plain entity types throughout without converting to proto intermediaries. Co-Authored-By: Claude Sonnet 5 --- entity/BUILD.bazel | 9 +++- entity/optimized_target.go | 49 ++++++++++++++++++ entity/optimized_target_test.go | 74 ++++++++++++++++++++++++++++ internal/streaming/BUILD.bazel | 5 +- internal/streaming/streaming.go | 20 ++------ internal/streaming/streaming_test.go | 26 +++++----- 6 files changed, 148 insertions(+), 35 deletions(-) create mode 100644 entity/optimized_target_test.go diff --git a/entity/BUILD.bazel b/entity/BUILD.bazel index 3e9b3a9a..6d34fbb9 100644 --- a/entity/BUILD.bazel +++ b/entity/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 = "entity", @@ -13,3 +13,10 @@ go_library( importpath = "github.com/uber/tango/entity", visibility = ["//visibility:public"], ) + +go_test( + name = "entity_test", + srcs = ["optimized_target_test.go"], + embed = [":entity"], + deps = ["@com_github_stretchr_testify//assert"], +) diff --git a/entity/optimized_target.go b/entity/optimized_target.go index 36308ff9..6c8955be 100644 --- a/entity/optimized_target.go +++ b/entity/optimized_target.go @@ -14,6 +14,55 @@ type OptimizedTarget struct { Attributes map[int32]int32 `json:"attributes"` } +// Size returns an estimate of the protobuf wire size for this target. +// Used by streaming splitters to stay within gRPC message size limits. +func (t *OptimizedTarget) Size() int { + n := 0 + if t.ID != 0 { + n += 1 + varintSize(uint64(t.ID)) + } + if len(t.Hash) > 0 { + n += 1 + varintSize(uint64(len(t.Hash))) + len(t.Hash) + } + if len(t.DirectDependencies) > 0 { + dataSize := 0 + for _, d := range t.DirectDependencies { + dataSize += varintSize(uint64(d)) + } + n += 1 + varintSize(uint64(dataSize)) + dataSize + } + if t.RuleType != 0 { + n += 1 + varintSize(uint64(t.RuleType)) + } + if len(t.Tags) > 0 { + dataSize := 0 + for _, tag := range t.Tags { + dataSize += varintSize(uint64(tag)) + } + n += 1 + varintSize(uint64(dataSize)) + dataSize + } + if t.Root { + n += 1 + 1 + } + if t.External { + n += 1 + 1 + } + for k, v := range t.Attributes { + entrySize := 1 + varintSize(uint64(k)) + 1 + varintSize(uint64(v)) + n += 1 + varintSize(uint64(entrySize)) + entrySize + } + return n +} + +func varintSize(x uint64) int { + n := 1 + for x >= 0x80 { + x >>= 7 + n++ + } + return n +} + // Metadata holds the ID-to-string mappings that accompany a set of // OptimizedTarget entries. Consumers merge metadata across chunks before // resolving IDs. diff --git a/entity/optimized_target_test.go b/entity/optimized_target_test.go new file mode 100644 index 00000000..e6fb0a97 --- /dev/null +++ b/entity/optimized_target_test.go @@ -0,0 +1,74 @@ +package entity + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestOptimizedTarget_Size(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + target OptimizedTarget + wantSize int + }{ + { + name: "zero value", + target: OptimizedTarget{}, + wantSize: 0, + }, + { + name: "id only", + target: OptimizedTarget{ID: 1}, + // field tag (1 byte) + varint(1) (1 byte) = 2 + wantSize: 2, + }, + { + name: "with hash", + target: OptimizedTarget{ID: 42, Hash: "abcd"}, + // id: tag(1) + varint(42)(1) = 2 + // hash: tag(1) + len_prefix(1) + 4 bytes = 6 + wantSize: 8, + }, + { + name: "with deps packed", + target: OptimizedTarget{ID: 1, DirectDependencies: []int32{2, 3}}, + // id: tag(1) + varint(1)(1) = 2 + // deps packed: tag(1) + len_prefix(1) + varint(2)(1) + varint(3)(1) = 4 + wantSize: 6, + }, + { + name: "booleans", + target: OptimizedTarget{ID: 1, Root: true, External: true}, + // id: 2, root: tag(1) + 1 = 2, external: tag(1) + 1 = 2 + wantSize: 6, + }, + { + name: "false booleans omitted", + target: OptimizedTarget{ID: 1, Root: false, External: false}, + wantSize: 2, + }, + { + name: "with attributes", + target: OptimizedTarget{ID: 1, Attributes: map[int32]int32{1: 2}}, + // id: 2 + // map entry: tag(1) + len_prefix(1) + [key: tag(1)+varint(1)(1) + val: tag(1)+varint(2)(1)] = 6 + wantSize: 8, + }, + { + name: "tags packed", + target: OptimizedTarget{Tags: []int32{5, 6}}, + // tags packed: tag(1) + len_prefix(1) + varint(5)(1) + varint(6)(1) = 4 + wantSize: 4, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.wantSize, tt.target.Size()) + }) + } +} diff --git a/internal/streaming/BUILD.bazel b/internal/streaming/BUILD.bazel index d6ea663a..89115024 100644 --- a/internal/streaming/BUILD.bazel +++ b/internal/streaming/BUILD.bazel @@ -5,10 +5,7 @@ go_library( srcs = ["streaming.go"], importpath = "github.com/uber/tango/internal/streaming", visibility = ["//:__subpackages__"], - deps = [ - "//entity", - "//tangopb", - ], + deps = ["//entity"], ) go_test( diff --git a/internal/streaming/streaming.go b/internal/streaming/streaming.go index 4dc5c49a..e1907b5b 100644 --- a/internal/streaming/streaming.go +++ b/internal/streaming/streaming.go @@ -4,7 +4,6 @@ import ( "fmt" "github.com/uber/tango/entity" - "github.com/uber/tango/tangopb" ) // Sizer is satisfied by any type that reports its serialized byte length. @@ -130,12 +129,12 @@ func varintSize(x uint64) int { // SplitTargetGraph splits targets and metadata into wire-safe // entity.GetTargetGraphResponse chunks bounded by maxMessageBytes. func SplitTargetGraph(targets []entity.OptimizedTarget, meta *entity.Metadata, maxMessageBytes int) ([]entity.GetTargetGraphResponse, error) { - protoTargets := make([]*tangopb.OptimizedTarget, len(targets)) + sizers := make([]*entity.OptimizedTarget, len(targets)) for i := range targets { - protoTargets[i] = optimizedTargetToProto(&targets[i]) + sizers[i] = &targets[i] } - targetGroups, err := SplitBySize(protoTargets, maxMessageBytes) + targetGroups, err := SplitBySize(sizers, maxMessageBytes) if err != nil { return nil, err } @@ -166,16 +165,3 @@ func SplitTargetGraph(targets []entity.OptimizedTarget, meta *entity.Metadata, m return chunks, nil } - -func optimizedTargetToProto(t *entity.OptimizedTarget) *tangopb.OptimizedTarget { - return &tangopb.OptimizedTarget{ - Id: t.ID, - Hash: t.Hash, - DirectDependencies: t.DirectDependencies, - RuleType: t.RuleType, - Tags: t.Tags, - Root: t.Root, - External: t.External, - Attributes: t.Attributes, - } -} diff --git a/internal/streaming/streaming_test.go b/internal/streaming/streaming_test.go index d56cbfec..3e7636fb 100644 --- a/internal/streaming/streaming_test.go +++ b/internal/streaming/streaming_test.go @@ -13,9 +13,9 @@ import ( func TestSplitBySize(t *testing.T) { t.Parallel() - targets := make([]*pb.OptimizedTarget, 25) + targets := make([]*entity.OptimizedTarget, 25) for i := range targets { - targets[i] = &pb.OptimizedTarget{Id: int32(i + 1)} + targets[i] = &entity.OptimizedTarget{ID: int32(i + 1)} } maxBytes := targets[0].Size() * 10 @@ -29,7 +29,7 @@ func TestSplitBySize(t *testing.T) { var total int for _, g := range groups { for _, target := range g { - assert.Equal(t, int32(total+1), target.Id) + assert.Equal(t, int32(total+1), target.ID) total++ } } @@ -39,21 +39,21 @@ func TestSplitBySize(t *testing.T) { func TestSplitBySize_SingleOversizedItemShipsAlone(t *testing.T) { t.Parallel() - oversized := &pb.OptimizedTarget{Id: 1, Hash: strings.Repeat("a", 1000)} - small := &pb.OptimizedTarget{Id: 2} + oversized := &entity.OptimizedTarget{ID: 1, Hash: strings.Repeat("a", 1000)} + small := &entity.OptimizedTarget{ID: 2} maxBytes := small.Size() - groups, err := SplitBySize([]*pb.OptimizedTarget{oversized, small}, maxBytes) + groups, err := SplitBySize([]*entity.OptimizedTarget{oversized, small}, maxBytes) require.NoError(t, err) require.Len(t, groups, 2) - assert.Equal(t, []*pb.OptimizedTarget{oversized}, groups[0]) - assert.Equal(t, []*pb.OptimizedTarget{small}, groups[1]) + assert.Equal(t, []*entity.OptimizedTarget{oversized}, groups[0]) + assert.Equal(t, []*entity.OptimizedTarget{small}, groups[1]) } func TestSplitBySize_EmptyInputReturnsOneEmptyGroup(t *testing.T) { t.Parallel() - groups, err := SplitBySize([]*pb.OptimizedTarget{}, 100) + groups, err := SplitBySize([]*entity.OptimizedTarget{}, 100) require.NoError(t, err) require.Len(t, groups, 1) assert.Empty(t, groups[0]) @@ -127,16 +127,16 @@ func TestSplitTargetGraph(t *testing.T) { RuleTypeMapping: map[int32]string{1: "go_library"}, } - protoSize := optimizedTargetToProto(&targets[0]).Size() + targetSize := targets[0].Size() tests := []struct { name string maxMessageBytes int wantTargetChunks int }{ - {name: "25 per chunk", maxMessageBytes: protoSize * 25, wantTargetChunks: 2}, - {name: "10 per chunk", maxMessageBytes: protoSize * 10, wantTargetChunks: 5}, - {name: "all in one", maxMessageBytes: protoSize * 100, wantTargetChunks: 1}, + {name: "25 per chunk", maxMessageBytes: targetSize * 25, wantTargetChunks: 2}, + {name: "10 per chunk", maxMessageBytes: targetSize * 10, wantTargetChunks: 5}, + {name: "all in one", maxMessageBytes: targetSize * 100, wantTargetChunks: 1}, } for _, tt := range tests { From 30d505173a6f1f42b85eee1bf155740e40e12fe5 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Mon, 20 Jul 2026 15:31:36 -0700 Subject: [PATCH 3/4] perf(streaming): avoid extra allocation in SplitTargetGraph Inline the size-based grouping directly over the targets slice instead of allocating a parallel []*OptimizedTarget pointer slice to satisfy the Sizer generic constraint. For large monorepos with hundreds of thousands of targets, this eliminates a needless pointer-per-target allocation. Co-Authored-By: Claude Sonnet 5 --- internal/streaming/streaming.go | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/internal/streaming/streaming.go b/internal/streaming/streaming.go index e1907b5b..381ca7e8 100644 --- a/internal/streaming/streaming.go +++ b/internal/streaming/streaming.go @@ -129,24 +129,23 @@ func varintSize(x uint64) int { // SplitTargetGraph splits targets and metadata into wire-safe // entity.GetTargetGraphResponse chunks bounded by maxMessageBytes. func SplitTargetGraph(targets []entity.OptimizedTarget, meta *entity.Metadata, maxMessageBytes int) ([]entity.GetTargetGraphResponse, error) { - sizers := make([]*entity.OptimizedTarget, len(targets)) - for i := range targets { - sizers[i] = &targets[i] - } - - targetGroups, err := SplitBySize(sizers, maxMessageBytes) - if err != nil { - return nil, err - } - var chunks []entity.GetTargetGraphResponse - idx := 0 - for _, g := range targetGroups { - chunks = append(chunks, entity.GetTargetGraphResponse{ - Targets: targets[idx : idx+len(g)], - }) - idx += len(g) + start := 0 + currentBytes := 0 + for i := range targets { + itemBytes := targets[i].Size() + if i > start && currentBytes+itemBytes > maxMessageBytes { + chunks = append(chunks, entity.GetTargetGraphResponse{ + Targets: targets[start:i], + }) + start = i + currentBytes = 0 + } + currentBytes += itemBytes } + chunks = append(chunks, entity.GetTargetGraphResponse{ + Targets: targets[start:], + }) metaGroups, err := SplitMetadata( meta.TargetIDMapping, From 3f052c18d501c22d95b82ac4450b093055948b29 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Tue, 21 Jul 2026 15:01:47 -0700 Subject: [PATCH 4/4] refactor(entity): deduplicate varintSize into internal/streaming/wire Move varintSize to an exported wire.VarintSize in a shared package so both entity and internal/streaming import it instead of each maintaining a copy. Co-Authored-By: Claude Sonnet 5 --- entity/BUILD.bazel | 1 + entity/optimized_target.go | 29 +++++++++++------------------ internal/streaming/BUILD.bazel | 5 ++++- internal/streaming/streaming.go | 14 +++----------- internal/streaming/wire/BUILD.bazel | 8 ++++++++ internal/streaming/wire/wire.go | 11 +++++++++++ 6 files changed, 38 insertions(+), 30 deletions(-) create mode 100644 internal/streaming/wire/BUILD.bazel create mode 100644 internal/streaming/wire/wire.go diff --git a/entity/BUILD.bazel b/entity/BUILD.bazel index 6d34fbb9..f6c53680 100644 --- a/entity/BUILD.bazel +++ b/entity/BUILD.bazel @@ -12,6 +12,7 @@ go_library( ], importpath = "github.com/uber/tango/entity", visibility = ["//visibility:public"], + deps = ["//internal/streaming/wire"], ) go_test( diff --git a/entity/optimized_target.go b/entity/optimized_target.go index 6c8955be..30194dfd 100644 --- a/entity/optimized_target.go +++ b/entity/optimized_target.go @@ -1,5 +1,7 @@ package entity +import "github.com/uber/tango/internal/streaming/wire" + // OptimizedTarget is the compact, ID-mapped representation of a target used // for streaming and storage. String fields are replaced with int32 IDs that // reference the accompanying Metadata maps. @@ -19,27 +21,27 @@ type OptimizedTarget struct { func (t *OptimizedTarget) Size() int { n := 0 if t.ID != 0 { - n += 1 + varintSize(uint64(t.ID)) + n += 1 + wire.VarintSize(uint64(t.ID)) } if len(t.Hash) > 0 { - n += 1 + varintSize(uint64(len(t.Hash))) + len(t.Hash) + n += 1 + wire.VarintSize(uint64(len(t.Hash))) + len(t.Hash) } if len(t.DirectDependencies) > 0 { dataSize := 0 for _, d := range t.DirectDependencies { - dataSize += varintSize(uint64(d)) + dataSize += wire.VarintSize(uint64(d)) } - n += 1 + varintSize(uint64(dataSize)) + dataSize + n += 1 + wire.VarintSize(uint64(dataSize)) + dataSize } if t.RuleType != 0 { - n += 1 + varintSize(uint64(t.RuleType)) + n += 1 + wire.VarintSize(uint64(t.RuleType)) } if len(t.Tags) > 0 { dataSize := 0 for _, tag := range t.Tags { - dataSize += varintSize(uint64(tag)) + dataSize += wire.VarintSize(uint64(tag)) } - n += 1 + varintSize(uint64(dataSize)) + dataSize + n += 1 + wire.VarintSize(uint64(dataSize)) + dataSize } if t.Root { n += 1 + 1 @@ -48,17 +50,8 @@ func (t *OptimizedTarget) Size() int { n += 1 + 1 } for k, v := range t.Attributes { - entrySize := 1 + varintSize(uint64(k)) + 1 + varintSize(uint64(v)) - n += 1 + varintSize(uint64(entrySize)) + entrySize - } - return n -} - -func varintSize(x uint64) int { - n := 1 - for x >= 0x80 { - x >>= 7 - n++ + entrySize := 1 + wire.VarintSize(uint64(k)) + 1 + wire.VarintSize(uint64(v)) + n += 1 + wire.VarintSize(uint64(entrySize)) + entrySize } return n } diff --git a/internal/streaming/BUILD.bazel b/internal/streaming/BUILD.bazel index 89115024..304c2769 100644 --- a/internal/streaming/BUILD.bazel +++ b/internal/streaming/BUILD.bazel @@ -5,7 +5,10 @@ go_library( srcs = ["streaming.go"], importpath = "github.com/uber/tango/internal/streaming", visibility = ["//:__subpackages__"], - deps = ["//entity"], + deps = [ + "//entity", + "//internal/streaming/wire", + ], ) go_test( diff --git a/internal/streaming/streaming.go b/internal/streaming/streaming.go index 381ca7e8..84203a8d 100644 --- a/internal/streaming/streaming.go +++ b/internal/streaming/streaming.go @@ -4,6 +4,7 @@ import ( "fmt" "github.com/uber/tango/entity" + "github.com/uber/tango/internal/streaming/wire" ) // Sizer is satisfied by any type that reports its serialized byte length. @@ -113,17 +114,8 @@ func splitMapByBytes(m map[int32]string, maxBytes int) []map[int32]string { } func mapEntryWireSize(k int32, v string) int { - mapEntrySize := 1 + varintSize(uint64(k)) + 1 + len(v) + varintSize(uint64(len(v))) - return mapEntrySize + 1 + varintSize(uint64(mapEntrySize)) -} - -func varintSize(x uint64) int { - n := 1 - for x >= 0x80 { - x >>= 7 - n++ - } - return n + mapEntrySize := 1 + wire.VarintSize(uint64(k)) + 1 + len(v) + wire.VarintSize(uint64(len(v))) + return mapEntrySize + 1 + wire.VarintSize(uint64(mapEntrySize)) } // SplitTargetGraph splits targets and metadata into wire-safe diff --git a/internal/streaming/wire/BUILD.bazel b/internal/streaming/wire/BUILD.bazel new file mode 100644 index 00000000..11345314 --- /dev/null +++ b/internal/streaming/wire/BUILD.bazel @@ -0,0 +1,8 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "wire", + srcs = ["wire.go"], + importpath = "github.com/uber/tango/internal/streaming/wire", + visibility = ["//:__subpackages__"], +) diff --git a/internal/streaming/wire/wire.go b/internal/streaming/wire/wire.go new file mode 100644 index 00000000..f452c7a4 --- /dev/null +++ b/internal/streaming/wire/wire.go @@ -0,0 +1,11 @@ +package wire + +// VarintSize returns the number of bytes needed to encode x as a protobuf varint. +func VarintSize(x uint64) int { + n := 1 + for x >= 0x80 { + x >>= 7 + n++ + } + return n +}