From ce1d0c936bc0de2fa53aa9337b97c65c77cd1cdb Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Wed, 15 Jul 2026 10:57:38 -0700 Subject: [PATCH 1/6] feat(mapper): add wire-size-based chunking and ResultToGetTargetGraphResponse MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce internal/mapper/chunk.go with BySize (generic proto chunker bounded by real serialized size) and ChunkMetadata (map splitter). Add ResultToGetTargetGraphResponse to internal/mapper/target_graph.go, converting a targethasher.Result into chunked stream responses. Both old (core/common) and new (internal/mapper) implementations coexist — callsite migration follows in a subsequent PR. Co-Authored-By: Claude Sonnet 5 --- internal/mapper/BUILD.bazel | 6 ++ internal/mapper/chunk.go | 147 +++++++++++++++++++++++++++ internal/mapper/chunk_test.go | 128 +++++++++++++++++++++++ internal/mapper/target_graph.go | 113 ++++++++++++++++++++ internal/mapper/target_graph_test.go | 72 +++++++++++++ 5 files changed, 466 insertions(+) create mode 100644 internal/mapper/chunk.go create mode 100644 internal/mapper/chunk_test.go diff --git a/internal/mapper/BUILD.bazel b/internal/mapper/BUILD.bazel index 77903eb8..929ad750 100644 --- a/internal/mapper/BUILD.bazel +++ b/internal/mapper/BUILD.bazel @@ -4,13 +4,17 @@ go_library( name = "mapper", srcs = [ "build_description.go", + "chunk.go", "target_graph.go", ], importpath = "github.com/uber/tango/internal/mapper", visibility = ["//visibility:public"], deps = [ + "//core/targethasher", + "//internal/mapper/idmapper", "//entity", "//tangopb", + "@com_github_bazelbuild_buildtools//build_proto", ], ) @@ -18,10 +22,12 @@ go_test( name = "mapper_test", srcs = [ "build_description_test.go", + "chunk_test.go", "target_graph_test.go", ], embed = [":mapper"], deps = [ + "//core/targethasher", "//entity", "//tangopb", "@com_github_stretchr_testify//assert", diff --git a/internal/mapper/chunk.go b/internal/mapper/chunk.go new file mode 100644 index 00000000..c2cf073b --- /dev/null +++ b/internal/mapper/chunk.go @@ -0,0 +1,147 @@ +// 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 mapper + +import ( + "fmt" + + "github.com/uber/tango/tangopb" +) + +// sizer is satisfied by generated proto message types, which all expose a +// Size() method computing their exact serialized byte length with no +// allocation or marshaling. +type sizer interface { + Size() int +} + +// BySize splits items into consecutive runs whose cumulative Size() stays +// at or under maxMessageBytes. A single item larger than the budget ships +// alone since it can't be split further. Always returns at least one chunk: +// empty input yields a single empty chunk so callers always have a message +// to send on the stream. Returns an error if a multi-chunk split produces +// any empty chunk after the first (which would indicate a bug in the +// splitting logic). +func BySize[T sizer](items []T, maxMessageBytes int) ([][]T, error) { + if len(items) == 0 { + return [][]T{nil}, nil + } + chunks := make([][]T, 0, 1) + var current []T + currentBytes := 0 + for _, item := range items { + itemBytes := item.Size() + if len(current) > 0 && currentBytes+itemBytes > maxMessageBytes { + chunks = append(chunks, current) + current = nil + currentBytes = 0 + } + current = append(current, item) + currentBytes += itemBytes + } + chunks = append(chunks, current) + for i := 1; i < len(chunks); i++ { + if len(chunks[i]) == 0 { + return nil, fmt.Errorf("internal error: chunk %d of %d is empty", i, len(chunks)) + } + } + return chunks, nil +} + +// ChunkMetadata splits the metadata maps into multiple Metadata messages so +// each stays at or under maxMessageBytes. The two large maps (target names, +// attribute string values) are split independently by measured entry wire +// size; consumers merge all Metadata messages before use. The small maps +// (rule_type, tag, attribute_name) are sent in the first chunk. +// Always returns at least one chunk so callers always have a message to +// send on the stream. Returns an error if a non-first chunk is completely +// empty (all maps nil/empty). +func ChunkMetadata( + targetIDToName map[int32]string, + ruleTypeIDToName map[int32]string, + tagIDToName map[int32]string, + attrNameIDToName map[int32]string, + attrStrValIDToVal map[int32]string, + maxMessageBytes int, +) ([]*tangopb.Metadata, error) { + targetChunks := splitMapByBytes(targetIDToName, maxMessageBytes) + attrValChunks := splitMapByBytes(attrStrValIDToVal, maxMessageBytes) + + chunks := make([]*tangopb.Metadata, 0, max(1, len(targetChunks)+len(attrValChunks))) + for _, c := range targetChunks { + chunks = append(chunks, &tangopb.Metadata{TargetIdMapping: c}) + } + for _, c := range attrValChunks { + chunks = append(chunks, &tangopb.Metadata{AttributeStringValueMapping: c}) + } + if len(chunks) == 0 { + chunks = append(chunks, &tangopb.Metadata{}) + } + chunks[0].RuleTypeMapping = ruleTypeIDToName + chunks[0].TagMapping = tagIDToName + chunks[0].AttributeNameMapping = attrNameIDToName + + for i := 1; i < len(chunks); i++ { + c := chunks[i] + if len(c.TargetIdMapping) == 0 && + len(c.RuleTypeMapping) == 0 && + len(c.TagMapping) == 0 && + len(c.AttributeNameMapping) == 0 && + len(c.AttributeStringValueMapping) == 0 { + return nil, fmt.Errorf("internal error: metadata chunk %d of %d is empty", i, len(chunks)) + } + } + + return chunks, nil +} + +func splitMapByBytes(m map[int32]string, maxBytes int) []map[int32]string { + if len(m) == 0 { + return nil + } + var chunks []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 { + chunks = append(chunks, current) + current = make(map[int32]string) + currentBytes = 0 + } + current[k] = v + currentBytes += entryBytes + } + if len(current) > 0 { + chunks = append(chunks, current) + } + return chunks +} + +// mapEntryWireSize returns the serialized size of one map[int32]string entry, +// mirroring the generated Metadata.Size() computation exactly. +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 +} diff --git a/internal/mapper/chunk_test.go b/internal/mapper/chunk_test.go new file mode 100644 index 00000000..c6fc1f11 --- /dev/null +++ b/internal/mapper/chunk_test.go @@ -0,0 +1,128 @@ +// 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 mapper + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + pb "github.com/uber/tango/tangopb" +) + +func TestBySize(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 + + chunks, err := BySize(targets, maxBytes) + require.NoError(t, err) + require.Len(t, chunks, 3) + assert.Len(t, chunks[0], 10) + assert.Len(t, chunks[1], 10) + assert.Len(t, chunks[2], 5) + + var total int + for _, c := range chunks { + for _, target := range c { + assert.Equal(t, int32(total+1), target.Id) + total++ + } + } + assert.Equal(t, 25, total) +} + +func TestBySize_SingleOversizedItemShipsAlone(t *testing.T) { + t.Parallel() + + oversized := &pb.OptimizedTarget{Id: 1, Hash: strings.Repeat("a", 1000)} + small := &pb.OptimizedTarget{Id: 2} + maxBytes := small.Size() + + chunks, err := BySize([]*pb.OptimizedTarget{oversized, small}, maxBytes) + require.NoError(t, err) + require.Len(t, chunks, 2) + assert.Equal(t, []*pb.OptimizedTarget{oversized}, chunks[0]) + assert.Equal(t, []*pb.OptimizedTarget{small}, chunks[1]) +} + +func TestBySize_EmptyInputReturnsOneEmptyChunk(t *testing.T) { + t.Parallel() + + chunks, err := BySize([]*pb.OptimizedTarget{}, 100) + require.NoError(t, err) + require.Len(t, chunks, 1) + assert.Empty(t, chunks[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 TestChunkMetadata_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") + + chunks, err := ChunkMetadata(targetMap, ruleType, nil, nil, nil, entryBytes*2) + require.NoError(t, err) + require.Len(t, chunks, 2) + + merged := map[int32]string{} + for i, meta := range chunks { + 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 TestChunkMetadata_AllEmptyMapsReturnsOneEmptyChunk(t *testing.T) { + t.Parallel() + + chunks, err := ChunkMetadata(nil, nil, nil, nil, nil, 100) + require.NoError(t, err) + require.Len(t, chunks, 1) + assert.Empty(t, chunks[0].TargetIdMapping) + assert.Empty(t, chunks[0].RuleTypeMapping) +} + +func TestChunkMetadata_EmptyBigMapsStillCarrySmallMaps(t *testing.T) { + t.Parallel() + + ruleType := map[int32]string{1: "go_library"} + chunks, err := ChunkMetadata(nil, ruleType, nil, nil, nil, 100) + require.NoError(t, err) + require.Len(t, chunks, 1) + assert.Equal(t, ruleType, chunks[0].RuleTypeMapping) +} diff --git a/internal/mapper/target_graph.go b/internal/mapper/target_graph.go index 45717d82..e3c551af 100644 --- a/internal/mapper/target_graph.go +++ b/internal/mapper/target_graph.go @@ -1,12 +1,19 @@ package mapper import ( + "context" + "encoding/hex" "errors" + buildpb "github.com/bazelbuild/buildtools/build_proto" + "github.com/uber/tango/core/targethasher" "github.com/uber/tango/entity" + "github.com/uber/tango/internal/mapper/idmapper" "github.com/uber/tango/tangopb" ) +const cancelCheckInterval = 4096 + // ProtoToGetTargetGraphRequest converts a proto GetTargetGraphRequest to the // domain type. Returns an error if req is nil or its BuildDescription fails // validation (see ProtoToBuildDescription). @@ -24,3 +31,109 @@ func ProtoToGetTargetGraphRequest(req *tangopb.GetTargetGraphRequest) (entity.Ge BypassCache: req.GetBypassCache(), }, nil } + +// ResultToGetTargetGraphResponse converts a targethasher.Result into chunked +// GetTargetGraphResponse messages ready for streaming or storage. Each message +// is bounded by maxMessageBytes of real wire size (see chunkTargets). +func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result, maxMessageBytes int) ([]*tangopb.GetTargetGraphResponse, error) { + 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() + + 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, + } + + if t.RuleType != "" { + ot.RuleType = ruleTypeMapper.ID(t.RuleType) + } + + 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 { + 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) + } + + targetIDToName := make(map[int32]string, len(targetNamesMapping)) + for s, id := range targetNamesMapping { + targetIDToName[id] = s + } + + var responses []*tangopb.GetTargetGraphResponse + targetChunks, err := BySize(optimizedTargets, maxMessageBytes) + if err != nil { + return nil, err + } + for _, chunk := range targetChunks { + responses = append(responses, &tangopb.GetTargetGraphResponse{ + Item: &tangopb.GetTargetGraphResponse_Targets{ + Targets: &tangopb.OptimizedTargets{Targets: chunk}, + }, + }) + } + metaChunks, err := ChunkMetadata( + targetIDToName, + ruleTypeMapper.Invert(), + tagMapper.Invert(), + attrNameMapper.Invert(), + attrStrValMapper.Invert(), + maxMessageBytes, + ) + if err != nil { + return nil, err + } + for _, meta := range metaChunks { + responses = append(responses, &tangopb.GetTargetGraphResponse{ + Item: &tangopb.GetTargetGraphResponse_Metadata{Metadata: meta}, + }) + } + + return responses, nil +} diff --git a/internal/mapper/target_graph_test.go b/internal/mapper/target_graph_test.go index 0d0d92fb..f86759c5 100644 --- a/internal/mapper/target_graph_test.go +++ b/internal/mapper/target_graph_test.go @@ -1,11 +1,14 @@ package mapper import ( + "context" + "fmt" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/uber/tango/core/targethasher" "github.com/uber/tango/entity" "github.com/uber/tango/tangopb" ) @@ -86,3 +89,72 @@ func TestProtoToGetTargetGraphRequest(t *testing.T) { }) } } + +func TestResultToGetTargetGraphResponse_EmptyResult(t *testing.T) { + t.Parallel() + + responses, err := ResultToGetTargetGraphResponse(context.Background(), targethasher.Result{}, 1<<20) + require.NoError(t, err) + require.Len(t, responses, 2) + + targets, ok := responses[0].Item.(*tangopb.GetTargetGraphResponse_Targets) + require.True(t, ok) + assert.Empty(t, targets.Targets.Targets) + + _, ok = responses[1].Item.(*tangopb.GetTargetGraphResponse_Metadata) + assert.True(t, ok) +} + +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"} + } + + baseline, err := ResultToGetTargetGraphResponse(context.Background(), result, 1<<30) + require.NoError(t, err) + var targetBytes int + for _, resp := range baseline { + if targets, ok := resp.Item.(*tangopb.GetTargetGraphResponse_Targets); ok { + targetBytes = targets.Targets.Targets[0].Size() + break + } + } + require.NotZero(t, targetBytes) + + tests := []struct { + name string + maxMessageBytes int + wantTargetChunks int + }{ + {name: "25 per chunk", maxMessageBytes: targetBytes * 25, wantTargetChunks: 2}, + {name: "10 per chunk", maxMessageBytes: targetBytes * 10, wantTargetChunks: 5}, + {name: "all in one chunk", maxMessageBytes: targetBytes * 100, wantTargetChunks: 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + responses, err := ResultToGetTargetGraphResponse(context.Background(), result, tt.maxMessageBytes) + require.NoError(t, err) + + var targetChunks, totalTargets int + for _, resp := range responses { + if targets, ok := resp.Item.(*tangopb.GetTargetGraphResponse_Targets); ok { + targetChunks++ + totalTargets += len(targets.Targets.Targets) + } + } + assert.Equal(t, tt.wantTargetChunks, targetChunks) + assert.Equal(t, numTargets, totalTargets) + }) + } +} From 436746a93112b45aaf4cff320e1072c6e5793f53 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Wed, 15 Jul 2026 11:05:42 -0700 Subject: [PATCH 2/6] chore: run gazelle to sort BUILD deps Co-Authored-By: Claude Sonnet 5 --- internal/mapper/BUILD.bazel | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/mapper/BUILD.bazel b/internal/mapper/BUILD.bazel index 929ad750..c4877cee 100644 --- a/internal/mapper/BUILD.bazel +++ b/internal/mapper/BUILD.bazel @@ -11,8 +11,8 @@ go_library( visibility = ["//visibility:public"], deps = [ "//core/targethasher", - "//internal/mapper/idmapper", "//entity", + "//internal/mapper/idmapper", "//tangopb", "@com_github_bazelbuild_buildtools//build_proto", ], From 061e7afdd46e7a6a25d7085eed0b2575435978e7 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Wed, 15 Jul 2026 21:19:06 -0700 Subject: [PATCH 3/6] refactor(mapper): split ResultToGetTargetGraphResponse into entity and proto steps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add entity.OptimizedTarget and entity.TargetGraph as proto-free domain types for target graphs. Split the monolithic ResultToGetTargetGraphResponse into: - ResultToTargetGraph: targethasher.Result → entity.TargetGraph (proto-free) - TargetGraphToProto: entity.TargetGraph → chunked proto responses This lets the orchestrator work with entity types while keeping proto conversion confined to the mapper and controller layers. Co-Authored-By: Claude Sonnet 5 --- entity/BUILD.bazel | 1 + entity/optimized_target.go | 21 ++++++ internal/mapper/target_graph.go | 95 ++++++++++++++++++++-------- internal/mapper/target_graph_test.go | 42 +++++++++++- 4 files changed, 130 insertions(+), 29 deletions(-) create mode 100644 entity/optimized_target.go diff --git a/entity/BUILD.bazel b/entity/BUILD.bazel index deff8984..c64bc795 100644 --- a/entity/BUILD.bazel +++ b/entity/BUILD.bazel @@ -5,6 +5,7 @@ go_library( srcs = [ "build_description.go", "computation_strategy.go", + "optimized_target.go", "target_graph.go", ], importpath = "github.com/uber/tango/entity", diff --git a/entity/optimized_target.go b/entity/optimized_target.go new file mode 100644 index 00000000..f654b2da --- /dev/null +++ b/entity/optimized_target.go @@ -0,0 +1,21 @@ +package entity + +// OptimizedTarget is the domain representation of a single build target +// with its hash, dependencies, and metadata. All references use string +// names — ID assignment is a serialization concern handled by the mapper. +type OptimizedTarget struct { + Name string + Hash string + Dependencies []string + RuleType string + Tags []string + Root bool + External bool + Attributes map[string]string +} + +// TargetGraph is the domain representation of a complete target graph: +// a topologically sorted list of optimized targets. +type TargetGraph struct { + Targets []OptimizedTarget +} diff --git a/internal/mapper/target_graph.go b/internal/mapper/target_graph.go index e3c551af..b6daf64c 100644 --- a/internal/mapper/target_graph.go +++ b/internal/mapper/target_graph.go @@ -32,13 +32,54 @@ func ProtoToGetTargetGraphRequest(req *tangopb.GetTargetGraphRequest) (entity.Ge }, nil } -// ResultToGetTargetGraphResponse converts a targethasher.Result into chunked -// GetTargetGraphResponse messages ready for streaming or storage. Each message -// is bounded by maxMessageBytes of real wire size (see chunkTargets). -func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result, maxMessageBytes int) ([]*tangopb.GetTargetGraphResponse, error) { - targetNamesMapping := make(map[string]int32, len(result.TargetNames)) +// ResultToTargetGraph converts a targethasher.Result into a proto-free +// entity.TargetGraph. Targets are emitted in the topological order given +// by result.TargetNames. Only STRING attributes with non-nil name and +// value are included. +func ResultToTargetGraph(ctx context.Context, result targethasher.Result) (entity.TargetGraph, error) { + targets := make([]entity.OptimizedTarget, 0, len(result.TargetNames)) for i, name := range result.TargetNames { - targetNamesMapping[name] = int32(i + 1) + if i%cancelCheckInterval == 0 { + if err := ctx.Err(); err != nil { + return entity.TargetGraph{}, err + } + } + t, ok := result.Targets[name] + if !ok { + continue + } + ot := entity.OptimizedTarget{ + Name: name, + Hash: hex.EncodeToString(t.Hash), + Dependencies: t.Deps, + RuleType: t.RuleType, + Tags: t.Tags, + Root: t.Root, + External: t.External, + } + if len(t.Attributes) > 0 { + attrs := make(map[string]string, len(t.Attributes)) + for _, attr := range t.Attributes { + if attr.GetType() == buildpb.Attribute_STRING && attr.Name != nil && attr.StringValue != nil { + attrs[*attr.Name] = *attr.StringValue + } + } + if len(attrs) > 0 { + ot.Attributes = attrs + } + } + targets = append(targets, ot) + } + return entity.TargetGraph{Targets: targets}, nil +} + +// TargetGraphToProto converts an entity.TargetGraph into chunked +// GetTargetGraphResponse messages ready for streaming. Each message is +// bounded by maxMessageBytes of real wire size. +func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessageBytes int) ([]*tangopb.GetTargetGraphResponse, error) { + targetNamesMapping := make(map[string]int32, len(graph.Targets)) + for i, t := range graph.Targets { + targetNamesMapping[t.Name] = int32(i + 1) } ruleTypeMapper := idmapper.NewMapper() @@ -46,20 +87,15 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res attrNameMapper := idmapper.NewMapper() attrStrValMapper := idmapper.NewMapper() - optimizedTargets := make([]*tangopb.OptimizedTarget, 0, len(result.Targets)) - - n := 0 - for _, t := range result.Targets { - if n%cancelCheckInterval == 0 { + optimizedTargets := make([]*tangopb.OptimizedTarget, 0, len(graph.Targets)) + for i, t := range graph.Targets { + if i%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 { + depIDs := make([]int32, 0, len(t.Dependencies)) + for _, depName := range t.Dependencies { depID, ok := targetNamesMapping[depName] if !ok { continue @@ -68,15 +104,16 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res } ot := &tangopb.OptimizedTarget{ - Id: nameID, - Hash: hex.EncodeToString(t.Hash), + Id: targetNamesMapping[t.Name], + Hash: t.Hash, DirectDependencies: depIDs, + Root: t.Root, + External: t.External, } if t.RuleType != "" { ot.RuleType = ruleTypeMapper.ID(t.RuleType) } - if len(t.Tags) > 0 { tagIDs := make([]int32, 0, len(t.Tags)) for _, tag := range t.Tags { @@ -84,16 +121,10 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res } 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 { - 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 - } + for k, v := range t.Attributes { + attrs[attrNameMapper.ID(k)] = attrStrValMapper.ID(v) } ot.Attributes = attrs } @@ -137,3 +168,13 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res return responses, nil } + +// ResultToGetTargetGraphResponse is a convenience that composes +// ResultToTargetGraph and TargetGraphToProto. +func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result, maxMessageBytes int) ([]*tangopb.GetTargetGraphResponse, error) { + graph, err := ResultToTargetGraph(ctx, result) + if err != nil { + return nil, err + } + return TargetGraphToProto(ctx, graph, maxMessageBytes) +} diff --git a/internal/mapper/target_graph_test.go b/internal/mapper/target_graph_test.go index f86759c5..f3d1679d 100644 --- a/internal/mapper/target_graph_test.go +++ b/internal/mapper/target_graph_test.go @@ -90,10 +90,48 @@ func TestProtoToGetTargetGraphRequest(t *testing.T) { } } -func TestResultToGetTargetGraphResponse_EmptyResult(t *testing.T) { +func TestResultToTargetGraph(t *testing.T) { t.Parallel() - responses, err := ResultToGetTargetGraphResponse(context.Background(), targethasher.Result{}, 1<<20) + result := targethasher.Result{ + TargetNames: []string{"//pkg:a", "//pkg:b"}, + Targets: map[string]*targethasher.Target{ + "//pkg:a": {Name: "//pkg:a", Hash: []byte{0xab}, RuleType: "go_library", Deps: []string{"//pkg:b"}, Tags: []string{"manual"}, Root: true}, + "//pkg:b": {Name: "//pkg:b", Hash: []byte{0xcd}, RuleType: "go_test", External: true}, + }, + } + + graph, err := ResultToTargetGraph(context.Background(), result) + require.NoError(t, err) + require.Len(t, graph.Targets, 2) + + a := graph.Targets[0] + assert.Equal(t, "//pkg:a", a.Name) + assert.Equal(t, "ab", a.Hash) + assert.Equal(t, "go_library", a.RuleType) + assert.Equal(t, []string{"//pkg:b"}, a.Dependencies) + assert.Equal(t, []string{"manual"}, a.Tags) + assert.True(t, a.Root) + assert.False(t, a.External) + + b := graph.Targets[1] + assert.Equal(t, "//pkg:b", b.Name) + assert.Equal(t, "cd", b.Hash) + assert.True(t, b.External) +} + +func TestResultToTargetGraph_EmptyResult(t *testing.T) { + t.Parallel() + + graph, err := ResultToTargetGraph(context.Background(), targethasher.Result{}) + require.NoError(t, err) + assert.Empty(t, graph.Targets) +} + +func TestTargetGraphToProto_EmptyGraph(t *testing.T) { + t.Parallel() + + responses, err := TargetGraphToProto(context.Background(), entity.TargetGraph{}, 1<<20) require.NoError(t, err) require.Len(t, responses, 2) From f91f68b5ca4afb98ad19fda59eb53e01f09df5ba Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 16 Jul 2026 13:27:05 -0700 Subject: [PATCH 4/6] refactor(mapper): remove ResultToGetTargetGraphResponse convenience function Callers should use the two-step flow (ResultToTargetGraph then TargetGraphToProto) directly, keeping entity and proto concerns explicitly separate. Co-Authored-By: Claude Sonnet 5 --- internal/mapper/target_graph.go | 10 ---------- internal/mapper/target_graph_test.go | 9 ++++++--- 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/internal/mapper/target_graph.go b/internal/mapper/target_graph.go index b6daf64c..bcb9a19f 100644 --- a/internal/mapper/target_graph.go +++ b/internal/mapper/target_graph.go @@ -168,13 +168,3 @@ func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessag return responses, nil } - -// ResultToGetTargetGraphResponse is a convenience that composes -// ResultToTargetGraph and TargetGraphToProto. -func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result, maxMessageBytes int) ([]*tangopb.GetTargetGraphResponse, error) { - graph, err := ResultToTargetGraph(ctx, result) - if err != nil { - return nil, err - } - return TargetGraphToProto(ctx, graph, maxMessageBytes) -} diff --git a/internal/mapper/target_graph_test.go b/internal/mapper/target_graph_test.go index f3d1679d..c07cd9f7 100644 --- a/internal/mapper/target_graph_test.go +++ b/internal/mapper/target_graph_test.go @@ -143,7 +143,7 @@ func TestTargetGraphToProto_EmptyGraph(t *testing.T) { assert.True(t, ok) } -func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { +func TestTargetGraphToProto_Chunking(t *testing.T) { t.Parallel() numTargets := 50 @@ -157,7 +157,10 @@ func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { result.Targets[name] = &targethasher.Target{Name: name, Hash: []byte{0}, RuleType: "go_library"} } - baseline, err := ResultToGetTargetGraphResponse(context.Background(), result, 1<<30) + graph, err := ResultToTargetGraph(context.Background(), result) + require.NoError(t, err) + + baseline, err := TargetGraphToProto(context.Background(), graph, 1<<30) require.NoError(t, err) var targetBytes int for _, resp := range baseline { @@ -181,7 +184,7 @@ func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - responses, err := ResultToGetTargetGraphResponse(context.Background(), result, tt.maxMessageBytes) + responses, err := TargetGraphToProto(context.Background(), graph, tt.maxMessageBytes) require.NoError(t, err) var targetChunks, totalTargets int From eab64f1a4e5b5722dced1a88ac22b098b4e83263 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 16 Jul 2026 13:37:16 -0700 Subject: [PATCH 5/6] feat(entity): add ID-mapped target types and chunk/proto converters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add entity.IDTarget, entity.GraphMetadata, and entity.GraphChunk as compact ID-mapped representations for streaming and storage, mirroring the proto shape without proto dependencies. Add TargetGraphToChunks (entity.TargetGraph → []entity.GraphChunk), GraphChunkToProto/ProtoToGraphChunk converters. TargetGraphToProto now composes TargetGraphToChunks + GraphChunkToProto. Co-Authored-By: Claude Sonnet 5 --- entity/optimized_target.go | 32 +++++++ internal/mapper/target_graph.go | 149 +++++++++++++++++++++++++++----- 2 files changed, 160 insertions(+), 21 deletions(-) diff --git a/entity/optimized_target.go b/entity/optimized_target.go index f654b2da..c97e1de1 100644 --- a/entity/optimized_target.go +++ b/entity/optimized_target.go @@ -19,3 +19,35 @@ type OptimizedTarget struct { type TargetGraph struct { Targets []OptimizedTarget } + +// IDTarget 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 GraphMetadata maps. +type IDTarget struct { + ID int32 + Hash string + DirectDependencies []int32 + RuleType int32 + Tags []int32 + Root bool + External bool + Attributes map[int32]int32 +} + +// GraphMetadata holds the ID-to-string mappings that accompany a set of +// IDTarget entries. Consumers merge metadata across chunks before +// resolving IDs. +type GraphMetadata struct { + TargetIDMapping map[int32]string + RuleTypeMapping map[int32]string + TagMapping map[int32]string + AttributeNameMapping map[int32]string + AttributeStringValueMapping map[int32]string +} + +// GraphChunk is one piece of a streamed target graph — either a batch of +// ID-mapped targets or a metadata mapping. Exactly one field is non-nil. +type GraphChunk struct { + Targets []IDTarget + Metadata *GraphMetadata +} diff --git a/internal/mapper/target_graph.go b/internal/mapper/target_graph.go index bcb9a19f..96d33f4e 100644 --- a/internal/mapper/target_graph.go +++ b/internal/mapper/target_graph.go @@ -73,10 +73,11 @@ func ResultToTargetGraph(ctx context.Context, result targethasher.Result) (entit return entity.TargetGraph{Targets: targets}, nil } -// TargetGraphToProto converts an entity.TargetGraph into chunked -// GetTargetGraphResponse messages ready for streaming. Each message is -// bounded by maxMessageBytes of real wire size. -func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessageBytes int) ([]*tangopb.GetTargetGraphResponse, error) { +// TargetGraphToChunks converts a string-based entity.TargetGraph into +// ID-mapped entity.GraphChunk slices, ready for storage or proto +// conversion. Each chunk is bounded by maxMessageBytes of estimated wire +// size. +func TargetGraphToChunks(ctx context.Context, graph entity.TargetGraph, maxMessageBytes int) ([]entity.GraphChunk, error) { targetNamesMapping := make(map[string]int32, len(graph.Targets)) for i, t := range graph.Targets { targetNamesMapping[t.Name] = int32(i + 1) @@ -87,7 +88,8 @@ func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessag attrNameMapper := idmapper.NewMapper() attrStrValMapper := idmapper.NewMapper() - optimizedTargets := make([]*tangopb.OptimizedTarget, 0, len(graph.Targets)) + protoTargets := make([]*tangopb.OptimizedTarget, 0, len(graph.Targets)) + idTargets := make([]entity.IDTarget, 0, len(graph.Targets)) for i, t := range graph.Targets { if i%cancelCheckInterval == 0 { if err := ctx.Err(); err != nil { @@ -103,8 +105,9 @@ func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessag depIDs = append(depIDs, depID) } - ot := &tangopb.OptimizedTarget{ - Id: targetNamesMapping[t.Name], + id := targetNamesMapping[t.Name] + idt := entity.IDTarget{ + ID: id, Hash: t.Hash, DirectDependencies: depIDs, Root: t.Root, @@ -112,24 +115,25 @@ func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessag } if t.RuleType != "" { - ot.RuleType = ruleTypeMapper.ID(t.RuleType) + idt.RuleType = ruleTypeMapper.ID(t.RuleType) } 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 + idt.Tags = tagIDs } if len(t.Attributes) > 0 { attrs := make(map[int32]int32, len(t.Attributes)) for k, v := range t.Attributes { attrs[attrNameMapper.ID(k)] = attrStrValMapper.ID(v) } - ot.Attributes = attrs + idt.Attributes = attrs } - optimizedTargets = append(optimizedTargets, ot) + idTargets = append(idTargets, idt) + protoTargets = append(protoTargets, idTargetToProtoTarget(&idt)) } targetIDToName := make(map[int32]string, len(targetNamesMapping)) @@ -137,18 +141,22 @@ func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessag targetIDToName[id] = s } - var responses []*tangopb.GetTargetGraphResponse - targetChunks, err := BySize(optimizedTargets, maxMessageBytes) + // Use proto Size() for accurate chunking, then build entity chunks. + targetChunks, err := BySize(protoTargets, maxMessageBytes) if err != nil { return nil, err } - for _, chunk := range targetChunks { - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Targets{ - Targets: &tangopb.OptimizedTargets{Targets: chunk}, - }, - }) + + var chunks []entity.GraphChunk + idx := 0 + for _, pc := range targetChunks { + chunk := entity.GraphChunk{ + Targets: idTargets[idx : idx+len(pc)], + } + idx += len(pc) + chunks = append(chunks, chunk) } + metaChunks, err := ChunkMetadata( targetIDToName, ruleTypeMapper.Invert(), @@ -161,10 +169,109 @@ func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessag return nil, err } for _, meta := range metaChunks { - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Metadata{Metadata: meta}, + chunks = append(chunks, entity.GraphChunk{ + Metadata: protoMetadataToEntity(meta), }) } + return chunks, nil +} + +// TargetGraphToProto converts an entity.TargetGraph into chunked +// GetTargetGraphResponse messages ready for streaming. Each message is +// bounded by maxMessageBytes of real wire size. +func TargetGraphToProto(ctx context.Context, graph entity.TargetGraph, maxMessageBytes int) ([]*tangopb.GetTargetGraphResponse, error) { + chunks, err := TargetGraphToChunks(ctx, graph, maxMessageBytes) + if err != nil { + return nil, err + } + responses := make([]*tangopb.GetTargetGraphResponse, 0, len(chunks)) + for i := range chunks { + responses = append(responses, GraphChunkToProto(&chunks[i])) + } return responses, nil } + +// GraphChunkToProto converts an entity.GraphChunk to the corresponding +// proto GetTargetGraphResponse. +func GraphChunkToProto(chunk *entity.GraphChunk) *tangopb.GetTargetGraphResponse { + if chunk.Metadata != nil { + return &tangopb.GetTargetGraphResponse{ + Item: &tangopb.GetTargetGraphResponse_Metadata{ + Metadata: entityMetadataToProto(chunk.Metadata), + }, + } + } + targets := make([]*tangopb.OptimizedTarget, len(chunk.Targets)) + for i := range chunk.Targets { + targets[i] = idTargetToProtoTarget(&chunk.Targets[i]) + } + return &tangopb.GetTargetGraphResponse{ + Item: &tangopb.GetTargetGraphResponse_Targets{ + Targets: &tangopb.OptimizedTargets{Targets: targets}, + }, + } +} + +// ProtoToGraphChunk converts a proto GetTargetGraphResponse to an +// entity.GraphChunk. +func ProtoToGraphChunk(resp *tangopb.GetTargetGraphResponse) entity.GraphChunk { + switch item := resp.GetItem().(type) { + case *tangopb.GetTargetGraphResponse_Targets: + targets := make([]entity.IDTarget, len(item.Targets.GetTargets())) + for i, t := range item.Targets.GetTargets() { + targets[i] = protoTargetToIDTarget(t) + } + return entity.GraphChunk{Targets: targets} + case *tangopb.GetTargetGraphResponse_Metadata: + return entity.GraphChunk{Metadata: protoMetadataToEntity(item.Metadata)} + default: + return entity.GraphChunk{} + } +} + +func idTargetToProtoTarget(t *entity.IDTarget) *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, + } +} + +func protoTargetToIDTarget(t *tangopb.OptimizedTarget) entity.IDTarget { + return entity.IDTarget{ + ID: t.GetId(), + Hash: t.GetHash(), + DirectDependencies: t.GetDirectDependencies(), + RuleType: t.GetRuleType(), + Tags: t.GetTags(), + Root: t.GetRoot(), + External: t.GetExternal(), + Attributes: t.GetAttributes(), + } +} + +func entityMetadataToProto(m *entity.GraphMetadata) *tangopb.Metadata { + return &tangopb.Metadata{ + TargetIdMapping: m.TargetIDMapping, + RuleTypeMapping: m.RuleTypeMapping, + TagMapping: m.TagMapping, + AttributeNameMapping: m.AttributeNameMapping, + AttributeStringValueMapping: m.AttributeStringValueMapping, + } +} + +func protoMetadataToEntity(m *tangopb.Metadata) *entity.GraphMetadata { + return &entity.GraphMetadata{ + TargetIDMapping: m.GetTargetIdMapping(), + RuleTypeMapping: m.GetRuleTypeMapping(), + TagMapping: m.GetTagMapping(), + AttributeNameMapping: m.GetAttributeNameMapping(), + AttributeStringValueMapping: m.GetAttributeStringValueMapping(), + } +} From ce290fd39604e298152b22bb276800cdb04bd80b Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 16 Jul 2026 15:09:01 -0700 Subject: [PATCH 6/6] refactor: move chunking to internal/streaming with renamed APIs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move streaming-split logic from internal/mapper/chunk.go to internal/streaming/streaming.go with clearer names: - BySize → streaming.SplitBySize (generic, proto-free) - ChunkMetadata → streaming.SplitMetadata (returns entity.GraphMetadata) The streaming package has no proto dependency — it works with the Sizer interface and entity types only. Proto conversion stays in the mapper. Co-Authored-By: Claude Sonnet 5 --- internal/mapper/BUILD.bazel | 3 +- internal/mapper/chunk.go | 147 --------------------------- internal/mapper/chunk_test.go | 128 ----------------------- internal/mapper/target_graph.go | 17 ++-- internal/streaming/BUILD.bazel | 20 ++++ internal/streaming/streaming.go | 129 +++++++++++++++++++++++ internal/streaming/streaming_test.go | 114 +++++++++++++++++++++ 7 files changed, 273 insertions(+), 285 deletions(-) delete mode 100644 internal/mapper/chunk.go delete mode 100644 internal/mapper/chunk_test.go 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/mapper/BUILD.bazel b/internal/mapper/BUILD.bazel index c4877cee..4a3ec64a 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", "target_graph.go", ], importpath = "github.com/uber/tango/internal/mapper", @@ -13,6 +12,7 @@ go_library( "//core/targethasher", "//entity", "//internal/mapper/idmapper", + "//internal/streaming", "//tangopb", "@com_github_bazelbuild_buildtools//build_proto", ], @@ -22,7 +22,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/chunk.go b/internal/mapper/chunk.go deleted file mode 100644 index c2cf073b..00000000 --- a/internal/mapper/chunk.go +++ /dev/null @@ -1,147 +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 mapper - -import ( - "fmt" - - "github.com/uber/tango/tangopb" -) - -// sizer is satisfied by generated proto message types, which all expose a -// Size() method computing their exact serialized byte length with no -// allocation or marshaling. -type sizer interface { - Size() int -} - -// BySize splits items into consecutive runs whose cumulative Size() stays -// at or under maxMessageBytes. A single item larger than the budget ships -// alone since it can't be split further. Always returns at least one chunk: -// empty input yields a single empty chunk so callers always have a message -// to send on the stream. Returns an error if a multi-chunk split produces -// any empty chunk after the first (which would indicate a bug in the -// splitting logic). -func BySize[T sizer](items []T, maxMessageBytes int) ([][]T, error) { - if len(items) == 0 { - return [][]T{nil}, nil - } - chunks := make([][]T, 0, 1) - var current []T - currentBytes := 0 - for _, item := range items { - itemBytes := item.Size() - if len(current) > 0 && currentBytes+itemBytes > maxMessageBytes { - chunks = append(chunks, current) - current = nil - currentBytes = 0 - } - current = append(current, item) - currentBytes += itemBytes - } - chunks = append(chunks, current) - for i := 1; i < len(chunks); i++ { - if len(chunks[i]) == 0 { - return nil, fmt.Errorf("internal error: chunk %d of %d is empty", i, len(chunks)) - } - } - return chunks, nil -} - -// ChunkMetadata splits the metadata maps into multiple Metadata messages so -// each stays at or under maxMessageBytes. The two large maps (target names, -// attribute string values) are split independently by measured entry wire -// size; consumers merge all Metadata messages before use. The small maps -// (rule_type, tag, attribute_name) are sent in the first chunk. -// Always returns at least one chunk so callers always have a message to -// send on the stream. Returns an error if a non-first chunk is completely -// empty (all maps nil/empty). -func ChunkMetadata( - targetIDToName map[int32]string, - ruleTypeIDToName map[int32]string, - tagIDToName map[int32]string, - attrNameIDToName map[int32]string, - attrStrValIDToVal map[int32]string, - maxMessageBytes int, -) ([]*tangopb.Metadata, error) { - targetChunks := splitMapByBytes(targetIDToName, maxMessageBytes) - attrValChunks := splitMapByBytes(attrStrValIDToVal, maxMessageBytes) - - chunks := make([]*tangopb.Metadata, 0, max(1, len(targetChunks)+len(attrValChunks))) - for _, c := range targetChunks { - chunks = append(chunks, &tangopb.Metadata{TargetIdMapping: c}) - } - for _, c := range attrValChunks { - chunks = append(chunks, &tangopb.Metadata{AttributeStringValueMapping: c}) - } - if len(chunks) == 0 { - chunks = append(chunks, &tangopb.Metadata{}) - } - chunks[0].RuleTypeMapping = ruleTypeIDToName - chunks[0].TagMapping = tagIDToName - chunks[0].AttributeNameMapping = attrNameIDToName - - for i := 1; i < len(chunks); i++ { - c := chunks[i] - if len(c.TargetIdMapping) == 0 && - len(c.RuleTypeMapping) == 0 && - len(c.TagMapping) == 0 && - len(c.AttributeNameMapping) == 0 && - len(c.AttributeStringValueMapping) == 0 { - return nil, fmt.Errorf("internal error: metadata chunk %d of %d is empty", i, len(chunks)) - } - } - - return chunks, nil -} - -func splitMapByBytes(m map[int32]string, maxBytes int) []map[int32]string { - if len(m) == 0 { - return nil - } - var chunks []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 { - chunks = append(chunks, current) - current = make(map[int32]string) - currentBytes = 0 - } - current[k] = v - currentBytes += entryBytes - } - if len(current) > 0 { - chunks = append(chunks, current) - } - return chunks -} - -// mapEntryWireSize returns the serialized size of one map[int32]string entry, -// mirroring the generated Metadata.Size() computation exactly. -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 -} diff --git a/internal/mapper/chunk_test.go b/internal/mapper/chunk_test.go deleted file mode 100644 index c6fc1f11..00000000 --- a/internal/mapper/chunk_test.go +++ /dev/null @@ -1,128 +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 mapper - -import ( - "strings" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - pb "github.com/uber/tango/tangopb" -) - -func TestBySize(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 - - chunks, err := BySize(targets, maxBytes) - require.NoError(t, err) - require.Len(t, chunks, 3) - assert.Len(t, chunks[0], 10) - assert.Len(t, chunks[1], 10) - assert.Len(t, chunks[2], 5) - - var total int - for _, c := range chunks { - for _, target := range c { - assert.Equal(t, int32(total+1), target.Id) - total++ - } - } - assert.Equal(t, 25, total) -} - -func TestBySize_SingleOversizedItemShipsAlone(t *testing.T) { - t.Parallel() - - oversized := &pb.OptimizedTarget{Id: 1, Hash: strings.Repeat("a", 1000)} - small := &pb.OptimizedTarget{Id: 2} - maxBytes := small.Size() - - chunks, err := BySize([]*pb.OptimizedTarget{oversized, small}, maxBytes) - require.NoError(t, err) - require.Len(t, chunks, 2) - assert.Equal(t, []*pb.OptimizedTarget{oversized}, chunks[0]) - assert.Equal(t, []*pb.OptimizedTarget{small}, chunks[1]) -} - -func TestBySize_EmptyInputReturnsOneEmptyChunk(t *testing.T) { - t.Parallel() - - chunks, err := BySize([]*pb.OptimizedTarget{}, 100) - require.NoError(t, err) - require.Len(t, chunks, 1) - assert.Empty(t, chunks[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 TestChunkMetadata_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") - - chunks, err := ChunkMetadata(targetMap, ruleType, nil, nil, nil, entryBytes*2) - require.NoError(t, err) - require.Len(t, chunks, 2) - - merged := map[int32]string{} - for i, meta := range chunks { - 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 TestChunkMetadata_AllEmptyMapsReturnsOneEmptyChunk(t *testing.T) { - t.Parallel() - - chunks, err := ChunkMetadata(nil, nil, nil, nil, nil, 100) - require.NoError(t, err) - require.Len(t, chunks, 1) - assert.Empty(t, chunks[0].TargetIdMapping) - assert.Empty(t, chunks[0].RuleTypeMapping) -} - -func TestChunkMetadata_EmptyBigMapsStillCarrySmallMaps(t *testing.T) { - t.Parallel() - - ruleType := map[int32]string{1: "go_library"} - chunks, err := ChunkMetadata(nil, ruleType, nil, nil, nil, 100) - require.NoError(t, err) - require.Len(t, chunks, 1) - assert.Equal(t, ruleType, chunks[0].RuleTypeMapping) -} diff --git a/internal/mapper/target_graph.go b/internal/mapper/target_graph.go index 96d33f4e..8ca4d459 100644 --- a/internal/mapper/target_graph.go +++ b/internal/mapper/target_graph.go @@ -9,6 +9,7 @@ import ( "github.com/uber/tango/core/targethasher" "github.com/uber/tango/entity" "github.com/uber/tango/internal/mapper/idmapper" + "github.com/uber/tango/internal/streaming" "github.com/uber/tango/tangopb" ) @@ -141,23 +142,23 @@ func TargetGraphToChunks(ctx context.Context, graph entity.TargetGraph, maxMessa targetIDToName[id] = s } - // Use proto Size() for accurate chunking, then build entity chunks. - targetChunks, err := BySize(protoTargets, maxMessageBytes) + // Use proto Size() for accurate splitting, then build entity chunks. + targetGroups, err := streaming.SplitBySize(protoTargets, maxMessageBytes) if err != nil { return nil, err } var chunks []entity.GraphChunk idx := 0 - for _, pc := range targetChunks { + for _, g := range targetGroups { chunk := entity.GraphChunk{ - Targets: idTargets[idx : idx+len(pc)], + Targets: idTargets[idx : idx+len(g)], } - idx += len(pc) + idx += len(g) chunks = append(chunks, chunk) } - metaChunks, err := ChunkMetadata( + metaGroups, err := streaming.SplitMetadata( targetIDToName, ruleTypeMapper.Invert(), tagMapper.Invert(), @@ -168,9 +169,9 @@ func TargetGraphToChunks(ctx context.Context, graph entity.TargetGraph, maxMessa if err != nil { return nil, err } - for _, meta := range metaChunks { + for _, meta := range metaGroups { chunks = append(chunks, entity.GraphChunk{ - Metadata: protoMetadataToEntity(meta), + Metadata: meta, }) } diff --git a/internal/streaming/BUILD.bazel b/internal/streaming/BUILD.bazel new file mode 100644 index 00000000..f6014da8 --- /dev/null +++ b/internal/streaming/BUILD.bazel @@ -0,0 +1,20 @@ +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"], +) + +go_test( + name = "streaming_test", + srcs = ["streaming_test.go"], + embed = [":streaming"], + deps = [ + "//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..1f4072e9 --- /dev/null +++ b/internal/streaming/streaming.go @@ -0,0 +1,129 @@ +package streaming + +import ( + "fmt" + + "github.com/uber/tango/entity" +) + +// 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 GraphMetadata +// 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 messages 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.GraphMetadata, error) { + targetGroups := splitMapByBytes(targetIDToName, maxBytes) + attrValGroups := splitMapByBytes(attrStrValIDToVal, maxBytes) + + metas := make([]*entity.GraphMetadata, 0, max(1, len(targetGroups)+len(attrValGroups))) + for _, g := range targetGroups { + metas = append(metas, &entity.GraphMetadata{TargetIDMapping: g}) + } + for _, g := range attrValGroups { + metas = append(metas, &entity.GraphMetadata{AttributeStringValueMapping: g}) + } + if len(metas) == 0 { + metas = append(metas, &entity.GraphMetadata{}) + } + 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 +} + +// MapEntryWireSize returns the serialized size of one map[int32]string +// entry, mirroring the protobuf Metadata.Size() computation exactly. +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 +} diff --git a/internal/streaming/streaming_test.go b/internal/streaming/streaming_test.go new file mode 100644 index 00000000..f4335cee --- /dev/null +++ b/internal/streaming/streaming_test.go @@ -0,0 +1,114 @@ +package streaming + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + 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) +}