feat(streaming): add streaming split utilities for gRPC message sizing

entity/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@rules_go//go:def.bzl", "go_library")
+load("@rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "entity",
@@ -12,4 +12,12 @@ go_library(
     ],
     importpath = "github.com/uber/tango/entity",
     visibility = ["//visibility:public"],
+    deps = ["//internal/streaming/wire"],
+)
+
+go_test(
+    name = "entity_test",
+    srcs = ["optimized_target_test.go"],
+    embed = [":entity"],
+    deps = ["@com_github_stretchr_testify//assert"],
 )

entity/optimized_target.go

@@ -1,5 +1,7 @@
 package entity
 
+import "github.com/uber/tango/internal/streaming/wire"
+
 // OptimizedTarget is the compact, ID-mapped representation of a target used
 // for streaming and storage. String fields are replaced with int32 IDs that
 // reference the accompanying Metadata maps.
@@ -14,6 +16,46 @@ type OptimizedTarget struct {
 	Attributes         map[int32]int32 `json:"attributes"`
 }
 
+// Size returns an estimate of the protobuf wire size for this target.
+// Used by streaming splitters to stay within gRPC message size limits.
+func (t *OptimizedTarget) Size() int {
+	n := 0
+	if t.ID != 0 {
+		n += 1 + wire.VarintSize(uint64(t.ID))
+	}
+	if len(t.Hash) > 0 {
+		n += 1 + wire.VarintSize(uint64(len(t.Hash))) + len(t.Hash)
+	}
+	if len(t.DirectDependencies) > 0 {
+		dataSize := 0
+		for _, d := range t.DirectDependencies {
+			dataSize += wire.VarintSize(uint64(d))
+		}
+		n += 1 + wire.VarintSize(uint64(dataSize)) + dataSize
+	}
+	if t.RuleType != 0 {
+		n += 1 + wire.VarintSize(uint64(t.RuleType))
+	}
+	if len(t.Tags) > 0 {
+		dataSize := 0
+		for _, tag := range t.Tags {
+			dataSize += wire.VarintSize(uint64(tag))
+		}
+		n += 1 + wire.VarintSize(uint64(dataSize)) + dataSize
+	}
+	if t.Root {
+		n += 1 + 1
+	}
+	if t.External {
+		n += 1 + 1
+	}
+	for k, v := range t.Attributes {
+		entrySize := 1 + wire.VarintSize(uint64(k)) + 1 + wire.VarintSize(uint64(v))
+		n += 1 + wire.VarintSize(uint64(entrySize)) + entrySize
+	}
+	return n
+}
+
 // Metadata holds the ID-to-string mappings that accompany a set of
 // OptimizedTarget entries. Consumers merge metadata across chunks before
 // resolving IDs.

entity/optimized_target_test.go

@@ -0,0 +1,74 @@
+package entity
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestOptimizedTarget_Size(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		target   OptimizedTarget
+		wantSize int
+	}{
+		{
+			name:     "zero value",
+			target:   OptimizedTarget{},
+			wantSize: 0,
+		},
+		{
+			name:   "id only",
+			target: OptimizedTarget{ID: 1},
+			// field tag (1 byte) + varint(1) (1 byte) = 2
+			wantSize: 2,
+		},
+		{
+			name:   "with hash",
+			target: OptimizedTarget{ID: 42, Hash: "abcd"},
+			// id: tag(1) + varint(42)(1) = 2
+			// hash: tag(1) + len_prefix(1) + 4 bytes = 6
+			wantSize: 8,
+		},
+		{
+			name:   "with deps packed",
+			target: OptimizedTarget{ID: 1, DirectDependencies: []int32{2, 3}},
+			// id: tag(1) + varint(1)(1) = 2
+			// deps packed: tag(1) + len_prefix(1) + varint(2)(1) + varint(3)(1) = 4
+			wantSize: 6,
+		},
+		{
+			name:   "booleans",
+			target: OptimizedTarget{ID: 1, Root: true, External: true},
+			// id: 2, root: tag(1) + 1 = 2, external: tag(1) + 1 = 2
+			wantSize: 6,
+		},
+		{
+			name:     "false booleans omitted",
+			target:   OptimizedTarget{ID: 1, Root: false, External: false},
+			wantSize: 2,
+		},
+		{
+			name:   "with attributes",
+			target: OptimizedTarget{ID: 1, Attributes: map[int32]int32{1: 2}},
+			// id: 2
+			// map entry: tag(1) + len_prefix(1) + [key: tag(1)+varint(1)(1) + val: tag(1)+varint(2)(1)] = 6
+			wantSize: 8,
+		},
+		{
+			name:   "tags packed",
+			target: OptimizedTarget{Tags: []int32{5, 6}},
+			// tags packed: tag(1) + len_prefix(1) + varint(5)(1) + varint(6)(1) = 4
+			wantSize: 4,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			assert.Equal(t, tt.wantSize, tt.target.Size())
+		})
+	}
+}

internal/streaming/BUILD.bazel

@@ -0,0 +1,24 @@
+load("@rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "streaming",
+    srcs = ["streaming.go"],
+    importpath = "github.com/uber/tango/internal/streaming",
+    visibility = ["//:__subpackages__"],
+    deps = [
+        "//entity",
+        "//internal/streaming/wire",
+    ],
+)
+
+go_test(
+    name = "streaming_test",
+    srcs = ["streaming_test.go"],
+    embed = [":streaming"],
+    deps = [
+        "//entity",
+        "//tangopb",
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//require",
+    ],
+)

internal/streaming/streaming.go

@@ -0,0 +1,158 @@
+package streaming
+
+import (
+	"fmt"
+
+	"github.com/uber/tango/entity"
+	"github.com/uber/tango/internal/streaming/wire"
+)
+
+// Sizer is satisfied by any type that reports its serialized byte length.
+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 + wire.VarintSize(uint64(k)) + 1 + len(v) + wire.VarintSize(uint64(len(v)))
+	return mapEntrySize + 1 + wire.VarintSize(uint64(mapEntrySize))
+}
+
+// SplitTargetGraph splits targets and metadata into wire-safe
+// entity.GetTargetGraphResponse chunks bounded by maxMessageBytes.
+func SplitTargetGraph(targets []entity.OptimizedTarget, meta *entity.Metadata, maxMessageBytes int) ([]entity.GetTargetGraphResponse, error) {
+	var chunks []entity.GetTargetGraphResponse
+	start := 0
+	currentBytes := 0
+	for i := range targets {
+		itemBytes := targets[i].Size()
+		if i > start && currentBytes+itemBytes > maxMessageBytes {
+			chunks = append(chunks, entity.GetTargetGraphResponse{
+				Targets: targets[start:i],
+			})
+			start = i
+			currentBytes = 0
+		}
+		currentBytes += itemBytes
+	}
+	chunks = append(chunks, entity.GetTargetGraphResponse{
+		Targets: targets[start:],
+	})
+
+	metaGroups, err := SplitMetadata(
+		meta.TargetIDMapping,
+		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
+}