refactor(mapper): add ResultToGraphChunks convenience for the map+split flow

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",

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)

mapper/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
         "//core/targethasher",
         "//entity",
         "//internal/mapper/idmapper",
+        "//internal/streaming",
         "@com_github_bazelbuild_buildtools//build_proto",
     ],
 )

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
+}

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)
+	})
+}

orchestrator/BUILD.bazel

@@ -21,7 +21,6 @@ go_library(
         "//core/workspace",
         "//entity",
         "//graphrunner",
-        "//internal/streaming",
         "//internal/url",
         "//mapper",
         "//observability/metrics",

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 {