From 9a2515372b369c6f686ee082552f74a61ea41481 Mon Sep 17 00:00:00 2001 From: Xiaoyang Tan Date: Sat, 25 Jul 2026 00:05:15 -0700 Subject: [PATCH] refactor(mapper): add ResultToGraphChunks convenience for the map+split flow Reintroduce a single public entry point that turns a targethasher.Result into the ordered []entity.GetTargetGraphResponse chunks that storage.WriteGraphStream consumes. The map -> SplitBySize -> SplitMetadata recipe was hand-assembled identically in nativeOrchestrator and query-bench, and off-module callers (e.g. CI-integrated orchestrators) could not reproduce it at all because streaming.SplitBySize / SplitMetadata live under internal/. ResultToGraphChunks keeps the split helpers internal while exposing the common "map, split, hand to WriteGraphStream" path. The two-step ResultToTargetGraph API stays for callers that need to control chunking. Both existing callsites collapse onto the new function (net -40 LOC). --- example/cmd/query-bench/BUILD.bazel | 2 - example/cmd/query-bench/main.go | 30 ++------------- mapper/BUILD.bazel | 1 + mapper/mapper.go | 44 ++++++++++++++++++++++ mapper/mapper_test.go | 58 +++++++++++++++++++++++++++++ orchestrator/BUILD.bazel | 1 - orchestrator/native_orchestrator.go | 25 +------------ 7 files changed, 107 insertions(+), 54 deletions(-) diff --git a/example/cmd/query-bench/BUILD.bazel b/example/cmd/query-bench/BUILD.bazel index d5f616c0..1defc2d7 100644 --- a/example/cmd/query-bench/BUILD.bazel +++ b/example/cmd/query-bench/BUILD.bazel @@ -9,9 +9,7 @@ go_library( "//config", "//core/bazel", "//core/targethasher", - "//entity", "//internal/mapper", - "//internal/streaming", "//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 85c32d28..faed313f 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -34,9 +34,7 @@ import ( "github.com/uber/tango/config" "github.com/uber/tango/core/bazel" "github.com/uber/tango/core/targethasher" - "github.com/uber/tango/entity" "github.com/uber/tango/internal/mapper" - "github.com/uber/tango/internal/streaming" tgmapper "github.com/uber/tango/mapper" "go.uber.org/zap" ) @@ -116,34 +114,12 @@ 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() - targets, meta, err := tgmapper.ResultToTargetGraph(ctx, targethasherResult) + chunks, err := tgmapper.ResultToGraphChunks(ctx, targethasherResult, config.DefaultMaxMessageBytes) if err != nil { - return fmt.Errorf("run %d: converting to target graph: %w", i+1, err) - } - targetGroups, err := streaming.SplitBySize(targets, config.DefaultMaxMessageBytes) - if err != nil { - return fmt.Errorf("run %d: splitting target graph: %w", i+1, err) - } - chunks := make([]entity.GetTargetGraphResponse, 0, len(targetGroups)) - for _, g := range targetGroups { - chunks = append(chunks, entity.GetTargetGraphResponse{Targets: g}) - } - metaGroups, err := streaming.SplitMetadata( - meta.TargetIDMapping, - meta.RuleTypeMapping, - meta.TagMapping, - meta.AttributeNameMapping, - meta.AttributeStringValueMapping, - config.DefaultMaxMessageBytes, - ) - if err != nil { - return fmt.Errorf("run %d: splitting target graph metadata: %w", i+1, err) - } - for _, m := range metaGroups { - chunks = append(chunks, entity.GetTargetGraphResponse{Metadata: m}) + return fmt.Errorf("run %d: converting to graph chunks: %w", i+1, err) } elapsed = time.Since(start) - fmt.Printf("run %d: ResultToTargetGraph+Split: %v (%d chunks)\n", i+1, elapsed.Round(time.Millisecond), len(chunks)) + fmt.Printf("run %d: ResultToGraphChunks: %v (%d chunks)\n", i+1, elapsed.Round(time.Millisecond), len(chunks)) m := jsonpb.Marshaler{Indent: " "} for _, chunk := range chunks { protoResp := mapper.GetTargetGraphResponseToProto(&chunk) diff --git a/mapper/BUILD.bazel b/mapper/BUILD.bazel index ea640a95..0d04e656 100644 --- a/mapper/BUILD.bazel +++ b/mapper/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "//core/targethasher", "//entity", "//internal/mapper/idmapper", + "//internal/streaming", "@com_github_bazelbuild_buildtools//build_proto", ], ) diff --git a/mapper/mapper.go b/mapper/mapper.go index 92fa0f54..e5db35e6 100644 --- a/mapper/mapper.go +++ b/mapper/mapper.go @@ -8,6 +8,7 @@ import ( "github.com/uber/tango/core/targethasher" "github.com/uber/tango/entity" "github.com/uber/tango/internal/mapper/idmapper" + "github.com/uber/tango/internal/streaming" ) const cancelCheckInterval = 4096 @@ -95,3 +96,46 @@ func ResultToTargetGraph(ctx context.Context, result targethasher.Result) ([]ent return targets, meta, nil } + +// ResultToGraphChunks converts a targethasher.Result into the ordered +// sequence of GetTargetGraphResponse chunks ready to stream or persist. It +// maps the result to entity types (ResultToTargetGraph), then splits both the +// targets and the accompanying metadata so each chunk's serialized size stays +// at or under maxBytes. Target chunks come first, followed by metadata chunks; +// consumers merge metadata across chunks before resolving IDs. +// +// Callers that need to interleave or size chunks differently can use +// ResultToTargetGraph directly; this is the convenience path for the common +// "map, split, hand to WriteGraphStream" flow. +func ResultToGraphChunks(ctx context.Context, result targethasher.Result, maxBytes int) ([]entity.GetTargetGraphResponse, error) { + targets, meta, err := ResultToTargetGraph(ctx, result) + if err != nil { + return nil, err + } + + targetGroups, err := streaming.SplitBySize(targets, maxBytes) + if err != nil { + return nil, err + } + chunks := make([]entity.GetTargetGraphResponse, 0, len(targetGroups)) + for _, g := range targetGroups { + chunks = append(chunks, entity.GetTargetGraphResponse{Targets: g}) + } + + metaGroups, err := streaming.SplitMetadata( + meta.TargetIDMapping, + meta.RuleTypeMapping, + meta.TagMapping, + meta.AttributeNameMapping, + meta.AttributeStringValueMapping, + maxBytes, + ) + if err != nil { + return nil, err + } + for _, m := range metaGroups { + chunks = append(chunks, entity.GetTargetGraphResponse{Metadata: m}) + } + + return chunks, nil +} diff --git a/mapper/mapper_test.go b/mapper/mapper_test.go index 3641f7ea..96aa19e8 100644 --- a/mapper/mapper_test.go +++ b/mapper/mapper_test.go @@ -1,6 +1,7 @@ package mapper import ( + "context" "testing" "github.com/stretchr/testify/assert" @@ -17,3 +18,60 @@ func TestResultToTargetGraph_EmptyResult(t *testing.T) { assert.Empty(t, targets) assert.NotNil(t, meta) } + +func TestResultToGraphChunks(t *testing.T) { + t.Parallel() + + result := targethasher.Result{ + TargetNames: []string{"//a:a", "//b:b"}, + Targets: map[string]*targethasher.Target{ + "//a:a": {Hash: []byte{0x01}, RuleType: "go_library", Deps: []string{"//b:b"}}, + "//b:b": {Hash: []byte{0x02}, RuleType: "go_library"}, + }, + } + + t.Run("single chunk when under budget", func(t *testing.T) { + t.Parallel() + + chunks, err := ResultToGraphChunks(t.Context(), result, 1<<20) + require.NoError(t, err) + + var targets, metas int + for _, c := range chunks { + targets += len(c.Targets) + if c.Metadata != nil { + metas++ + } + } + assert.Equal(t, 2, targets) + assert.Positive(t, metas, "expected at least one metadata chunk") + }) + + t.Run("splits targets across chunks under a tiny budget", func(t *testing.T) { + t.Parallel() + + chunks, err := ResultToGraphChunks(t.Context(), result, 1) + require.NoError(t, err) + + targetChunks := 0 + total := 0 + for _, c := range chunks { + if len(c.Targets) > 0 { + targetChunks++ + total += len(c.Targets) + } + } + assert.Equal(t, 2, total) + assert.Greater(t, targetChunks, 1, "tiny budget should spread targets over multiple chunks") + }) + + t.Run("propagates context cancellation", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err := ResultToGraphChunks(ctx, result, 1<<20) + require.Error(t, err) + }) +} diff --git a/orchestrator/BUILD.bazel b/orchestrator/BUILD.bazel index 5e211f9c..eb48dff0 100644 --- a/orchestrator/BUILD.bazel +++ b/orchestrator/BUILD.bazel @@ -21,7 +21,6 @@ go_library( "//core/workspace", "//entity", "//graphrunner", - "//internal/streaming", "//internal/url", "//mapper", "//observability/metrics", diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index 509e2fcd..797e3e72 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -33,7 +33,6 @@ import ( "github.com/uber/tango/core/workspace" "github.com/uber/tango/entity" "github.com/uber/tango/graphrunner" - "github.com/uber/tango/internal/streaming" "github.com/uber/tango/internal/url" "github.com/uber/tango/mapper" "github.com/uber/tango/observability/metrics" @@ -213,32 +212,10 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, req entity.GetT if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - targets, meta, err := mapper.ResultToTargetGraph(ctx, result) + chunks, err := mapper.ResultToGraphChunks(ctx, result, b.config.Service.MaxMessageBytes) if err != nil { return nil, fmt.Errorf("convert target graph: %w", err) } - targetGroups, err := streaming.SplitBySize(targets, b.config.Service.MaxMessageBytes) - if err != nil { - return nil, fmt.Errorf("split target graph: %w", err) - } - chunks := make([]entity.GetTargetGraphResponse, 0, len(targetGroups)) - for _, g := range targetGroups { - chunks = append(chunks, entity.GetTargetGraphResponse{Targets: g}) - } - metaGroups, err := streaming.SplitMetadata( - meta.TargetIDMapping, - meta.RuleTypeMapping, - meta.TagMapping, - meta.AttributeNameMapping, - meta.AttributeStringValueMapping, - b.config.Service.MaxMessageBytes, - ) - if err != nil { - return nil, fmt.Errorf("split target graph metadata: %w", err) - } - for _, m := range metaGroups { - chunks = append(chunks, entity.GetTargetGraphResponse{Metadata: m}) - } cacheWriteStart := time.Now() err = storage.WriteGraphStream(ctx, b.storage, treehashPath, chunks) if err != nil {