From 9789ce128e991a14e1ed23eae328a38af5947c0c Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Sun, 19 Jul 2026 21:52:01 -0700 Subject: [PATCH 1/4] feat(entity): add entity types for target graph and changed targets Co-Authored-By: Claude Sonnet 5 --- entity/BUILD.bazel | 2 ++ entity/changed_targets.go | 17 +++++++++++++++++ entity/optimized_target.go | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+) create mode 100644 entity/changed_targets.go create mode 100644 entity/optimized_target.go diff --git a/entity/BUILD.bazel b/entity/BUILD.bazel index deff8984..58de4a19 100644 --- a/entity/BUILD.bazel +++ b/entity/BUILD.bazel @@ -4,7 +4,9 @@ go_library( name = "entity", srcs = [ "build_description.go", + "changed_targets.go", "computation_strategy.go", + "optimized_target.go", "target_graph.go", ], importpath = "github.com/uber/tango/entity", diff --git a/entity/changed_targets.go b/entity/changed_targets.go new file mode 100644 index 00000000..cfe19317 --- /dev/null +++ b/entity/changed_targets.go @@ -0,0 +1,17 @@ +package entity + +// ChangedTarget represents a target that differs between two revisions. +type ChangedTarget struct { + ChangeType int32 `json:"change_type"` + OldTarget *OptimizedTarget `json:"old_target,omitempty"` + NewTarget *OptimizedTarget `json:"new_target,omitempty"` + Distance int32 `json:"distance"` +} + +// GetChangedTargetsResponse is one piece of a streamed changed-targets +// result — either a batch of changed targets or a metadata mapping. +// Exactly one field is non-nil. +type GetChangedTargetsResponse struct { + ChangedTargets []ChangedTarget `json:"changed_targets"` + Metadata *Metadata `json:"metadata,omitempty"` +} diff --git a/entity/optimized_target.go b/entity/optimized_target.go new file mode 100644 index 00000000..36308ff9 --- /dev/null +++ b/entity/optimized_target.go @@ -0,0 +1,34 @@ +package entity + +// 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. +type OptimizedTarget struct { + ID int32 `json:"id"` + Hash string `json:"hash"` + DirectDependencies []int32 `json:"direct_dependencies"` + RuleType int32 `json:"rule_type"` + Tags []int32 `json:"tags"` + Root bool `json:"root"` + External bool `json:"external"` + Attributes map[int32]int32 `json:"attributes"` +} + +// Metadata holds the ID-to-string mappings that accompany a set of +// OptimizedTarget entries. Consumers merge metadata across chunks before +// resolving IDs. +type Metadata struct { + TargetIDMapping map[int32]string `json:"target_id_mapping"` + RuleTypeMapping map[int32]string `json:"rule_type_mapping"` + TagMapping map[int32]string `json:"tag_mapping"` + AttributeNameMapping map[int32]string `json:"attribute_name_mapping"` + AttributeStringValueMapping map[int32]string `json:"attribute_string_value_mapping"` +} + +// GetTargetGraphResponse 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 GetTargetGraphResponse struct { + Targets []OptimizedTarget `json:"targets"` + Metadata *Metadata `json:"metadata,omitempty"` +} From 7f55e8537a290dcaa99b8d488634da23d0fa7198 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Sun, 19 Jul 2026 21:52:05 -0700 Subject: [PATCH 2/4] feat(mapper): add entity-proto mappers and streaming splitters Co-Authored-By: Claude Sonnet 5 --- internal/mapper/BUILD.bazel | 6 + internal/mapper/changed_targets.go | 40 ++++++ internal/mapper/target_graph.go | 200 +++++++++++++++++++++++++++ internal/mapper/target_graph_test.go | 111 +++++++++++++++ internal/streaming/BUILD.bazel | 20 +++ internal/streaming/streaming.go | 127 +++++++++++++++++ internal/streaming/streaming_test.go | 114 +++++++++++++++ 7 files changed, 618 insertions(+) create mode 100644 internal/mapper/changed_targets.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 77903eb8..caa6c1c4 100644 --- a/internal/mapper/BUILD.bazel +++ b/internal/mapper/BUILD.bazel @@ -4,13 +4,18 @@ go_library( name = "mapper", srcs = [ "build_description.go", + "changed_targets.go", "target_graph.go", ], importpath = "github.com/uber/tango/internal/mapper", visibility = ["//visibility:public"], deps = [ + "//core/targethasher", "//entity", + "//internal/mapper/idmapper", + "//internal/streaming", "//tangopb", + "@com_github_bazelbuild_buildtools//build_proto", ], ) @@ -22,6 +27,7 @@ go_test( ], embed = [":mapper"], deps = [ + "//core/targethasher", "//entity", "//tangopb", "@com_github_stretchr_testify//assert", diff --git a/internal/mapper/changed_targets.go b/internal/mapper/changed_targets.go new file mode 100644 index 00000000..8cfac69a --- /dev/null +++ b/internal/mapper/changed_targets.go @@ -0,0 +1,40 @@ +package mapper + +import ( + "github.com/uber/tango/entity" + "github.com/uber/tango/tangopb" +) + +// ChangedTargetsResponseToProto converts an entity.GetChangedTargetsResponse +// to its proto equivalent for gRPC streaming. +func ChangedTargetsResponseToProto(resp *entity.GetChangedTargetsResponse) *tangopb.GetChangedTargetsResponse { + if resp.Metadata != nil { + return &tangopb.GetChangedTargetsResponse{ + Item: &tangopb.GetChangedTargetsResponse_Metadata{ + Metadata: metadataToProto(resp.Metadata), + }, + } + } + changed := make([]*tangopb.ChangedTarget, len(resp.ChangedTargets)) + for i := range resp.ChangedTargets { + ct := &resp.ChangedTargets[i] + changed[i] = &tangopb.ChangedTarget{ + ChangeType: tangopb.ChangeType(ct.ChangeType), + OldTarget: optionalTargetToProto(ct.OldTarget), + NewTarget: optionalTargetToProto(ct.NewTarget), + Distance: ct.Distance, + } + } + return &tangopb.GetChangedTargetsResponse{ + Item: &tangopb.GetChangedTargetsResponse_ChangedTargets{ + ChangedTargets: &tangopb.ChangedTargets{ChangedTargets: changed}, + }, + } +} + +func optionalTargetToProto(t *entity.OptimizedTarget) *tangopb.OptimizedTarget { + if t == nil { + return nil + } + return optimizedTargetToProto(t) +} diff --git a/internal/mapper/target_graph.go b/internal/mapper/target_graph.go index 45717d82..6b621bcb 100644 --- a/internal/mapper/target_graph.go +++ b/internal/mapper/target_graph.go @@ -1,12 +1,20 @@ 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/internal/streaming" "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 +32,195 @@ func ProtoToGetTargetGraphRequest(req *tangopb.GetTargetGraphRequest) (entity.Ge BypassCache: req.GetBypassCache(), }, nil } + +// ResultToTargetGraph converts a targethasher.Result into ID-mapped entity +// types. Returns a flat list of targets and their accompanying metadata. +// No chunking or proto conversion is performed. +func ResultToTargetGraph(ctx context.Context, result targethasher.Result) ([]entity.OptimizedTarget, *entity.Metadata, 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() + + targets := make([]entity.OptimizedTarget, 0, len(result.TargetNames)) + + n := 0 + for _, name := range result.TargetNames { + t, ok := result.Targets[name] + if !ok { + continue + } + if n%cancelCheckInterval == 0 { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + } + n++ + + depIDs := make([]int32, 0, len(t.Deps)) + for _, depName := range t.Deps { + if depID, ok := targetNamesMapping[depName]; ok { + depIDs = append(depIDs, depID) + } + } + + idt := entity.OptimizedTarget{ + ID: targetNamesMapping[name], + Hash: hex.EncodeToString(t.Hash), + DirectDependencies: depIDs, + Root: t.Root, + External: t.External, + } + if 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)) + } + idt.Tags = tagIDs + } + 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 { + attrs[attrNameMapper.ID(*attr.Name)] = attrStrValMapper.ID(*attr.StringValue) + } + } + if len(attrs) > 0 { + idt.Attributes = attrs + } + } + + targets = append(targets, idt) + } + + targetIDToName := make(map[int32]string, len(targetNamesMapping)) + for s, id := range targetNamesMapping { + targetIDToName[id] = s + } + + meta := &entity.Metadata{ + TargetIDMapping: targetIDToName, + RuleTypeMapping: ruleTypeMapper.Invert(), + TagMapping: tagMapper.Invert(), + AttributeNameMapping: attrNameMapper.Invert(), + AttributeStringValueMapping: attrStrValMapper.Invert(), + } + + return targets, meta, nil +} + +// ChunkTargetGraph splits targets and metadata into wire-safe +// entity.GetTargetGraphResponse chunks bounded by maxMessageBytes. +func ChunkTargetGraph(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 := streaming.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 := streaming.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 +} + +// GetTargetGraphResponseToProto converts an entity.GetTargetGraphResponse to +// the corresponding proto GetTargetGraphResponse. +func GetTargetGraphResponseToProto(chunk *entity.GetTargetGraphResponse) *tangopb.GetTargetGraphResponse { + if chunk.Metadata != nil { + return &tangopb.GetTargetGraphResponse{ + Item: &tangopb.GetTargetGraphResponse_Metadata{ + Metadata: metadataToProto(chunk.Metadata), + }, + } + } + targets := make([]*tangopb.OptimizedTarget, len(chunk.Targets)) + for i := range chunk.Targets { + targets[i] = optimizedTargetToProto(&chunk.Targets[i]) + } + return &tangopb.GetTargetGraphResponse{ + Item: &tangopb.GetTargetGraphResponse_Targets{ + Targets: &tangopb.OptimizedTargets{Targets: targets}, + }, + } +} + +// metadataToProto converts an entity.Metadata to a proto Metadata. +func metadataToProto(m *entity.Metadata) *tangopb.Metadata { + return &tangopb.Metadata{ + TargetIdMapping: m.TargetIDMapping, + RuleTypeMapping: m.RuleTypeMapping, + TagMapping: m.TagMapping, + AttributeNameMapping: m.AttributeNameMapping, + AttributeStringValueMapping: m.AttributeStringValueMapping, + } +} + +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, + } +} + +func protoToOptimizedTarget(t *tangopb.OptimizedTarget) entity.OptimizedTarget { + return entity.OptimizedTarget{ + 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 protoToMetadata(m *tangopb.Metadata) *entity.Metadata { + return &entity.Metadata{ + TargetIDMapping: m.GetTargetIdMapping(), + RuleTypeMapping: m.GetRuleTypeMapping(), + TagMapping: m.GetTagMapping(), + AttributeNameMapping: m.GetAttributeNameMapping(), + AttributeStringValueMapping: m.GetAttributeStringValueMapping(), + } +} diff --git a/internal/mapper/target_graph_test.go b/internal/mapper/target_graph_test.go index 0d0d92fb..271bf7f6 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,111 @@ func TestProtoToGetTargetGraphRequest(t *testing.T) { }) } } + +func TestResultToTargetGraph_EmptyResult(t *testing.T) { + t.Parallel() + + targets, meta, err := ResultToTargetGraph(context.Background(), targethasher.Result{}) + require.NoError(t, err) + assert.Empty(t, targets) + assert.NotNil(t, meta) +} + +func TestChunkTargetGraph(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"} + } + + targets, meta, err := ResultToTargetGraph(context.Background(), result) + require.NoError(t, err) + + // Get baseline proto size per target for budget calculation. + baseline, err := ChunkTargetGraph(targets, meta, 1<<30) + require.NoError(t, err) + protoResp := GetTargetGraphResponseToProto(&baseline[0]) + protoTargets := protoResp.GetItem().(*tangopb.GetTargetGraphResponse_Targets) + targetBytes := protoTargets.Targets.Targets[0].Size() + 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() + chunks, err := ChunkTargetGraph(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) + }) + } +} + +func TestGetTargetGraphResponseToProto_RoundTrip(t *testing.T) { + t.Parallel() + + chunk := entity.GetTargetGraphResponse{ + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "ab", DirectDependencies: []int32{2}, RuleType: 10, Root: true}, + {ID: 2, Hash: "cd", Tags: []int32{5, 6}, External: true}, + }, + } + + proto := GetTargetGraphResponseToProto(&chunk) + roundTripped := protoToEntityResponse(proto) + assert.Equal(t, chunk, roundTripped) +} + +func TestGetTargetGraphResponseToProto_Metadata_RoundTrip(t *testing.T) { + t.Parallel() + + chunk := entity.GetTargetGraphResponse{ + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//pkg:a", 2: "//pkg:b"}, + RuleTypeMapping: map[int32]string{10: "go_library"}, + }, + } + + proto := GetTargetGraphResponseToProto(&chunk) + roundTripped := protoToEntityResponse(proto) + assert.Equal(t, chunk, roundTripped) +} + +func protoToEntityResponse(resp *tangopb.GetTargetGraphResponse) entity.GetTargetGraphResponse { + switch item := resp.GetItem().(type) { + case *tangopb.GetTargetGraphResponse_Targets: + targets := make([]entity.OptimizedTarget, len(item.Targets.GetTargets())) + for i, t := range item.Targets.GetTargets() { + targets[i] = protoToOptimizedTarget(t) + } + return entity.GetTargetGraphResponse{Targets: targets} + case *tangopb.GetTargetGraphResponse_Metadata: + return entity.GetTargetGraphResponse{Metadata: protoToMetadata(item.Metadata)} + default: + return entity.GetTargetGraphResponse{} + } +} 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..abdfdee9 --- /dev/null +++ b/internal/streaming/streaming.go @@ -0,0 +1,127 @@ +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 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 +} diff --git a/internal/streaming/streaming_test.go b/internal/streaming/streaming_test.go new file mode 100644 index 00000000..65a31faf --- /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) +} From 86773114ab578ab28ed8654165d66cbf608244f8 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Sun, 19 Jul 2026 21:52:08 -0700 Subject: [PATCH 3/4] refactor(storage): replace proto serialization with JSON entity encoding Co-Authored-By: Claude Sonnet 5 --- config/config.go | 3 ++ config/service_config.go | 25 ++++-------- core/storage/BUILD.bazel | 6 +-- core/storage/changedtargetsreader.go | 18 +++----- core/storage/graphreader.go | 16 ++++---- core/storage/graphwriter.go | 48 ++++++++++------------ core/storage/reader.go | 52 +++++++++--------------- core/storage/storage_test.go | 17 +++----- core/storage/storagemock/BUILD.bazel | 2 +- core/storage/storagemock/readermock.go | 6 +-- orchestrator/BUILD.bazel | 3 +- orchestrator/native_orchestrator.go | 14 ++++--- orchestrator/native_orchestrator_test.go | 34 ++++++---------- orchestrator/orchestrator.go | 1 - orchestrator/testdata/config.yaml | 4 -- 15 files changed, 95 insertions(+), 154 deletions(-) diff --git a/config/config.go b/config/config.go index 77ca4939..a1382e71 100644 --- a/config/config.go +++ b/config/config.go @@ -70,6 +70,9 @@ func Parse(configFilePath string) (*Config, error) { if config.Service.WorkerPoolSize <= 0 { return nil, fmt.Errorf("service.worker_pool_size must be > 0, got %d", config.Service.WorkerPoolSize) } + if config.Service.MaxMessageBytes <= 0 { + config.Service.MaxMessageBytes = defaultMaxMessageBytes + } config.repositoryByRemote = make(map[string]*RepositoryConfig, len(config.Repository)) for i := range config.Repository { remote := config.Repository[i].Remote diff --git a/config/service_config.go b/config/service_config.go index 2245f35f..77c9dade 100644 --- a/config/service_config.go +++ b/config/service_config.go @@ -16,23 +16,12 @@ package config // ServiceConfig holds operational configuration for the Tango service. type ServiceConfig struct { - WorkerPoolSize int `yaml:"worker_pool_size"` // number of worker workspaces per repo - RepoManagerClonePath string `yaml:"repo_manager_clone_path"` // root directory for origin repo clones - WorkerRootPath string `yaml:"worker_root_path"` // root directory for worker workspace checkouts; defaults to repo_manager_clone_path/.workers - Chunking ChunkConfig `yaml:"chunking"` // streaming chunk sizes; zero values fall back to package defaults + WorkerPoolSize int `yaml:"worker_pool_size"` // number of worker workspaces per repo + RepoManagerClonePath string `yaml:"repo_manager_clone_path"` // root directory for origin repo clones + WorkerRootPath string `yaml:"worker_root_path"` // root directory for worker workspace checkouts; defaults to repo_manager_clone_path/.workers + MaxMessageBytes int `yaml:"max_message_bytes"` // max serialized bytes per streamed gRPC message; 0 → DefaultMaxMessageBytes } -// 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 -// the 64MB default gRPC per-message limit. -type ChunkConfig struct { - // TargetChunkSize is the max number of OptimizedTarget entries per stream message. - TargetChunkSize int `yaml:"target_chunk_size"` - // ChangedTargetChunkSize is the max number of ChangedTarget entries per stream message. - // ChangedTarget carries both old and new targets (~2× the size of a regular target). - ChangedTargetChunkSize int `yaml:"changed_target_chunk_size"` - // MetadataMapChunkSize is the max number of entries per metadata map chunk. - // Applies to target_id_mapping and attribute_string_value_mapping. - MetadataMapChunkSize int `yaml:"metadata_map_chunk_size"` -} +// defaultMaxMessageBytes is the fallback max serialized size per streamed +// message (~4.25 MB), well under the 64 MB default gRPC limit. +const defaultMaxMessageBytes = 4_250_000 diff --git a/core/storage/BUILD.bazel b/core/storage/BUILD.bazel index 46f395aa..fdde12fb 100644 --- a/core/storage/BUILD.bazel +++ b/core/storage/BUILD.bazel @@ -13,11 +13,7 @@ go_library( ], importpath = "github.com/uber/tango/core/storage", visibility = ["//visibility:public"], - deps = [ - "//tangopb", - "@com_github_gogo_protobuf//io", - "@com_github_gogo_protobuf//proto", - ], + deps = ["//entity"], ) go_test( diff --git a/core/storage/changedtargetsreader.go b/core/storage/changedtargetsreader.go index d490a139..e009257a 100644 --- a/core/storage/changedtargetsreader.go +++ b/core/storage/changedtargetsreader.go @@ -1,4 +1,4 @@ -// Copyright (c) 2026 Uber Technologies, Inc. +// 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. @@ -17,28 +17,20 @@ package storage import ( "context" - pb "github.com/uber/tango/tangopb" + "github.com/uber/tango/entity" ) -// ChangedTargetsReader reads GetChangedTargetsResponse messages from storage. +// ChangedTargetsReader reads entity.GetChangedTargetsResponse values from storage. type ChangedTargetsReader interface { - Read() (*pb.GetChangedTargetsResponse, error) + Read() (entity.GetChangedTargetsResponse, error) Close() error } // NewChangedTargetsReader returns a ChangedTargetsReader that reads from storage at key. func NewChangedTargetsReader(ctx context.Context, st Storage, key string) (ChangedTargetsReader, error) { - r, err := newReader[pb.GetChangedTargetsResponse](ctx, st, key, 32<<20, func(m *pb.GetChangedTargetsResponse) bool { - return m.GetItem() == nil - }) + r, err := newReader[entity.GetChangedTargetsResponse](ctx, st, key) if err != nil { return nil, err } return r, nil } - -// WriteChangedTargetsStream writes a list of GetChangedTargetsResponse messages to storage. -// The messages are written as length-delimited protobuf, allowing streaming reads. -func WriteChangedTargetsStream(ctx context.Context, st Storage, key string, responses []*pb.GetChangedTargetsResponse) error { - return writeStream[pb.GetChangedTargetsResponse](ctx, st, key, responses) -} diff --git a/core/storage/graphreader.go b/core/storage/graphreader.go index fbb28cc3..6765aaa0 100644 --- a/core/storage/graphreader.go +++ b/core/storage/graphreader.go @@ -17,22 +17,20 @@ package storage import ( "context" - pb "github.com/uber/tango/tangopb" + "github.com/uber/tango/entity" ) +// GraphReader streams entity.GetTargetGraphResponse values from a stored +// target graph. type GraphReader interface { - // Read reads the next GetTargetGraphResponse message from the storage. - Read() (*pb.GetTargetGraphResponse, error) - // Close releases any underlying resources if supported by the implementation. - // Implementations that don't hold resources may return nil. + Read() (entity.GetTargetGraphResponse, error) Close() error } -// NewGraphReader returns a GraphReader that, when read, will fetch the stored graph at key +// NewGraphReader opens the stored target graph at key and returns a +// GraphReader that yields entity.GetTargetGraphResponse values. func NewGraphReader(ctx context.Context, st Storage, key string) (GraphReader, error) { - r, err := newReader[pb.GetTargetGraphResponse](ctx, st, key, 512<<20, func(m *pb.GetTargetGraphResponse) bool { // 512MB/message limit - return m.GetItem() == nil - }) + r, err := newReader[entity.GetTargetGraphResponse](ctx, st, key) if err != nil { return nil, err } diff --git a/core/storage/graphwriter.go b/core/storage/graphwriter.go index 6bd0a711..c558bc34 100644 --- a/core/storage/graphwriter.go +++ b/core/storage/graphwriter.go @@ -16,36 +16,40 @@ package storage import ( "context" + "encoding/json" "fmt" "io" - gogio "github.com/gogo/protobuf/io" - pb "github.com/uber/tango/tangopb" + "github.com/uber/tango/entity" ) -// writeStream marshals msgs as length-delimited protobuf and streams them to -// storage under key. It uses an io.Pipe so the serialized payload is never -// buffered in full a second time: a writer goroutine encodes into the pipe -// while Put consumes from it. -// -// The writer goroutine checks ctx before each message so a cancellation -// unwinds the encode loop promptly instead of waiting for Put to notice and -// stop reading; the context error is propagated to the reader. If Put returns -// before draining the pipe, the reader is closed to unblock the writer. The -// goroutine is joined before returning, and its error is returned when Put -// succeeds. -func writeStream[T any, PT protoMessage[T]](ctx context.Context, st Storage, key string, msgs []PT) error { +// WriteGraphStream writes entity.GetTargetGraphResponse values to storage +// as newline-delimited JSON, allowing streaming reads via NewGraphReader. +func WriteGraphStream(ctx context.Context, st Storage, key string, chunks []entity.GetTargetGraphResponse) error { + return writeStream(ctx, st, key, chunks) +} + +// WriteChangedTargetsStream writes entity.GetChangedTargetsResponse values to storage +// as newline-delimited JSON, allowing streaming reads via NewChangedTargetsReader. +func WriteChangedTargetsStream(ctx context.Context, st Storage, key string, responses []entity.GetChangedTargetsResponse) error { + return writeStream(ctx, st, key, responses) +} + +// writeStream JSON-encodes values and streams them to storage under key. +// It uses an io.Pipe so the serialized payload is never buffered in full: +// a writer goroutine encodes into the pipe while Put consumes from it. +func writeStream[T any](ctx context.Context, st Storage, key string, values []T) error { pr, pw := io.Pipe() writerErr := make(chan error, 1) go func() { - w := gogio.NewDelimitedWriter(pw) // varint-length-delimited + enc := json.NewEncoder(pw) var err error - for _, m := range msgs { + for i := range values { if err = ctx.Err(); err != nil { break } - if err = w.WriteMsg(m); err != nil { - err = fmt.Errorf("write delimited: %w", err) + if err = enc.Encode(&values[i]); err != nil { + err = fmt.Errorf("encode value: %w", err) break } } @@ -53,7 +57,6 @@ func writeStream[T any, PT protoMessage[T]](ctx context.Context, st Storage, key writerErr <- err }() putErr := st.Put(ctx, UploadRequest{Key: key, Reader: pr}) - // Unblock the writer goroutine if Put stopped reading early. pr.CloseWithError(putErr) writeErr := <-writerErr if putErr != nil { @@ -61,10 +64,3 @@ func writeStream[T any, PT protoMessage[T]](ctx context.Context, st Storage, key } return writeErr } - -// WriteGraphStream writes a list of GetTargetGraphResponse messages to the storage. -// The messages are written as length-delimited protobuf, allowing streaming reads. -// Typically this includes multiple OptimizedTargets chunks followed by Metadata. -func WriteGraphStream(ctx context.Context, st Storage, key string, responses []*pb.GetTargetGraphResponse) error { - return writeStream[pb.GetTargetGraphResponse](ctx, st, key, responses) -} diff --git a/core/storage/reader.go b/core/storage/reader.go index 2546c7e5..9162856a 100644 --- a/core/storage/reader.go +++ b/core/storage/reader.go @@ -1,4 +1,4 @@ -// Copyright (c) 2026 Uber Technologies, Inc. +// 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. @@ -16,40 +16,28 @@ package storage import ( "context" + "encoding/json" "io" - - gogio "github.com/gogo/protobuf/io" - gogoproto "github.com/gogo/protobuf/proto" ) -// protoMessage is the constraint satisfied by generated gogoproto message -// pointer types used with reader. -type protoMessage[T any] interface { - *T - gogoproto.Message -} - -// reader streams length-delimited protobuf messages of type T from storage, -// treating a message for which isEmpty returns true as the stream terminator. -type reader[T any, PT protoMessage[T]] struct { - rc gogio.ReadCloser - isEmpty func(PT) bool +// reader streams JSON-encoded values of type T from storage. +type reader[T any] struct { + rc io.ReadCloser + dec *json.Decoder } -// Read reads the next message from the storage. -func (r *reader[T, PT]) Read() (PT, error) { - m := PT(new(T)) - if err := r.rc.ReadMsg(m); err != nil { - return nil, err - } - if r.isEmpty(m) { - return nil, io.EOF +// Read decodes the next value from the stream. Returns io.EOF at end of stream. +func (r *reader[T]) Read() (T, error) { + var v T + if err := r.dec.Decode(&v); err != nil { + var zero T + return zero, err } - return m, nil + return v, nil } -// Close releases any underlying resources. -func (r *reader[T, PT]) Close() error { +// Close releases the underlying reader. +func (r *reader[T]) Close() error { if r.rc != nil { return r.rc.Close() } @@ -57,14 +45,14 @@ func (r *reader[T, PT]) Close() error { } // newReader opens the blob at key and returns a reader that decodes -// length-delimited T messages from it, up to maxMessageSize bytes/message. -func newReader[T any, PT protoMessage[T]](ctx context.Context, st Storage, key string, maxMessageSize int, isEmpty func(PT) bool) (*reader[T, PT], error) { +// JSON-encoded T values from it. +func newReader[T any](ctx context.Context, st Storage, key string) (*reader[T], error) { resp, err := st.Get(ctx, DownloadRequest{Key: key}) if err != nil { return nil, err } - return &reader[T, PT]{ - rc: gogio.NewDelimitedReader(resp.ReadCloser, maxMessageSize), - isEmpty: isEmpty, + return &reader[T]{ + rc: resp.ReadCloser, + dec: json.NewDecoder(resp.ReadCloser), }, nil } diff --git a/core/storage/storage_test.go b/core/storage/storage_test.go index eb471a03..36551b98 100644 --- a/core/storage/storage_test.go +++ b/core/storage/storage_test.go @@ -24,15 +24,8 @@ import ( "github.com/stretchr/testify/require" ) -var errMarshal = errors.New("marshal failed") - -type marshalErrorMessage struct{} - -func (*marshalErrorMessage) Reset() {} -func (*marshalErrorMessage) String() string { return "" } -func (*marshalErrorMessage) ProtoMessage() {} -func (*marshalErrorMessage) Marshal() ([]byte, error) { - return nil, errMarshal +type marshalErrorValue struct { + Value chan int // channels cannot be JSON-marshaled } type discardStorage struct{} @@ -48,14 +41,14 @@ func (discardStorage) Exists(context.Context, string) (bool, error) { return fal func (discardStorage) List(context.Context, string) ([]string, error) { return nil, nil } func TestWriteStreamReturnsWriterError(t *testing.T) { - err := writeStream[marshalErrorMessage]( + err := writeStream( context.Background(), discardStorage{}, "key", - []*marshalErrorMessage{{}}, + []marshalErrorValue{{Value: make(chan int)}}, ) - require.ErrorIs(t, err, errMarshal) + require.Error(t, err) } func TestMemoryStorage_List(t *testing.T) { diff --git a/core/storage/storagemock/BUILD.bazel b/core/storage/storagemock/BUILD.bazel index 36d7778e..51978333 100644 --- a/core/storage/storagemock/BUILD.bazel +++ b/core/storage/storagemock/BUILD.bazel @@ -10,7 +10,7 @@ go_library( visibility = ["//visibility:public"], deps = [ "//core/storage", - "//tangopb", + "//entity", "@org_uber_go_mock//gomock", ], ) diff --git a/core/storage/storagemock/readermock.go b/core/storage/storagemock/readermock.go index 19ae18df..17931a0f 100644 --- a/core/storage/storagemock/readermock.go +++ b/core/storage/storagemock/readermock.go @@ -12,7 +12,7 @@ package storagemock import ( reflect "reflect" - tangopb "github.com/uber/tango/tangopb" + entity "github.com/uber/tango/entity" gomock "go.uber.org/mock/gomock" ) @@ -55,10 +55,10 @@ func (mr *MockGraphReaderMockRecorder) Close() *gomock.Call { } // Read mocks base method. -func (m *MockGraphReader) Read() (*tangopb.GetTargetGraphResponse, error) { +func (m *MockGraphReader) Read() (entity.GetTargetGraphResponse, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "Read") - ret0, _ := ret[0].(*tangopb.GetTargetGraphResponse) + ret0, _ := ret[0].(entity.GetTargetGraphResponse) ret1, _ := ret[1].(error) return ret0, ret1 } diff --git a/orchestrator/BUILD.bazel b/orchestrator/BUILD.bazel index fe456af3..3b620f43 100644 --- a/orchestrator/BUILD.bazel +++ b/orchestrator/BUILD.bazel @@ -20,6 +20,7 @@ go_library( "//entity", "//graphrunner", "//internal/cachekey", + "//internal/mapper", "@com_github_uber_go_tally//:tally", "@org_uber_go_zap//:zap", ], @@ -41,8 +42,6 @@ go_test( "//core/workspace/workspacemock", "//entity", "//graphrunner/mock", - "//tangopb", - "@com_github_gogo_protobuf//io", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", "@org_uber_go_mock//gomock", diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index 0ca9e389..3637c2e0 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -34,6 +34,7 @@ import ( "github.com/uber/tango/entity" "github.com/uber/tango/graphrunner" "github.com/uber/tango/internal/cachekey" + "github.com/uber/tango/internal/mapper" "go.uber.org/zap" ) @@ -209,14 +210,15 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - responses, err := common.ResultToGetTargetGraphResponse(ctx, result, - b.config.Service.Chunking.TargetChunkSize, - b.config.Service.Chunking.MetadataMapChunkSize, - ) + targets, meta, err := mapper.ResultToTargetGraph(ctx, result) if err != nil { - return nil, fmt.Errorf("convert target graph to response: %w", err) + return nil, fmt.Errorf("convert target graph: %w", err) } - err = storage.WriteGraphStream(ctx, b.storage, treehashPath, responses) + chunks, err := mapper.ChunkTargetGraph(targets, meta, b.config.Service.MaxMessageBytes) + if err != nil { + return nil, fmt.Errorf("chunk target graph: %w", err) + } + err = storage.WriteGraphStream(ctx, b.storage, treehashPath, chunks) if err != nil { return nil, fmt.Errorf("write graph to storage at %s: %w", treehashPath, err) } diff --git a/orchestrator/native_orchestrator_test.go b/orchestrator/native_orchestrator_test.go index e3650e4d..38e68671 100644 --- a/orchestrator/native_orchestrator_test.go +++ b/orchestrator/native_orchestrator_test.go @@ -21,7 +21,8 @@ import ( "io" "testing" - gogio "github.com/gogo/protobuf/io" + "encoding/json" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/uber/tango/config" @@ -34,7 +35,6 @@ import ( workspacemock "github.com/uber/tango/core/workspace/workspacemock" "github.com/uber/tango/entity" graphmock "github.com/uber/tango/graphrunner/mock" - pb "github.com/uber/tango/tangopb" "go.uber.org/mock/gomock" "go.uber.org/zap/zaptest" ) @@ -45,10 +45,7 @@ func TestNative_GetTargetGraph_Success(t *testing.T) { st := storagemock.NewMockStorage(ctrl) var buf bytes.Buffer - err := gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }) - require.NoError(t, err) + require.NoError(t, json.NewEncoder(&buf).Encode(entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}})) // Single fetch by remote/treehash for the graph st.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ ReadCloser: io.NopCloser(bytes.NewReader(buf.Bytes())), @@ -79,13 +76,11 @@ func TestNative_GetTargetGraph_Success(t *testing.T) { require.NoError(t, err) require.NotNil(t, reader) defer reader.Close() - graph, rerr := reader.Read() + chunk, rerr := reader.Read() require.NoError(t, rerr) - require.NotNil(t, graph) - assert.NotNil(t, graph.GetTargets()) - graph, rerr = reader.Read() + require.NotNil(t, chunk.Targets) + _, rerr = reader.Read() assert.Equal(t, io.EOF, rerr) - assert.Nil(t, graph) } func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) { @@ -101,9 +96,7 @@ func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) { }).MinTimes(2) // After compute, second read returns a valid delimited stream with one message var buf bytes.Buffer - _ = gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }) + _ = json.NewEncoder(&buf).Encode(entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}}) st.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ ReadCloser: io.NopCloser(bytes.NewReader(buf.Bytes())), }, nil) @@ -136,9 +129,9 @@ func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) { require.NoError(t, err) require.NotNil(t, reader) defer reader.Close() - graph, rerr := reader.Read() + chunk, rerr := reader.Read() require.NoError(t, rerr) - require.NotNil(t, graph) + require.NotNil(t, chunk.Targets) } func TestNative_GetTargetGraph_RevParseError_Propagates(t *testing.T) { @@ -172,10 +165,7 @@ func TestNative_GetTargetGraph_AppliesGitHubPR(t *testing.T) { defer ctrl.Finish() st := storagemock.NewMockStorage(ctrl) var buf bytes.Buffer - err := gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }) - require.NoError(t, err) + require.NoError(t, json.NewEncoder(&buf).Encode(entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}})) // git mock must handle Apply sequence from workspace.NewRequest for PR 123 g := gitmock.NewMockInterface(ctrl) @@ -210,9 +200,9 @@ func TestNative_GetTargetGraph_AppliesGitHubPR(t *testing.T) { require.NoError(t, err) require.NotNil(t, reader) defer reader.Close() - graph, rerr := reader.Read() + chunk, rerr := reader.Read() require.NoError(t, rerr) - require.NotNil(t, graph) + require.NotNil(t, chunk.Targets) } func TestNewNativeOrchestrator_usesProvidedConfig(t *testing.T) { diff --git a/orchestrator/orchestrator.go b/orchestrator/orchestrator.go index ace1b322..0490766e 100644 --- a/orchestrator/orchestrator.go +++ b/orchestrator/orchestrator.go @@ -18,7 +18,6 @@ import ( "context" "github.com/uber/tango/core/storage" - "github.com/uber/tango/entity" ) diff --git a/orchestrator/testdata/config.yaml b/orchestrator/testdata/config.yaml index 34887e08..f420f74f 100644 --- a/orchestrator/testdata/config.yaml +++ b/orchestrator/testdata/config.yaml @@ -11,7 +11,3 @@ repository: service: worker_pool_size: 3 - chunking: - target_chunk_size: 250 - changed_target_chunk_size: 125 - metadata_map_chunk_size: 50000 From 3969f6781242fe5b3bcc0a3911d2e8d2677862af Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Sun, 19 Jul 2026 21:52:11 -0700 Subject: [PATCH 4/4] refactor: migrate controller to entity types and remove old chunking code Co-Authored-By: Claude Sonnet 5 --- controller/BUILD.bazel | 3 +- controller/controller.go | 63 +- controller/getchangedtargets.go | 180 +++--- controller/getchangedtargets_test.go | 910 +++++++++++---------------- controller/gettargetgraph.go | 5 +- controller/gettargetgraph_test.go | 36 +- controller/testhelper_test.go | 28 +- core/common/BUILD.bazel | 13 +- core/common/utils.go | 239 ------- core/common/utils_test.go | 95 --- example/cmd/query-bench/BUILD.bazel | 2 +- example/cmd/query-bench/main.go | 22 +- 12 files changed, 533 insertions(+), 1063 deletions(-) diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 9ec9c14d..9368b763 100644 --- a/controller/BUILD.bazel +++ b/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", @@ -49,7 +49,6 @@ go_test( "//orchestrator/orchestratormock", "//tangopb", "//tangopb/tangopbmock", - "@com_github_gogo_protobuf//io", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", "@com_github_uber_go_tally//:tally", diff --git a/controller/controller.go b/controller/controller.go index 34ca33c2..cbcee0f2 100644 --- a/controller/controller.go +++ b/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" @@ -31,25 +29,25 @@ import ( // Params are the parameters for the controller. type Params struct { fx.In - Logger *zap.Logger - Storage storage.Storage - Orchestrator orchestrator.Orchestrator - Scope tally.Scope `optional:"true"` - ChunkConfig config.ChunkConfig `optional:"true"` + Logger *zap.Logger + Storage storage.Storage + Orchestrator orchestrator.Orchestrator + Scope tally.Scope `optional:"true"` + MaxMessageBytes int `optional:"true"` } // _totalDurationBuckets covers 0–15 minutes in 10-second linear intervals. var _totalDurationBuckets = tally.MustMakeLinearDurationBuckets(10*time.Second, 10*time.Second, 90) +const _defaultMaxMessageBytes = 4_250_000 + 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 +62,18 @@ func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { if scope == nil { scope = tally.NoopScope } - targetChunkSize := p.ChunkConfig.TargetChunkSize - if targetChunkSize <= 0 { - targetChunkSize = common.DefaultTargetChunkSize - } - changedTargetChunkSize := p.ChunkConfig.ChangedTargetChunkSize - if changedTargetChunkSize <= 0 { - changedTargetChunkSize = common.DefaultChangedTargetChunkSize - } - metadataMapChunkSize := p.ChunkConfig.MetadataMapChunkSize - if metadataMapChunkSize <= 0 { - metadataMapChunkSize = common.DefaultMetadataMapChunkSize + maxMessageBytes := p.MaxMessageBytes + if maxMessageBytes <= 0 { + maxMessageBytes = _defaultMaxMessageBytes } 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: maxMessageBytes, + totalDurationBuckets: _totalDurationBuckets, + appCtx: appCtx, } } @@ -98,12 +86,7 @@ func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { // The returned cancel function MUST be deferred; it releases the // context.AfterFunc handle so we do not leak a watcher past the request. func (c *controller) linkRequestCtx(reqCtx context.Context) (context.Context, context.CancelFunc) { - // Derive a per-request ctx whose cancel only affects this ctx and its - // children — it never propagates up to reqCtx. ctx, cancel := context.WithCancel(reqCtx) - // Register a one-shot watcher that cancels the derived ctx if appCtx fires. - // AfterFunc only observes appCtx; it never cancels it. stop() deregisters - // the watcher so the closure is not retained past the request. stop := context.AfterFunc(c.appCtx, cancel) return ctx, func() { stop() diff --git a/controller/getchangedtargets.go b/controller/getchangedtargets.go index 287db153..1ac16ce3 100644 --- a/controller/getchangedtargets.go +++ b/controller/getchangedtargets.go @@ -28,6 +28,7 @@ import ( "github.com/uber/tango/internal/cachekey" "github.com/uber/tango/internal/mapper" "github.com/uber/tango/internal/mapper/idmapper" + "github.com/uber/tango/internal/streaming" "github.com/uber/tango/internal/targetdiff" pb "github.com/uber/tango/tangopb" "go.uber.org/zap" @@ -35,12 +36,12 @@ import ( // job represents a single goroutine of getting a target graph type job struct { - graphStreamChunks []*pb.GetTargetGraphResponse - err error - cancelled bool - completed bool - ctx context.Context - cancel context.CancelFunc + graphChunks []entity.GetTargetGraphResponse + err error + cancelled bool + completed bool + ctx context.Context + cancel context.CancelFunc } // GetChangedTargets returns the changed targets between two revisions. If the @@ -161,7 +162,7 @@ func (c *controller) serveChangedTargetsFromCache(ctx context.Context, scope tal // Buffer all responses before sending any. A concurrent goroutine write may have // left a partial blob in storage; buffering lets us detect corruption and fall // through to recompute before we've sent anything to the client. - var cached []*pb.GetChangedTargetsResponse + var cached []entity.GetChangedTargetsResponse var readErr error for { if err := ctx.Err(); err != nil { @@ -169,7 +170,7 @@ func (c *controller) serveChangedTargetsFromCache(ctx context.Context, scope tal // Client gave up while we were draining the cache. Surface as a user-cancelled error. return false, common.WithReason(common.FailureReasonCancelled, common.ErrorTypeUser, err) } - var resp *pb.GetChangedTargetsResponse + var resp entity.GetChangedTargetsResponse resp, readErr = cachedReader.Read() if readErr == io.EOF { readErr = nil @@ -212,7 +213,7 @@ func (c *controller) serveChangedTargetsFromCache(ctx context.Context, scope tal // Errors caused solely by that induced cancellation are dropped; only the // original failure is returned. A client disconnect surfaces as a user-cancelled // error. -func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, logger *zap.Logger, request *pb.GetChangedTargetsRequest) ([]*pb.GetTargetGraphResponse, []*pb.GetTargetGraphResponse, error) { +func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, logger *zap.Logger, request *pb.GetChangedTargetsRequest) ([]entity.GetTargetGraphResponse, []entity.GetTargetGraphResponse, error) { jobs := make([]*job, 2) for i := 0; i < 2; i++ { // create independent contexts for each job; if one of the jobs fails, the other one should be cancelled to save resources and improve latency @@ -225,7 +226,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, l type graphResult struct { // order is 0 or 1, 0 is the base (first) revision, 1 is the target (second) revision order int - chunks []*pb.GetTargetGraphResponse + chunks []entity.GetTargetGraphResponse err error } results := make(chan graphResult, len(jobs)) @@ -262,8 +263,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, l } defer graphReader.Close() - // Read all chunks from the stream - var chunks []*pb.GetTargetGraphResponse + var chunks []entity.GetTargetGraphResponse for { chunk, err := graphReader.Read() if err == io.EOF { @@ -282,7 +282,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, l // Wait for both results to complete, either successfully or with an error. for range jobs { res := <-results - jobs[res.order].graphStreamChunks = res.chunks + jobs[res.order].graphChunks = res.chunks jobs[res.order].completed = true jobs[res.order].err = res.err if res.chunks == nil && res.err == nil { @@ -329,11 +329,11 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, l return nil, nil, err } - firstGraph := jobs[0].graphStreamChunks - secondGraph := jobs[1].graphStreamChunks + firstGraph := jobs[0].graphChunks + secondGraph := jobs[1].graphChunks // Drop job references so the GC can reclaim them once the comparison is done. - jobs[0].graphStreamChunks = nil - jobs[1].graphStreamChunks = nil + jobs[0].graphChunks = nil + jobs[1].graphChunks = nil return firstGraph, secondGraph, nil } @@ -341,7 +341,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, scope tally.Scope, l // a fire-and-forget goroutine so it does not block the stream send. The responses // is only read (never mutated) by the goroutine and the foreground send, so // concurrent access is safe; the caller must not mutate it. This is best effort. -func (c *controller) cacheComparedTargets(logger *zap.Logger, request *pb.GetChangedTargetsRequest, responses []*pb.GetChangedTargetsResponse) { +func (c *controller) cacheComparedTargets(logger *zap.Logger, request *pb.GetChangedTargetsRequest, responses []entity.GetChangedTargetsResponse) { go func() { // Use c.appCtx directly: the cache write is fire-and-forget and must // outlive the request (so a client disconnect doesn't abort it) but @@ -378,7 +378,7 @@ func (c *controller) cacheComparedTargets(logger *zap.Logger, request *pb.GetCha // are re-mapped into a canonical per-call ID namespace so the response metadata // only carries the names actually referenced. See internal/targetdiff for the // classification and distance rules. -func (c *controller) compareTargetGraphs(ctx context.Context, scope tally.Scope, logger *zap.Logger, firstGraph, secondGraph []*pb.GetTargetGraphResponse, maxDist int32) ([]*pb.GetChangedTargetsResponse, error) { +func (c *controller) compareTargetGraphs(ctx context.Context, scope tally.Scope, logger *zap.Logger, firstGraph, secondGraph []entity.GetTargetGraphResponse, maxDist int32) ([]entity.GetChangedTargetsResponse, error) { start := time.Now() compareScope := scope.SubScope("compare_target_graphs") logger.Info("compareTargetGraphs: Computing differences between target graphs") @@ -435,9 +435,9 @@ func (c *controller) compareTargetGraphs(ctx context.Context, scope tally.Scope, // only assign IDs to names they actually see, so the emitted metadata is // pruned to what the changed targets reference. mappers := newCanonicalMappers() - changed := make([]*pb.ChangedTarget, 0, len(result.ChangedTargets)) + changed := make([]entity.ChangedTarget, 0, len(result.ChangedTargets)) for _, ct := range result.ChangedTargets { - changed = append(changed, &pb.ChangedTarget{ + changed = append(changed, entity.ChangedTarget{ ChangeType: toChangeType(ct.ChangeType), OldTarget: mappers.transpose(ct.Before), NewTarget: mappers.transpose(ct.After), @@ -446,39 +446,36 @@ 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], - }, - }, - }) + // Entity ChangedTarget doesn't implement Sizer, so convert to proto for size measurement + // (same pattern as ChunkTargetGraph) and map groups back by index. + tempResp := mapper.ChangedTargetsResponseToProto(&entity.GetChangedTargetsResponse{ChangedTargets: changed}) + protoChanged := tempResp.GetItem().(*pb.GetChangedTargetsResponse_ChangedTargets).ChangedTargets.GetChangedTargets() + var results []entity.GetChangedTargetsResponse + changedGroups, err := streaming.SplitBySize(protoChanged, c.maxMessageBytes) + if err != nil { + return nil, err } - if len(results) == 0 { - results = append(results, &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{}, - }, + idx := 0 + for _, group := range changedGroups { + results = append(results, entity.GetChangedTargetsResponse{ + ChangedTargets: changed[idx : idx+len(group)], }) + idx += len(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, - ) { - results = append(results, &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_Metadata{ - Metadata: meta, - }, + c.maxMessageBytes, + ) + if err != nil { + return nil, err + } + for _, meta := range metaGroups { + results = append(results, entity.GetChangedTargetsResponse{ + Metadata: meta, }) } totalDuration := time.Since(start) @@ -498,10 +495,10 @@ const cancelCheckInterval = 4096 // getTargetsAndMetadata builds ID->target maps and merges metadata from a target graph stream. // Metadata may arrive in multiple chunks (e.g. when target_id_mapping exceeds the gRPC message // size limit); all chunks are merged into a single Metadata so callers can use it uniformly. -func getTargetsAndMetadata(ctx context.Context, graph []*pb.GetTargetGraphResponse) (map[int32]*pb.OptimizedTarget, *pb.Metadata, error) { - targets := make(map[int32]*pb.OptimizedTarget) - merged := &pb.Metadata{ - TargetIdMapping: make(map[int32]string), +func getTargetsAndMetadata(ctx context.Context, graph []entity.GetTargetGraphResponse) (map[int32]*entity.OptimizedTarget, *entity.Metadata, error) { + targets := make(map[int32]*entity.OptimizedTarget) + merged := &entity.Metadata{ + TargetIDMapping: make(map[int32]string), RuleTypeMapping: make(map[int32]string), TagMapping: make(map[int32]string), AttributeNameMapping: make(map[int32]string), @@ -511,26 +508,24 @@ func getTargetsAndMetadata(ctx context.Context, graph []*pb.GetTargetGraphRespon if ctx.Err() != nil { return nil, nil, ctx.Err() } - switch item := chunk.GetItem().(type) { - case *pb.GetTargetGraphResponse_Targets: - for _, t := range item.Targets.GetTargets() { - targets[t.GetId()] = t - } - case *pb.GetTargetGraphResponse_Metadata: - m := item.Metadata - for k, v := range m.GetTargetIdMapping() { - merged.TargetIdMapping[k] = v + for i := range chunk.Targets { + t := &chunk.Targets[i] + targets[t.ID] = t + } + if m := chunk.Metadata; m != nil { + for k, v := range m.TargetIDMapping { + merged.TargetIDMapping[k] = v } - for k, v := range m.GetRuleTypeMapping() { + for k, v := range m.RuleTypeMapping { merged.RuleTypeMapping[k] = v } - for k, v := range m.GetTagMapping() { + for k, v := range m.TagMapping { merged.TagMapping[k] = v } - for k, v := range m.GetAttributeNameMapping() { + for k, v := range m.AttributeNameMapping { merged.AttributeNameMapping[k] = v } - for k, v := range m.GetAttributeStringValueMapping() { + for k, v := range m.AttributeStringValueMapping { merged.AttributeStringValueMapping[k] = v } } @@ -541,12 +536,12 @@ func getTargetsAndMetadata(ctx context.Context, graph []*pb.GetTargetGraphRespon // toDiffGraph resolves a stream's int32 IDs into a semantic targetdiff.Graph // keyed by canonical target name. Targets with no name mapping are skipped; // dependency, tag, and attribute IDs that don't resolve are dropped. -func toDiffGraph(ctx context.Context, targetsByID map[int32]*pb.OptimizedTarget, meta *pb.Metadata) (targetdiff.Graph, error) { - targetIDMap := meta.GetTargetIdMapping() - ruleTypeMap := meta.GetRuleTypeMapping() - tagMap := meta.GetTagMapping() - attrNameMap := meta.GetAttributeNameMapping() - attrValMap := meta.GetAttributeStringValueMapping() +func toDiffGraph(ctx context.Context, targetsByID map[int32]*entity.OptimizedTarget, meta *entity.Metadata) (targetdiff.Graph, error) { + targetIDMap := meta.TargetIDMapping + ruleTypeMap := meta.RuleTypeMapping + tagMap := meta.TagMapping + attrNameMap := meta.AttributeNameMapping + attrValMap := meta.AttributeStringValueMapping graph := make(targetdiff.Graph, len(targetsByID)) i := 0 @@ -561,12 +556,12 @@ func toDiffGraph(ctx context.Context, targetsByID map[int32]*pb.OptimizedTarget, } target := &targetdiff.Target{ Name: name, - Hash: t.GetHash(), - RuleType: ruleTypeMap[t.GetRuleType()], - Root: t.GetRoot(), - External: t.GetExternal(), + Hash: t.Hash, + RuleType: ruleTypeMap[t.RuleType], + Root: t.Root, + External: t.External, } - if deps := t.GetDirectDependencies(); len(deps) > 0 { + if deps := t.DirectDependencies; len(deps) > 0 { target.Dependencies = make([]string, 0, len(deps)) for _, depID := range deps { if depName := targetIDMap[depID]; depName != "" { @@ -574,7 +569,7 @@ func toDiffGraph(ctx context.Context, targetsByID map[int32]*pb.OptimizedTarget, } } } - if tags := t.GetTags(); len(tags) > 0 { + if tags := t.Tags; len(tags) > 0 { target.Tags = make([]string, 0, len(tags)) for _, tagID := range tags { if tagName := tagMap[tagID]; tagName != "" { @@ -582,7 +577,7 @@ func toDiffGraph(ctx context.Context, targetsByID map[int32]*pb.OptimizedTarget, } } } - if attrs := t.GetAttributes(); len(attrs) > 0 { + if attrs := t.Attributes; len(attrs) > 0 { target.Attributes = make(map[string]string, len(attrs)) for nameID, valID := range attrs { if attrName := attrNameMap[nameID]; attrName != "" { @@ -618,14 +613,14 @@ func newCanonicalMappers() *canonicalMappers { } } -// transpose converts a semantic targetdiff.Target into a wire OptimizedTarget, +// transpose converts a semantic targetdiff.Target into an entity OptimizedTarget, // assigning canonical IDs to every name it references. Returns nil for a nil src. -func (m *canonicalMappers) transpose(src *targetdiff.Target) *pb.OptimizedTarget { +func (m *canonicalMappers) transpose(src *targetdiff.Target) *entity.OptimizedTarget { if src == nil { return nil } - dst := &pb.OptimizedTarget{ - Id: m.target.ID(src.Name), + dst := &entity.OptimizedTarget{ + ID: m.target.ID(src.Name), Hash: src.Hash, Root: src.Root, External: src.External, @@ -657,31 +652,32 @@ func (m *canonicalMappers) transpose(src *targetdiff.Target) *pb.OptimizedTarget return dst } -// toChangeType maps a targetdiff.ChangeType to its wire equivalent. -func toChangeType(ct targetdiff.ChangeType) pb.ChangeType { +// toChangeType maps a targetdiff.ChangeType to its wire equivalent as an int32. +func toChangeType(ct targetdiff.ChangeType) int32 { switch ct { case targetdiff.ChangeTypeNew: - return pb.CHANGE_TYPE_NEW + return int32(pb.CHANGE_TYPE_NEW) case targetdiff.ChangeTypeDeleted: - return pb.CHANGE_TYPE_DELETED + return int32(pb.CHANGE_TYPE_DELETED) case targetdiff.ChangeTypeChanged: - return pb.CHANGE_TYPE_CHANGED + return int32(pb.CHANGE_TYPE_CHANGED) default: - return pb.CHANGE_TYPE_INVALID + return int32(pb.CHANGE_TYPE_INVALID) } } // sendTrimmedChangedTargets streams responses to the client, filtering changed targets to those // within maxDist from any distance-0 seed when maxDist >= 0, stripping per-target // hash/tags/attributes per outputConfig's include_* flags, and pruning metadata mappings -// whose IDs are no longer referenced. Filtering and sending are combined into a single pass -// to avoid an intermediate allocation. -func sendTrimmedChangedTargets(stream pb.TangoServiceGetChangedTargetsYARPCServer, responses []*pb.GetChangedTargetsResponse, maxDist int32, outputConfig *pb.OutputConfig) error { +// whose IDs are no longer referenced. Each entity response is converted to proto at the +// stream.Send boundary. +func sendTrimmedChangedTargets(stream pb.TangoServiceGetChangedTargetsYARPCServer, responses []entity.GetChangedTargetsResponse, maxDist int32, outputConfig *pb.OutputConfig) error { stripFields := optimizedTargetNeedsStripping(outputConfig) pruneMeta := metadataNeedsPruning(outputConfig) - for _, resp := range responses { - toSend := resp - switch item := resp.GetItem().(type) { + for i := range responses { + protoResp := mapper.ChangedTargetsResponseToProto(&responses[i]) + toSend := protoResp + switch item := protoResp.GetItem().(type) { case *pb.GetChangedTargetsResponse_ChangedTargets: if maxDist >= 0 || stripFields { kept := item.ChangedTargets.GetChangedTargets() diff --git a/controller/getchangedtargets_test.go b/controller/getchangedtargets_test.go index 1c363e1f..3fcb1817 100644 --- a/controller/getchangedtargets_test.go +++ b/controller/getchangedtargets_test.go @@ -17,6 +17,7 @@ package controller import ( "bytes" "context" + "encoding/json" "errors" "fmt" "io" @@ -24,7 +25,6 @@ import ( "testing" "time" - gogio "github.com/gogo/protobuf/io" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/uber/tango/core/common" @@ -136,18 +136,10 @@ func TestValidateGetChangedTargetsRequest(t *testing.T) { func TestCompareTargetGraphs(t *testing.T) { c := newTestController(zap.NewNop()) - firstGraph := &pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{}, - }, - } - secondGraph := &pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{}, - }, - } + firstGraph := entity.GetTargetGraphResponse{Metadata: &entity.Metadata{}} + secondGraph := entity.GetTargetGraphResponse{Metadata: &entity.Metadata{}} - response, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), []*pb.GetTargetGraphResponse{firstGraph}, []*pb.GetTargetGraphResponse{secondGraph}, -1) + response, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), []entity.GetTargetGraphResponse{firstGraph}, []entity.GetTargetGraphResponse{secondGraph}, -1) require.NoError(t, err) require.NotNil(t, response) } @@ -167,21 +159,12 @@ func TestGetChangedTargets_CacheHit(t *testing.T) { stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) stream.EXPECT().Context().Return(t.Context()) - // Build a cached response with one ChangedTargets message and one Metadata message. - cachedChanged := &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{}, - }, - } - cachedMeta := &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_Metadata{ - Metadata: &pb.Metadata{}, - }, - } + // Build a cached response with one ChangedTargets message and one Metadata message, + // JSON-encoded (the storage layer uses newline-delimited JSON). var buf bytes.Buffer - w := gogio.NewDelimitedWriter(&buf) - w.WriteMsg(cachedChanged) - w.WriteMsg(cachedMeta) + enc := json.NewEncoder(&buf) + enc.Encode(entity.GetChangedTargetsResponse{ChangedTargets: []entity.ChangedTarget{}}) + enc.Encode(entity.GetChangedTargetsResponse{Metadata: &entity.Metadata{}}) cachedBytes := buf.Bytes() storagemock := storagemock.NewMockStorage(ctrl) @@ -297,17 +280,18 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { storagemock := storagemock.NewMockStorage(ctrl) var buf bytes.Buffer - gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }) + json.NewEncoder(&buf).Encode(entity.GetTargetGraphResponse{}) storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req storage.DownloadRequest) (storage.DownloadResponse, error) { if strings.Contains(req.Key, "compared-targets") { return storage.DownloadResponse{}, storage.NewNotFoundError(req.Key) } - if strings.Contains(req.Key, "th") { + if strings.Contains(req.Key, "treehashes") { + return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader([]byte("th")))}, nil + } + if strings.Contains(req.Key, "graphs") { return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader(buf.Bytes()))}, nil } - return storage.DownloadResponse{ReadCloser: io.NopCloser(bytes.NewReader([]byte("th")))}, nil + return storage.DownloadResponse{}, fmt.Errorf("unexpected key: %s", req.Key) }).AnyTimes() // Put is launched in a goroutine — use a channel to wait for it before the test ends. @@ -354,46 +338,34 @@ func TestGetChangedTargets_streamChunks(t *testing.T) { // Build first revision graph (2 chunks: Targets + Metadata) var buf1 bytes.Buffer - w1 := gogio.NewDelimitedWriter(&buf1) - w1.WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 1, Hash: "h1", RuleType: 100}, - {Id: 2, Hash: "h2-old", RuleType: 300}, - }, - }, + enc1 := json.NewEncoder(&buf1) + enc1.Encode(entity.GetTargetGraphResponse{ + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "h1", RuleType: 100}, + {ID: 2, Hash: "h2-old", RuleType: 300}, }, }) - w1.WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:target1", 2: "//app:target2"}, - RuleTypeMapping: map[int32]string{100: "go_library", 300: "source file"}, - }, + enc1.Encode(entity.GetTargetGraphResponse{ + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:target1", 2: "//app:target2"}, + RuleTypeMapping: map[int32]string{100: "go_library", 300: "source file"}, }, }) graph1Bytes := buf1.Bytes() // Build second revision graph - target2 has different hash var buf2 bytes.Buffer - w2 := gogio.NewDelimitedWriter(&buf2) - w2.WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 1, Hash: "h1", RuleType: 100}, - {Id: 2, Hash: "h2-new", RuleType: 300}, // changed hash - }, - }, + enc2 := json.NewEncoder(&buf2) + enc2.Encode(entity.GetTargetGraphResponse{ + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "h1", RuleType: 100}, + {ID: 2, Hash: "h2-new", RuleType: 300}, // changed hash }, }) - w2.WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:target1", 2: "//app:target2"}, - RuleTypeMapping: map[int32]string{100: "go_library", 300: "source file"}, - }, + enc2.Encode(entity.GetTargetGraphResponse{ + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:target1", 2: "//app:target2"}, + RuleTypeMapping: map[int32]string{100: "go_library", 300: "source file"}, }, }) graph2Bytes := buf2.Bytes() @@ -478,20 +450,14 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { // goroutine runs. Both revisions share the same target so there are no // diffs to send beyond the metadata chunk. var graphBuf bytes.Buffer - w := gogio.NewDelimitedWriter(&graphBuf) - w.WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{{Id: 1, Hash: "h1", RuleType: 100}}, - }, - }, + enc := json.NewEncoder(&graphBuf) + enc.Encode(entity.GetTargetGraphResponse{ + Targets: []entity.OptimizedTarget{{ID: 1, Hash: "h1", RuleType: 100}}, }) - w.WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:t1"}, - RuleTypeMapping: map[int32]string{100: "go_library"}, - }, + enc.Encode(entity.GetTargetGraphResponse{ + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:t1"}, + RuleTypeMapping: map[int32]string{100: "go_library"}, }, }) graphBytes := graphBuf.Bytes() @@ -579,180 +545,151 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { func TestCompareTargetGraphs_NewTarget_CanonicalIDs(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{}, - RuleTypeMapping: map[int32]string{}, - TagMapping: map[int32]string{}, - AttributeNameMapping: map[int32]string{}, - AttributeStringValueMapping: map[int32]string{}, - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{}, + RuleTypeMapping: map[int32]string{}, + TagMapping: map[int32]string{}, + AttributeNameMapping: map[int32]string{}, + AttributeStringValueMapping: map[int32]string{}, }, }, } - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 10, Hash: "h2", RuleType: 1}, - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 10, Hash: "h2", RuleType: 1}, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{10: "//app:new"}, - RuleTypeMapping: map[int32]string{1: "rule"}, - TagMapping: map[int32]string{}, - AttributeNameMapping: map[int32]string{}, - AttributeStringValueMapping: map[int32]string{}, - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{10: "//app:new"}, + RuleTypeMapping: map[int32]string{1: "rule"}, + TagMapping: map[int32]string{}, + AttributeNameMapping: map[int32]string{}, + AttributeStringValueMapping: map[int32]string{}, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) require.Len(t, res, 2) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) - require.Len(t, cs.GetChangedTargets(), 1) - ct := cs.GetChangedTargets()[0] - require.Equal(t, pb.CHANGE_TYPE_NEW, ct.GetChangeType()) + require.Len(t, cs, 1) + ct := cs[0] + require.Equal(t, int32(pb.CHANGE_TYPE_NEW), ct.ChangeType) // ID used in target should match canonical metadata mapping - meta := res[1].GetMetadata() + meta := res[1].Metadata require.NotNil(t, meta) - newID := ct.GetNewTarget().GetId() - require.Equal(t, "//app:new", meta.GetTargetIdMapping()[newID]) + newID := ct.NewTarget.ID + require.Equal(t, "//app:new", meta.TargetIDMapping[newID]) } func TestCompareTargetGraphs_SourceFileDirectAndPropagation(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) // Old: source file A (id 1, hash h1), lib L (id 2, hash h1, dep -> A) - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 1, Hash: "h1", RuleType: 100}, // "source file" - {Id: 2, Hash: "h1", RuleType: 200, DirectDependencies: []int32{1}}, // "rule" - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "h1", RuleType: 100}, // "source file" + {ID: 2, Hash: "h1", RuleType: 200, DirectDependencies: []int32{1}}, // "rule" }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 1: "//app:A", - 2: "//app:L", - }, - RuleTypeMapping: map[int32]string{ - 100: "source file", - 200: "rule", - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 1: "//app:A", + 2: "//app:L", + }, + RuleTypeMapping: map[int32]string{ + 100: "source file", + 200: "rule", }, }, }, } // New: both change hashes; same structure - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 11, Hash: "h2", RuleType: 101}, // "source file" - {Id: 22, Hash: "h2", RuleType: 201, DirectDependencies: []int32{11}}, // "rule" - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 11, Hash: "h2", RuleType: 101}, // "source file" + {ID: 22, Hash: "h2", RuleType: 201, DirectDependencies: []int32{11}}, // "rule" }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 11: "//app:A", - 22: "//app:L", - }, - RuleTypeMapping: map[int32]string{ - 101: "source file", - 201: "rule", - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 11: "//app:A", + 22: "//app:L", + }, + RuleTypeMapping: map[int32]string{ + 101: "source file", + 201: "rule", }, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) // Expect 2 changed: A (source-file seed, distance 0) and L (rule whose own src changed, distance 0) - require.Len(t, cs.GetChangedTargets(), 2) - var aCT, lCT *pb.ChangedTarget - for _, ct := range cs.GetChangedTargets() { - name := res[1].GetMetadata().GetTargetIdMapping()[ct.GetNewTarget().GetId()] + require.Len(t, cs, 2) + var aCT, lCT *entity.ChangedTarget + for i := range cs { + if cs[i].NewTarget == nil { + continue + } + name := res[1].Metadata.TargetIDMapping[cs[i].NewTarget.ID] if name == "//app:A" { - aCT = ct + aCT = &cs[i] } if name == "//app:L" { - lCT = ct + lCT = &cs[i] } } require.NotNil(t, aCT) require.NotNil(t, lCT) - require.Equal(t, pb.CHANGE_TYPE_CHANGED, aCT.GetChangeType()) - require.Equal(t, pb.CHANGE_TYPE_CHANGED, lCT.GetChangeType()) - assert.Equal(t, int32(0), aCT.GetDistance(), "source-file A with hash change is a seed (distance 0)") - assert.Equal(t, int32(0), lCT.GetDistance(), "rule L whose own source A changed is a seed (distance 0)") + require.Equal(t, int32(pb.CHANGE_TYPE_CHANGED), aCT.ChangeType) + require.Equal(t, int32(pb.CHANGE_TYPE_CHANGED), lCT.ChangeType) + assert.Equal(t, int32(0), aCT.Distance, "source-file A with hash change is a seed (distance 0)") + assert.Equal(t, int32(0), lCT.Distance, "rule L whose own source A changed is a seed (distance 0)") // Old and new IDs must match for each changed target under canonical metadata - require.Equal(t, aCT.GetOldTarget().GetId(), aCT.GetNewTarget().GetId()) - require.Equal(t, lCT.GetOldTarget().GetId(), lCT.GetNewTarget().GetId()) + require.Equal(t, aCT.OldTarget.ID, aCT.NewTarget.ID) + require.Equal(t, lCT.OldTarget.ID, lCT.NewTarget.ID) } func TestCompareTargetGraphs_ChangedRuleUnreachableFromAnySeed(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) // Old: T (id 1, rule), no deps - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 1, Hash: "h1", RuleType: 200}, - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "h1", RuleType: 200}, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:T"}, - RuleTypeMapping: map[int32]string{100: "source file", 200: "rule"}, - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:T"}, + RuleTypeMapping: map[int32]string{100: "source file", 200: "rule"}, }, }, } // New: T hash changed, still no deps - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 2, Hash: "h2", RuleType: 201}, - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 2, Hash: "h2", RuleType: 201}, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{2: "//app:T"}, - RuleTypeMapping: map[int32]string{101: "source file", 201: "rule"}, - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{2: "//app:T"}, + RuleTypeMapping: map[int32]string{101: "source file", 201: "rule"}, }, }, } @@ -761,262 +698,224 @@ func TestCompareTargetGraphs_ChangedRuleUnreachableFromAnySeed(t *testing.T) { // no upstream explanation becomes a distance-0 seed itself. res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) - require.Len(t, cs.GetChangedTargets(), 1) - got := cs.GetChangedTargets()[0] - require.Equal(t, pb.CHANGE_TYPE_CHANGED, got.GetChangeType()) - assert.Equal(t, int32(0), got.GetDistance(), "orphan hash change is seeded by trust-the-hasher") + require.Len(t, cs, 1) + got := cs[0] + require.Equal(t, int32(pb.CHANGE_TYPE_CHANGED), got.ChangeType) + assert.Equal(t, int32(0), got.Distance, "orphan hash change is seeded by trust-the-hasher") } func TestCompareTargetGraphs_ChangedWhenDependenciesChanged(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) // Old: T (id 1, rule) with deps on A - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 1, Hash: "h1", RuleType: 200, DirectDependencies: []int32{10}}, - {Id: 10, Hash: "h1", RuleType: 200}, // Dependency A - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "h1", RuleType: 200, DirectDependencies: []int32{10}}, + {ID: 10, Hash: "h1", RuleType: 200}, // Dependency A }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 1: "//app:T", - 10: "//app:A", - }, - RuleTypeMapping: map[int32]string{ - 100: "source file", - 200: "rule", - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 1: "//app:T", + 10: "//app:A", + }, + RuleTypeMapping: map[int32]string{ + 100: "source file", + 200: "rule", }, }, }, } // New: T now depends on B instead of A (hash changed due to dep change) - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 2, Hash: "h2", RuleType: 201, DirectDependencies: []int32{20}}, - {Id: 20, Hash: "h1", RuleType: 201}, // Dependency B - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 2, Hash: "h2", RuleType: 201, DirectDependencies: []int32{20}}, + {ID: 20, Hash: "h1", RuleType: 201}, // Dependency B }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 2: "//app:T", - 20: "//app:B", - }, - RuleTypeMapping: map[int32]string{ - 101: "source file", - 201: "rule", - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 2: "//app:T", + 20: "//app:B", + }, + RuleTypeMapping: map[int32]string{ + 101: "source file", + 201: "rule", }, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) // Find target T in the changed targets - var targetT *pb.ChangedTarget - for _, ct := range cs.GetChangedTargets() { - name := res[1].GetMetadata().GetTargetIdMapping()[ct.GetNewTarget().GetId()] + var targetT *entity.ChangedTarget + for i := range cs { + if cs[i].NewTarget == nil { + continue + } + name := res[1].Metadata.TargetIDMapping[cs[i].NewTarget.ID] if name == "//app:T" { - targetT = ct + targetT = &cs[i] break } } require.NotNil(t, targetT) - require.Equal(t, pb.CHANGE_TYPE_CHANGED, targetT.GetChangeType(), "Target with changed dependencies should be marked as CHANGED") - assert.Equal(t, int32(0), targetT.GetDistance(), "Target whose dep-name set changed is a seed (distance 0)") + require.Equal(t, int32(pb.CHANGE_TYPE_CHANGED), targetT.ChangeType, "Target with changed dependencies should be marked as CHANGED") + assert.Equal(t, int32(0), targetT.Distance, "Target whose dep-name set changed is a seed (distance 0)") } func TestCompareTargetGraphs_ChangedWhenAttributesChanged(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) // Old: T with attribute "key1" -> "value1" - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - { - Id: 1, - Hash: "h1", - RuleType: 200, - Attributes: map[int32]int32{1: 10}, // attr name 1 -> attr value 10 - }, - }, + Targets: []entity.OptimizedTarget{ + { + ID: 1, + Hash: "h1", + RuleType: 200, + Attributes: map[int32]int32{1: 10}, // attr name 1 -> attr value 10 }, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:T"}, - RuleTypeMapping: map[int32]string{ - 100: "source file", - 200: "rule", - }, - AttributeNameMapping: map[int32]string{1: "key1"}, - AttributeStringValueMapping: map[int32]string{10: "value1"}, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:T"}, + RuleTypeMapping: map[int32]string{ + 100: "source file", + 200: "rule", }, + AttributeNameMapping: map[int32]string{1: "key1"}, + AttributeStringValueMapping: map[int32]string{10: "value1"}, }, }, } // New: T with attribute "key1" -> "value2" (changed value) - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - { - Id: 2, - Hash: "h2", - RuleType: 201, - Attributes: map[int32]int32{2: 20}, // attr name 2 -> attr value 20 - }, - }, + Targets: []entity.OptimizedTarget{ + { + ID: 2, + Hash: "h2", + RuleType: 201, + Attributes: map[int32]int32{2: 20}, // attr name 2 -> attr value 20 }, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{2: "//app:T"}, - RuleTypeMapping: map[int32]string{ - 101: "source file", - 201: "rule", - }, - AttributeNameMapping: map[int32]string{2: "key1"}, - AttributeStringValueMapping: map[int32]string{20: "value2"}, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{2: "//app:T"}, + RuleTypeMapping: map[int32]string{ + 101: "source file", + 201: "rule", }, + AttributeNameMapping: map[int32]string{2: "key1"}, + AttributeStringValueMapping: map[int32]string{20: "value2"}, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) - require.Len(t, cs.GetChangedTargets(), 1) - got := cs.GetChangedTargets()[0] - require.Equal(t, pb.CHANGE_TYPE_CHANGED, got.GetChangeType(), "Target with changed attributes should be marked as CHANGED") - assert.Equal(t, int32(0), got.GetDistance(), "Target with own-config (attrs) change is a seed (distance 0)") + require.Len(t, cs, 1) + got := cs[0] + require.Equal(t, int32(pb.CHANGE_TYPE_CHANGED), got.ChangeType, "Target with changed attributes should be marked as CHANGED") + assert.Equal(t, int32(0), got.Distance, "Target with own-config (attrs) change is a seed (distance 0)") } func TestCompareTargetGraphs_ChangedWhenNewAttributeAdded(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) // Old: T with one attribute - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - { - Id: 1, - Hash: "h1", - RuleType: 200, - Attributes: map[int32]int32{1: 10}, - }, - }, + Targets: []entity.OptimizedTarget{ + { + ID: 1, + Hash: "h1", + RuleType: 200, + Attributes: map[int32]int32{1: 10}, }, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:T"}, - RuleTypeMapping: map[int32]string{ - 100: "source file", - 200: "rule", - }, - AttributeNameMapping: map[int32]string{1: "key1"}, - AttributeStringValueMapping: map[int32]string{10: "value1"}, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:T"}, + RuleTypeMapping: map[int32]string{ + 100: "source file", + 200: "rule", }, + AttributeNameMapping: map[int32]string{1: "key1"}, + AttributeStringValueMapping: map[int32]string{10: "value1"}, }, }, } // New: T with two attributes (added key2) - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - { - Id: 2, - Hash: "h2", - RuleType: 201, - Attributes: map[int32]int32{ - 2: 20, // key1 -> value1 - 3: 30, // key2 -> value2 (NEW) - }, - }, + Targets: []entity.OptimizedTarget{ + { + ID: 2, + Hash: "h2", + RuleType: 201, + Attributes: map[int32]int32{ + 2: 20, // key1 -> value1 + 3: 30, // key2 -> value2 (NEW) }, }, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{2: "//app:T"}, - RuleTypeMapping: map[int32]string{ - 101: "source file", - 201: "rule", - }, - AttributeNameMapping: map[int32]string{ - 2: "key1", - 3: "key2", - }, - AttributeStringValueMapping: map[int32]string{ - 20: "value1", - 30: "value2", - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{2: "//app:T"}, + RuleTypeMapping: map[int32]string{ + 101: "source file", + 201: "rule", + }, + AttributeNameMapping: map[int32]string{ + 2: "key1", + 3: "key2", + }, + AttributeStringValueMapping: map[int32]string{ + 20: "value1", + 30: "value2", }, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) - require.Len(t, cs.GetChangedTargets(), 1) - got := cs.GetChangedTargets()[0] - require.Equal(t, pb.CHANGE_TYPE_CHANGED, got.GetChangeType(), "Target with new attribute added should be marked as CHANGED") - assert.Equal(t, int32(0), got.GetDistance(), "Target with own-config (attrs) change is a seed (distance 0)") + require.Len(t, cs, 1) + got := cs[0] + require.Equal(t, int32(pb.CHANGE_TYPE_CHANGED), got.ChangeType, "Target with new attribute added should be marked as CHANGED") + assert.Equal(t, int32(0), got.Distance, "Target with own-config (attrs) change is a seed (distance 0)") } func TestSendTrimmedChangedTargets_MetadataAlwaysForwarded(t *testing.T) { ctrl := gomock.NewController(t) stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) - meta := &pb.Metadata{TargetIdMapping: map[int32]string{1: "//app:T"}} - responses := []*pb.GetChangedTargetsResponse{ + responses := []entity.GetChangedTargetsResponse{ { - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{ - ChangedTargets: []*pb.ChangedTarget{ - {Distance: 5, ChangeType: pb.CHANGE_TYPE_CHANGED}, - }, - }, + ChangedTargets: []entity.ChangedTarget{ + {Distance: 5, ChangeType: int32(pb.CHANGE_TYPE_CHANGED)}, }, }, { - Item: &pb.GetChangedTargetsResponse_Metadata{Metadata: meta}, + Metadata: &entity.Metadata{TargetIDMapping: map[int32]string{1: "//app:T"}}, }, } @@ -1032,18 +931,16 @@ func TestSendTrimmedChangedTargets_MetadataAlwaysForwarded(t *testing.T) { // First response: target filtered out (distance 5 > maxDist 1) assert.Empty(t, sent[0].GetChangedTargets().GetChangedTargets()) // Second response: metadata always forwarded - assert.Equal(t, meta, sent[1].GetMetadata()) + assert.NotNil(t, sent[1].GetMetadata()) } func TestSendTrimmedChangedTargets_SendError(t *testing.T) { ctrl := gomock.NewController(t) stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) - responses := []*pb.GetChangedTargetsResponse{ + responses := []entity.GetChangedTargetsResponse{ { - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{}, - }, + ChangedTargets: []entity.ChangedTarget{}, }, } @@ -1059,23 +956,15 @@ func TestGetChangedTargets_CacheHitWithDistanceFilter(t *testing.T) { stream.EXPECT().Context().Return(t.Context()) // Cached response: two targets at distances 0 and 2, plus metadata. - cachedChanged := &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{ - ChangedTargets: []*pb.ChangedTarget{ - {Distance: 0, ChangeType: pb.CHANGE_TYPE_CHANGED}, - {Distance: 2, ChangeType: pb.CHANGE_TYPE_CHANGED}, - }, - }, - }, - } - cachedMeta := &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_Metadata{Metadata: &pb.Metadata{}}, - } var buf bytes.Buffer - w := gogio.NewDelimitedWriter(&buf) - w.WriteMsg(cachedChanged) - w.WriteMsg(cachedMeta) + enc := json.NewEncoder(&buf) + enc.Encode(entity.GetChangedTargetsResponse{ + ChangedTargets: []entity.ChangedTarget{ + {Distance: 0, ChangeType: int32(pb.CHANGE_TYPE_CHANGED)}, + {Distance: 2, ChangeType: int32(pb.CHANGE_TYPE_CHANGED)}, + }, + }) + enc.Encode(entity.GetChangedTargetsResponse{Metadata: &entity.Metadata{}}) cachedBytes := buf.Bytes() storagemock := storagemock.NewMockStorage(ctrl) @@ -1121,93 +1010,84 @@ func TestCompareTargetGraphs_HashOnlyChangePropagatesViaBFS(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) // Old: T (rule) with deps on source file A (id 10) and attributes - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - { - Id: 1, - Hash: "h1", - RuleType: 200, - DirectDependencies: []int32{10}, - Attributes: map[int32]int32{1: 10}, - }, - {Id: 10, Hash: "h1", RuleType: 100}, // source file A - }, + Targets: []entity.OptimizedTarget{ + { + ID: 1, + Hash: "h1", + RuleType: 200, + DirectDependencies: []int32{10}, + Attributes: map[int32]int32{1: 10}, }, + {ID: 10, Hash: "h1", RuleType: 100}, // source file A }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 1: "//app:T", - 10: "//app:A", - }, - RuleTypeMapping: map[int32]string{ - 100: "source file", - 200: "rule", - }, - AttributeNameMapping: map[int32]string{1: "key1"}, - AttributeStringValueMapping: map[int32]string{10: "value1"}, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 1: "//app:T", + 10: "//app:A", }, + RuleTypeMapping: map[int32]string{ + 100: "source file", + 200: "rule", + }, + AttributeNameMapping: map[int32]string{1: "key1"}, + AttributeStringValueMapping: map[int32]string{10: "value1"}, }, }, } // New: source file A's hash changed (a seed); T's own config (deps, attrs) // is unchanged but its hash differs because of A. - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - { - Id: 2, - Hash: "h2", // Changed - RuleType: 201, - DirectDependencies: []int32{20}, // Same dep name (//app:A) - Attributes: map[int32]int32{2: 20}, // Same attribute - }, - {Id: 20, Hash: "h2", RuleType: 101}, // source file A, hash changed - }, + Targets: []entity.OptimizedTarget{ + { + ID: 2, + Hash: "h2", // Changed + RuleType: 201, + DirectDependencies: []int32{20}, // Same dep name (//app:A) + Attributes: map[int32]int32{2: 20}, // Same attribute }, + {ID: 20, Hash: "h2", RuleType: 101}, // source file A, hash changed }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 2: "//app:T", - 20: "//app:A", - }, - RuleTypeMapping: map[int32]string{ - 101: "source file", - 201: "rule", - }, - AttributeNameMapping: map[int32]string{2: "key1"}, - AttributeStringValueMapping: map[int32]string{20: "value1"}, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 2: "//app:T", + 20: "//app:A", + }, + RuleTypeMapping: map[int32]string{ + 101: "source file", + 201: "rule", }, + AttributeNameMapping: map[int32]string{2: "key1"}, + AttributeStringValueMapping: map[int32]string{20: "value1"}, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) // Find target T - var targetT *pb.ChangedTarget - for _, ct := range cs.GetChangedTargets() { - name := res[1].GetMetadata().GetTargetIdMapping()[ct.GetNewTarget().GetId()] + var targetT *entity.ChangedTarget + for i := range cs { + if cs[i].NewTarget == nil { + continue + } + name := res[1].Metadata.TargetIDMapping[cs[i].NewTarget.ID] if name == "//app:T" { - targetT = ct + targetT = &cs[i] break } } require.NotNil(t, targetT) - require.Equal(t, pb.CHANGE_TYPE_CHANGED, targetT.GetChangeType(), "Target with only hash change (not deps/attrs) is CHANGED") - assert.Equal(t, int32(0), targetT.GetDistance(), "T owns changed source file A so is a seed (distance 0)") + require.Equal(t, int32(pb.CHANGE_TYPE_CHANGED), targetT.ChangeType, "Target with only hash change (not deps/attrs) is CHANGED") + assert.Equal(t, int32(0), targetT.Distance, "T owns changed source file A so is a seed (distance 0)") } func TestCompareTargetGraphs_SiblingRuleNotPromotedToSeed(t *testing.T) { @@ -1217,128 +1097,109 @@ func TestCompareTargetGraphs_SiblingRuleNotPromotedToSeed(t *testing.T) { // Rule T (id 3) depends on L (sibling rule), NOT directly on A. // When A changes, L should be distance 0 (owns its changed src), // but T should be distance 1 (depends on changed rule, not its own src). - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 1, Hash: "h1", RuleType: 100}, // source file A - {Id: 2, Hash: "h1", RuleType: 200, DirectDependencies: []int32{1}}, // rule L -> A - {Id: 3, Hash: "h1", RuleType: 200, DirectDependencies: []int32{2}}, // rule T -> L - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "h1", RuleType: 100}, // source file A + {ID: 2, Hash: "h1", RuleType: 200, DirectDependencies: []int32{1}}, // rule L -> A + {ID: 3, Hash: "h1", RuleType: 200, DirectDependencies: []int32{2}}, // rule T -> L }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 1: "//pkg:A", - 2: "//pkg:L", - 3: "//pkg:T", - }, - RuleTypeMapping: map[int32]string{ - 100: "source file", - 200: "rule", - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 1: "//pkg:A", + 2: "//pkg:L", + 3: "//pkg:T", + }, + RuleTypeMapping: map[int32]string{ + 100: "source file", + 200: "rule", }, }, }, } - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 11, Hash: "h2", RuleType: 101}, // source file A changed - {Id: 22, Hash: "h2", RuleType: 201, DirectDependencies: []int32{11}}, // rule L -> A - {Id: 33, Hash: "h2", RuleType: 201, DirectDependencies: []int32{22}}, // rule T -> L - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 11, Hash: "h2", RuleType: 101}, // source file A changed + {ID: 22, Hash: "h2", RuleType: 201, DirectDependencies: []int32{11}}, // rule L -> A + {ID: 33, Hash: "h2", RuleType: 201, DirectDependencies: []int32{22}}, // rule T -> L }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{ - 11: "//pkg:A", - 22: "//pkg:L", - 33: "//pkg:T", - }, - RuleTypeMapping: map[int32]string{ - 101: "source file", - 201: "rule", - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{ + 11: "//pkg:A", + 22: "//pkg:L", + 33: "//pkg:T", + }, + RuleTypeMapping: map[int32]string{ + 101: "source file", + 201: "rule", }, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) - require.Len(t, cs.GetChangedTargets(), 3) + require.Len(t, cs, 3) - byName := make(map[string]*pb.ChangedTarget) - for _, ct := range cs.GetChangedTargets() { - name := res[1].GetMetadata().GetTargetIdMapping()[ct.GetNewTarget().GetId()] - byName[name] = ct + byName := make(map[string]*entity.ChangedTarget) + for i := range cs { + if cs[i].NewTarget == nil { + continue + } + name := res[1].Metadata.TargetIDMapping[cs[i].NewTarget.ID] + byName[name] = &cs[i] } - assert.Equal(t, int32(0), byName["//pkg:A"].GetDistance(), "source file A is a seed") - assert.Equal(t, int32(0), byName["//pkg:L"].GetDistance(), "rule L owns changed source A → seed") - assert.Equal(t, int32(1), byName["//pkg:T"].GetDistance(), "rule T depends on sibling rule L, not its own src → distance 1") + assert.Equal(t, int32(0), byName["//pkg:A"].Distance, "source file A is a seed") + assert.Equal(t, int32(0), byName["//pkg:L"].Distance, "rule L owns changed source A → seed") + assert.Equal(t, int32(1), byName["//pkg:T"].Distance, "rule T depends on sibling rule L, not its own src → distance 1") } func TestCompareTargetGraphs_DeletedTargetEmitted(t *testing.T) { c := newTestController(zaptest.NewLogger(t)) // Old: T (rule) exists; New: T is gone. - first := []*pb.GetTargetGraphResponse{ + first := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{ - Targets: []*pb.OptimizedTarget{ - {Id: 1, Hash: "h1", RuleType: 200}, - }, - }, + Targets: []entity.OptimizedTarget{ + {ID: 1, Hash: "h1", RuleType: 200}, }, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:T"}, - RuleTypeMapping: map[int32]string{100: "source file", 200: "rule"}, - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:T"}, + RuleTypeMapping: map[int32]string{100: "source file", 200: "rule"}, }, }, } - second := []*pb.GetTargetGraphResponse{ + second := []entity.GetTargetGraphResponse{ { - Item: &pb.GetTargetGraphResponse_Targets{ - Targets: &pb.OptimizedTargets{Targets: []*pb.OptimizedTarget{}}, - }, + Targets: []entity.OptimizedTarget{}, }, { - Item: &pb.GetTargetGraphResponse_Metadata{ - Metadata: &pb.Metadata{ - TargetIdMapping: map[int32]string{}, - RuleTypeMapping: map[int32]string{101: "source file", 201: "rule"}, - }, + Metadata: &entity.Metadata{ + TargetIDMapping: map[int32]string{}, + RuleTypeMapping: map[int32]string{101: "source file", 201: "rule"}, }, }, } res, err := c.compareTargetGraphs(t.Context(), c.scope, zap.NewNop(), first, second, -1) require.NoError(t, err) - cs := res[0].GetChangedTargets() + cs := res[0].ChangedTargets require.NotNil(t, cs) - require.Len(t, cs.GetChangedTargets(), 1) - got := cs.GetChangedTargets()[0] - require.Equal(t, pb.CHANGE_TYPE_DELETED, got.GetChangeType()) - require.NotNil(t, got.GetOldTarget(), "DELETED entry must carry OldTarget") - assert.Nil(t, got.GetNewTarget(), "DELETED entry must not carry NewTarget") - assert.Equal(t, int32(0), got.GetDistance(), "DELETED targets are seeds (distance 0)") + require.Len(t, cs, 1) + got := cs[0] + require.Equal(t, int32(pb.CHANGE_TYPE_DELETED), got.ChangeType) + require.NotNil(t, got.OldTarget, "DELETED entry must carry OldTarget") + assert.Nil(t, got.NewTarget, "DELETED entry must not carry NewTarget") + assert.Equal(t, int32(0), got.Distance, "DELETED targets are seeds (distance 0)") // Old id is remapped into the canonical id space; metadata must resolve back to the deleted name. - assert.Equal(t, "//app:T", res[1].GetMetadata().GetTargetIdMapping()[got.GetOldTarget().GetId()]) + assert.Equal(t, "//app:T", res[1].Metadata.TargetIDMapping[got.OldTarget.ID]) } func TestSendTrimmedChangedTargets_RetainsDeletedAtMaxDistanceOne(t *testing.T) { @@ -1346,16 +1207,12 @@ func TestSendTrimmedChangedTargets_RetainsDeletedAtMaxDistanceOne(t *testing.T) stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) // DELETED entries are seeds (distance 0) and must survive max_distance=1. - responses := []*pb.GetChangedTargetsResponse{ + responses := []entity.GetChangedTargetsResponse{ { - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{ - ChangedTargets: []*pb.ChangedTarget{ - {Distance: 0, ChangeType: pb.CHANGE_TYPE_DELETED}, - {Distance: 1, ChangeType: pb.CHANGE_TYPE_CHANGED}, - {Distance: 5, ChangeType: pb.CHANGE_TYPE_CHANGED}, - }, - }, + ChangedTargets: []entity.ChangedTarget{ + {Distance: 0, ChangeType: int32(pb.CHANGE_TYPE_DELETED)}, + {Distance: 1, ChangeType: int32(pb.CHANGE_TYPE_CHANGED)}, + {Distance: 5, ChangeType: int32(pb.CHANGE_TYPE_CHANGED)}, }, }, } @@ -1380,21 +1237,6 @@ func TestSendTrimmedChangedTargets_RetainsDeletedAtMaxDistanceOne(t *testing.T) assert.True(t, gotDeleted, "DELETED entry at distance 0 must survive max_distance=1") } -func newMockGraphReader(ctrl *gomock.Controller, chunks ...*pb.GetTargetGraphResponse) *storagemock.MockGraphReader { - r := storagemock.NewMockGraphReader(ctrl) - idx := 0 - r.EXPECT().Read().DoAndReturn(func() (*pb.GetTargetGraphResponse, error) { - if idx >= len(chunks) { - return nil, io.EOF - } - c := chunks[idx] - idx++ - return c, nil - }).AnyTimes() - r.EXPECT().Close().Return(nil).AnyTimes() - return r -} - func changedTargetsRequest() *pb.GetChangedTargetsRequest { return &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -1423,14 +1265,16 @@ func TestServeChangedTargetsFromCache(t *testing.T) { t.Run("corrupt cached blob falls through to recompute", func(t *testing.T) { ctrl := gomock.NewController(t) - // A valid single-message blob truncated by one byte — mimics an - // incomplete concurrent write. The reader errors partway, and the - // caller must fall through (served=false) without sending anything. + // A two-message JSON blob with the second message truncated — mimics an + // incomplete concurrent write. The reader returns the first message fine + // but errors on the second, and the caller must fall through + // (served=false) without sending anything. var buf bytes.Buffer - require.NoError(t, gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ChangedTargets: &pb.ChangedTargets{}}, - })) - truncated := buf.Bytes()[:buf.Len()-1] + enc := json.NewEncoder(&buf) + enc.Encode(entity.GetChangedTargetsResponse{ChangedTargets: []entity.ChangedTarget{}}) + enc.Encode(entity.GetChangedTargetsResponse{Metadata: &entity.Metadata{}}) + // Truncate well into the second JSON object to guarantee corruption. + truncated := buf.Bytes()[:buf.Len()-5] st := storagemock.NewMockStorage(ctrl) st.EXPECT().Get(gomock.Any(), gomock.Any()).DoAndReturn( @@ -1461,13 +1305,9 @@ func TestServeChangedTargetsFromCache(t *testing.T) { ctrl := gomock.NewController(t) var buf bytes.Buffer - w := gogio.NewDelimitedWriter(&buf) - require.NoError(t, w.WriteMsg(&pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ChangedTargets: &pb.ChangedTargets{}}, - })) - require.NoError(t, w.WriteMsg(&pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_Metadata{Metadata: &pb.Metadata{}}, - })) + enc := json.NewEncoder(&buf) + enc.Encode(entity.GetChangedTargetsResponse{ChangedTargets: []entity.ChangedTarget{}}) + enc.Encode(entity.GetChangedTargetsResponse{Metadata: &entity.Metadata{}}) cached := buf.Bytes() st := storagemock.NewMockStorage(ctrl) @@ -1504,16 +1344,14 @@ func TestFetchTargetGraphs(t *testing.T) { r.BypassCache = true return r } - chunk := &pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Metadata{Metadata: &pb.Metadata{}}, - } + entityChunk := entity.GetTargetGraphResponse{Metadata: &entity.Metadata{}} t.Run("returns both graphs on success", func(t *testing.T) { ctrl := gomock.NewController(t) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, _ entity.GetTargetGraphRequest) (storage.GraphReader, error) { - return newMockGraphReader(ctrl, chunk), nil + return newGraphReader(t, entityChunk), nil }).Times(2) c := newTestController(zaptest.NewLogger(t)) @@ -1534,7 +1372,7 @@ func TestFetchTargetGraphs(t *testing.T) { if p.Build.BaseSha == "sha1" { return nil, injected } - return newMockGraphReader(ctrl, chunk), nil + return newGraphReader(t, entityChunk), nil }).Times(2) c := newTestController(zaptest.NewLogger(t)) @@ -1552,7 +1390,7 @@ func TestFetchTargetGraphs(t *testing.T) { orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, _ entity.GetTargetGraphRequest) (storage.GraphReader, error) { - return newMockGraphReader(ctrl), nil + return newGraphReader(t), nil }).Times(2) c := newTestController(zaptest.NewLogger(t)) @@ -1579,7 +1417,7 @@ func TestFetchTargetGraphs(t *testing.T) { } func TestToDiffGraph_SkipsUnresolvedIDs(t *testing.T) { - targetsByID := map[int32]*pb.OptimizedTarget{ + targetsByID := map[int32]*entity.OptimizedTarget{ 1: { Hash: "h1", RuleType: 100, @@ -1590,8 +1428,8 @@ func TestToDiffGraph_SkipsUnresolvedIDs(t *testing.T) { 2: {Hash: "h2", RuleType: 100}, 3: {Hash: "h3", RuleType: 100}, } - meta := &pb.Metadata{ - TargetIdMapping: map[int32]string{1: "//app:a", 2: "//app:b"}, + meta := &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//app:a", 2: "//app:b"}, RuleTypeMapping: map[int32]string{100: "go_library"}, TagMapping: map[int32]string{10: "tag_a"}, AttributeNameMapping: map[int32]string{20: "attr_a"}, diff --git a/controller/gettargetgraph.go b/controller/gettargetgraph.go index 630e16d0..662a1c13 100644 --- a/controller/gettargetgraph.go +++ b/controller/gettargetgraph.go @@ -67,7 +67,7 @@ func (c *controller) GetTargetGraph(request *pb.GetTargetGraphRequest, stream pb sendStart := time.Now() outputConfig := request.GetOutputConfig() for { - graphStreamChunk, err := graphReader.Read() + chunk, err := graphReader.Read() if err == io.EOF { sendDuration := time.Since(sendStart) totalDuration := time.Since(start) @@ -82,7 +82,8 @@ func (c *controller) GetTargetGraph(request *pb.GetTargetGraphRequest, stream pb if err != nil { return common.WithReason(failureReasonGraphFetch, common.ErrorTypeInfra, err) } - toSend := applyOptimizedTargetsOutputConfigToChunk(graphStreamChunk, outputConfig) + protoResp := mapper.GetTargetGraphResponseToProto(&chunk) + toSend := applyOptimizedTargetsOutputConfigToChunk(protoResp, outputConfig) err = stream.Send(toSend) if err != nil { return common.WithReason(failureReasonSend, common.ErrorTypeInfra, fmt.Errorf("send graph: %w", err)) diff --git a/controller/gettargetgraph_test.go b/controller/gettargetgraph_test.go index 4a1ff02e..cd91e507 100644 --- a/controller/gettargetgraph_test.go +++ b/controller/gettargetgraph_test.go @@ -21,12 +21,14 @@ import ( "io" "testing" - gogio "github.com/gogo/protobuf/io" + "encoding/json" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/uber/tango/core/common" "github.com/uber/tango/core/storage" storagemock "github.com/uber/tango/core/storage/storagemock" + "github.com/uber/tango/entity" orchestratormock "github.com/uber/tango/orchestrator/orchestratormock" pb "github.com/uber/tango/tangopb" tangomock "github.com/uber/tango/tangopb/tangopbmock" @@ -121,10 +123,7 @@ func TestGetTargetGraph_SendsWhenItemPresent(t *testing.T) { stream.EXPECT().Send(gomock.Any()).Return(nil) store := storagemock.NewMockStorage(ctrl) var buf bytes.Buffer - err := gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }) - require.NoError(t, err) + require.NoError(t, json.NewEncoder(&buf).Encode(entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}})) gomock.InOrder( store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-xyz"))}, nil), @@ -134,7 +133,7 @@ func TestGetTargetGraph_SendsWhenItemPresent(t *testing.T) { Logger: zaptest.NewLogger(t), Storage: store, }) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", BaseSha: "sha", @@ -191,16 +190,7 @@ func TestGetTargetGraph_TreehashNotFound_NoError(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, storage.NewNotFoundError("x")) orchestrator := orchestratormock.NewMockOrchestrator(ctrl) - // Provide a fake GraphReader that yields one message then EOF - graphReader := storagemock.NewMockGraphReader(ctrl) - graphReader.EXPECT().Read().DoAndReturn(func() (*pb.GetTargetGraphResponse, error) { - return &pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }, nil - }).Times(1) - // Controller may call Read again to observe EOF - graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) - graphReader.EXPECT().Close().Return(nil) + graphReader := newGraphReader(t, entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}}) orchestrator.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) c := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), @@ -282,10 +272,7 @@ func TestGetTargetGraph_StreamSendError(t *testing.T) { storagemock := storagemock.NewMockStorage(ctrl) var buf bytes.Buffer - err := gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }) - require.NoError(t, err) + require.NoError(t, json.NewEncoder(&buf).Encode(entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}})) stream.EXPECT().Send(gomock.Any()).Return(errors.New("send fail")) gomock.InOrder( @@ -296,7 +283,7 @@ func TestGetTargetGraph_StreamSendError(t *testing.T) { Logger: zaptest.NewLogger(t), Storage: storagemock, }) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) assert.Error(t, err) @@ -314,12 +301,7 @@ func TestGetTargetGraph_GraphNotFound_FallsThrough(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, storage.NewNotFoundError("graphs/abc")), ) orch := orchestratormock.NewMockOrchestrator(ctrl) - graphReader := storagemock.NewMockGraphReader(ctrl) - graphReader.EXPECT().Read().Return(&pb.GetTargetGraphResponse{ - Item: &pb.GetTargetGraphResponse_Targets{Targets: &pb.OptimizedTargets{}}, - }, nil).Times(1) - graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) - graphReader.EXPECT().Close().Return(nil) + graphReader := newGraphReader(t, entity.GetTargetGraphResponse{Targets: []entity.OptimizedTarget{}}) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) c := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go index 56ff4396..1b8d4d77 100644 --- a/controller/testhelper_test.go +++ b/controller/testhelper_test.go @@ -16,20 +16,32 @@ package controller import ( "context" + "testing" + "github.com/stretchr/testify/require" "github.com/uber-go/tally" - "github.com/uber/tango/core/common" + "github.com/uber/tango/core/storage" + "github.com/uber/tango/entity" "go.uber.org/zap" ) func newTestController(logger *zap.Logger) *controller { return &controller{ - logger: logger, - scope: tally.NoopScope, - targetChunkSize: common.DefaultTargetChunkSize, - changedTargetChunkSize: common.DefaultChangedTargetChunkSize, - metadataMapChunkSize: common.DefaultMetadataMapChunkSize, - totalDurationBuckets: _totalDurationBuckets, - appCtx: context.Background(), + logger: logger, + scope: tally.NoopScope, + maxMessageBytes: _defaultMaxMessageBytes, + totalDurationBuckets: _totalDurationBuckets, + appCtx: context.Background(), } } + +// newGraphReader builds a storage.GraphReader from entity chunks +// by writing JSON to in-memory storage and reading back. +func newGraphReader(t *testing.T, chunks ...entity.GetTargetGraphResponse) storage.GraphReader { + t.Helper() + st := storage.NewMemoryStorage() + require.NoError(t, storage.WriteGraphStream(context.Background(), st, "test-graph", chunks)) + reader, err := storage.NewGraphReader(context.Background(), st, "test-graph") + require.NoError(t, err) + return reader +} diff --git a/core/common/BUILD.bazel b/core/common/BUILD.bazel index 11f4184d..518441f8 100644 --- a/core/common/BUILD.bazel +++ b/core/common/BUILD.bazel @@ -8,22 +8,11 @@ go_library( ], importpath = "github.com/uber/tango/core/common", visibility = ["//visibility:public"], - deps = [ - "//core/targethasher", - "//internal/mapper/idmapper", - "//tangopb", - "@com_github_bazelbuild_buildtools//build_proto", - ], ) go_test( name = "common_test", srcs = ["utils_test.go"], embed = [":common"], - deps = [ - "//core/targethasher", - "//tangopb", - "@com_github_stretchr_testify//assert", - "@com_github_stretchr_testify//require", - ], + deps = ["@com_github_stretchr_testify//assert"], ) diff --git a/core/common/utils.go b/core/common/utils.go index e0f47d47..d8062035 100644 --- a/core/common/utils.go +++ b/core/common/utils.go @@ -15,32 +15,7 @@ package common import ( - "context" - "encoding/hex" "strings" - - buildpb "github.com/bazelbuild/buildtools/build_proto" - "github.com/uber/tango/core/targethasher" - "github.com/uber/tango/internal/mapper/idmapper" - "github.com/uber/tango/tangopb" -) - -const ( - // DefaultTargetChunkSize is the default number of OptimizedTarget entries per stream message. - // Sized conservatively: at ~40KB/target worst-case (target with ~10K direct deps × 4 bytes), - // 250 targets ≈ 10MB — well under the 64MB default gRPC per-message limit. - DefaultTargetChunkSize = 250 - - // DefaultChangedTargetChunkSize is the default number of ChangedTarget entries per stream message. - // A ChangedTarget carries both old_target and new_target (2× an OptimizedTarget), so we use - // half the regular chunk size to stay within the same byte budget. - DefaultChangedTargetChunkSize = 125 - - // DefaultMetadataMapChunkSize is the max entries per metadata message chunk. - // target_id_mapping and attribute_string_value_mapping scale with repo size and can exceed - // the 64MB gRPC message limit for large monorepos, so they are split across multiple messages. - // At ~85 bytes/entry (60-char avg target name + proto overhead), 50 000 entries ≈ 4.25MB per chunk. - DefaultMetadataMapChunkSize = 50_000 ) // ToShortRemote returns the short remote name given a git ssh remote string. @@ -49,217 +24,3 @@ func ToShortRemote(remote string) string { strs := strings.Split(remote, ":") return strs[len(strs)-1] } - -// cancelCheckInterval is how often we poll ctx.Err() inside per-target hot loops. -// Picked to keep overhead negligible while still surfacing cancellation in <100ms -// for typical target rates. -const cancelCheckInterval = 4096 - -// ResultToGetTargetGraphResponse converts a Result to a GetTargetGraphResponse. -// targetChunkSize controls how many OptimizedTarget entries per stream message. -// metadataMapChunkSize controls how many entries per metadata map chunk. -// TODO: move this function to internal/mapper -func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result, targetChunkSize, metadataMapChunkSize int) ([]*tangopb.GetTargetGraphResponse, error) { - // Map target names to ids. This list is topologically sorted, so the ids are stable. - // IDs start at 1 — 0 is reserved as the proto3 "unset" sentinel so consumers using - // encoding/json (which honors `omitempty` on int32 fields) never silently lose a target. - targetNamesMapping := make(map[string]int32, len(result.TargetNames)) - for i, name := range result.TargetNames { - targetNamesMapping[name] = int32(i + 1) - } - - ruleTypeMapper := idmapper.NewMapper() - tagMapper := idmapper.NewMapper() - attrNameMapper := idmapper.NewMapper() - attrStrValMapper := idmapper.NewMapper() - - // Build the optimized targets slice - optimizedTargets := make([]*tangopb.OptimizedTarget, 0, len(result.Targets)) - - n := 0 - for _, t := range result.Targets { - if n%cancelCheckInterval == 0 { - if err := ctx.Err(); err != nil { - return nil, err - } - } - n++ - nameID := targetNamesMapping[t.Name] - - depIDs := make([]int32, 0, len(t.Deps)) - for _, depName := range t.Deps { - depID, ok := targetNamesMapping[depName] - if !ok { - continue - } - depIDs = append(depIDs, depID) - } - - ot := &tangopb.OptimizedTarget{ - Id: nameID, - Hash: hex.EncodeToString(t.Hash), - DirectDependencies: depIDs, - } - - // RuleType - if t.RuleType != "" { - id := ruleTypeMapper.ID(t.RuleType) - ot.RuleType = id - } - - // Tags - if len(t.Tags) > 0 { - tagIDs := make([]int32, 0, len(t.Tags)) - for _, tag := range t.Tags { - tagIDs = append(tagIDs, tagMapper.ID(tag)) - } - ot.Tags = tagIDs - } - ot.Root = t.Root - ot.External = t.External - if len(t.Attributes) > 0 { - attrs := make(map[int32]int32, len(t.Attributes)) - for _, attr := range t.Attributes { - // Only include STRING attributes with non-nil name and value to avoid nil dereferences. - if attr.GetType() == buildpb.Attribute_STRING && attr.Name != nil && attr.StringValue != nil { - nameID := attrNameMapper.ID(*attr.Name) - valID := attrStrValMapper.ID(*attr.StringValue) - attrs[nameID] = valID - } - } - ot.Attributes = attrs - } - - optimizedTargets = append(optimizedTargets, ot) - } - - // Invert mappings: string -> id => id -> string - targetIDToName := make(map[int32]string, len(targetNamesMapping)) - for s, id := range targetNamesMapping { - targetIDToName[id] = s - } - - ruleTypeIDToName := ruleTypeMapper.Invert() - tagIDToName := tagMapper.Invert() - attrNameIDToName := attrNameMapper.Invert() - attrStrValIDToVal := attrStrValMapper.Invert() - - // chunk targets into multiple messages for streaming - responses := chunkTargets(optimizedTargets, targetChunkSize) - for _, meta := range ChunkMetadata( - targetIDToName, - ruleTypeIDToName, - tagIDToName, - attrNameIDToName, - attrStrValIDToVal, - metadataMapChunkSize, - ) { - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Metadata{Metadata: meta}, - }) - } - - return responses, nil -} - -func chunkTargets(targets []*tangopb.OptimizedTarget, chunkSize int) []*tangopb.GetTargetGraphResponse { - if chunkSize <= 0 { - chunkSize = DefaultTargetChunkSize - } - - // at least one chunk - numChunks := max(1, (len(targets)+chunkSize-1)/chunkSize) - - responses := make([]*tangopb.GetTargetGraphResponse, 0, numChunks) - - for i := 0; i < len(targets); i += chunkSize { - end := i + chunkSize - if end > len(targets) { - end = len(targets) - } - - chunk := targets[i:end] - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Targets{ - Targets: &tangopb.OptimizedTargets{ - Targets: chunk, - }, - }, - }) - } - - // Handle empty targets case - if len(responses) == 0 { - responses = append(responses, &tangopb.GetTargetGraphResponse{ - Item: &tangopb.GetTargetGraphResponse_Targets{ - Targets: &tangopb.OptimizedTargets{ - Targets: []*tangopb.OptimizedTarget{}, - }, - }, - }) - } - - return responses -} - -// ChunkMetadata splits the metadata maps into multiple Metadata messages. -// target_id_mapping and attribute_string_value_mapping scale with repo size and can exceed the -// 64MB gRPC per-message limit for large monorepos; they are split across chunks of chunkSize entries. -// The small maps (rule_type, tag, attribute_name) always fit in one message and are sent in the first chunk. -func ChunkMetadata( - targetIDToName map[int32]string, - ruleTypeIDToName map[int32]string, - tagIDToName map[int32]string, - attrNameIDToName map[int32]string, - attrStrValIDToVal map[int32]string, - chunkSize int, -) []*tangopb.Metadata { - if chunkSize <= 0 { - chunkSize = DefaultMetadataMapChunkSize - } - - targetChunks := splitMap(targetIDToName, chunkSize) - attrValChunks := splitMap(attrStrValIDToVal, chunkSize) - - numChunks := max(1, max(len(targetChunks), len(attrValChunks))) - chunks := make([]*tangopb.Metadata, 0, numChunks) - - for i := range numChunks { - meta := &tangopb.Metadata{} - // Small maps are always small enough to fit in one message; include them in the first chunk. - if i == 0 { - meta.RuleTypeMapping = ruleTypeIDToName - meta.TagMapping = tagIDToName - meta.AttributeNameMapping = attrNameIDToName - } - if i < len(targetChunks) { - meta.TargetIdMapping = targetChunks[i] - } - if i < len(attrValChunks) { - meta.AttributeStringValueMapping = attrValChunks[i] - } - chunks = append(chunks, meta) - } - - return chunks -} - -// splitMap splits a map[int32]string into slices of at most size entries each. -func splitMap(m map[int32]string, size int) []map[int32]string { - if len(m) == 0 { - return nil - } - chunks := make([]map[int32]string, 0, (len(m)+size-1)/size) - current := make(map[int32]string, size) - for k, v := range m { - current[k] = v - if len(current) >= size { - chunks = append(chunks, current) - current = make(map[int32]string, size) - } - } - if len(current) > 0 { - chunks = append(chunks, current) - } - return chunks -} diff --git a/core/common/utils_test.go b/core/common/utils_test.go index 63e5eb24..ddcd9969 100644 --- a/core/common/utils_test.go +++ b/core/common/utils_test.go @@ -15,14 +15,9 @@ package common import ( - "context" - "fmt" "testing" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "github.com/uber/tango/core/targethasher" - pb "github.com/uber/tango/tangopb" ) func TestToShortRemote(t *testing.T) { @@ -55,93 +50,3 @@ func TestToShortRemote(t *testing.T) { }) } } - -func TestChunkTargets(t *testing.T) { - t.Parallel() - - // Create 25 targets, chunk by 10 → expect 3 chunks (10, 10, 5) - targets := make([]*pb.OptimizedTarget, 25) - for i := range targets { - targets[i] = &pb.OptimizedTarget{Id: int32(i)} - } - - responses := chunkTargets(targets, 10) - - require.Len(t, responses, 3) - - // Verify total count and order preserved - var total int - for _, resp := range responses { - item := resp.Item.(*pb.GetTargetGraphResponse_Targets) - for _, target := range item.Targets.Targets { - assert.Equal(t, int32(total), target.Id) - total++ - } - } - assert.Equal(t, 25, total) -} - -func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { - t.Parallel() - - numTargets := 50 - result := targethasher.Result{ - TargetNames: make([]string, numTargets), - Targets: make(map[string]*targethasher.Target, numTargets), - } - for i := 0; i < numTargets; i++ { - name := fmt.Sprintf("//pkg:target%d", i) - result.TargetNames[i] = name - result.Targets[name] = &targethasher.Target{Name: name, Hash: []byte{0}, RuleType: "go_library"} - } - - tests := []struct { - name string - targetChunkSize int - metadataMapChunkSize int - wantTargetChunks int - wantMetadataChunks int - }{ - { - name: "25 per chunk", - targetChunkSize: 25, - metadataMapChunkSize: 20, - wantTargetChunks: 2, - wantMetadataChunks: 3, - }, - { - name: "10 per chunk", - targetChunkSize: 10, - metadataMapChunkSize: 10, - wantTargetChunks: 5, - wantMetadataChunks: 5, - }, - { - name: "all in one chunk", - targetChunkSize: 100, - metadataMapChunkSize: 5_000, - wantTargetChunks: 1, - wantMetadataChunks: 1, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - responses, err := ResultToGetTargetGraphResponse(context.Background(), result, tt.targetChunkSize, tt.metadataMapChunkSize) - require.NoError(t, err) - - var targetChunks, metadataChunks int - for _, resp := range responses { - switch resp.Item.(type) { - case *pb.GetTargetGraphResponse_Targets: - targetChunks++ - case *pb.GetTargetGraphResponse_Metadata: - metadataChunks++ - } - } - assert.Equal(t, tt.wantTargetChunks, targetChunks) - assert.Equal(t, tt.wantMetadataChunks, metadataChunks) - }) - } -} diff --git a/example/cmd/query-bench/BUILD.bazel b/example/cmd/query-bench/BUILD.bazel index ba71ce1c..07ac5907 100644 --- a/example/cmd/query-bench/BUILD.bazel +++ b/example/cmd/query-bench/BUILD.bazel @@ -7,8 +7,8 @@ go_library( visibility = ["//visibility:private"], deps = [ "//core/bazel", - "//core/common", "//core/targethasher", + "//internal/mapper", "@com_github_gogo_protobuf//jsonpb", "@org_uber_go_zap//:zap", ], diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 28e4e52b..41b0f6c7 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -32,11 +32,13 @@ import ( "github.com/gogo/protobuf/jsonpb" "github.com/uber/tango/core/bazel" - "github.com/uber/tango/core/common" "github.com/uber/tango/core/targethasher" + "github.com/uber/tango/internal/mapper" "go.uber.org/zap" ) +const defaultMaxMessageBytes = 4_250_000 + func main() { if err := run(); err != nil { fmt.Fprintf(os.Stderr, "error: %v\n", err) @@ -112,18 +114,20 @@ func run() error { totalDuration += elapsed fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) start = time.Now() - response, err := common.ResultToGetTargetGraphResponse( - ctx, targethasherResult, - common.DefaultTargetChunkSize, common.DefaultMetadataMapChunkSize, - ) + targets, meta, err := mapper.ResultToTargetGraph(ctx, targethasherResult) + if err != nil { + return fmt.Errorf("run %d: converting to target graph: %w", i+1, err) + } + chunks, err := mapper.ChunkTargetGraph(targets, meta, defaultMaxMessageBytes) if err != nil { - return fmt.Errorf("run %d: converting to GetTargetGraphResponse: %w", i+1, err) + return fmt.Errorf("run %d: chunking target graph: %w", i+1, err) } elapsed = time.Since(start) - fmt.Printf("run %d: ResultToGetTargetGraphResponse: %v (%d responses)\n", i+1, elapsed.Round(time.Millisecond), len(response)) + fmt.Printf("run %d: ResultToTargetGraph+Chunk: %v (%d chunks)\n", i+1, elapsed.Round(time.Millisecond), len(chunks)) m := jsonpb.Marshaler{Indent: " "} - for _, r := range response { - if err := m.Marshal(os.Stdout, r); err != nil { + for _, chunk := range chunks { + protoResp := mapper.GetTargetGraphResponseToProto(&chunk) + if err := m.Marshal(os.Stdout, protoResp); err != nil { return fmt.Errorf("run %d: encoding response: %w", i+1, err) } fmt.Println()