diff --git a/entity/BUILD.bazel b/entity/BUILD.bazel index 3e9b3a9a..f6c53680 100644 --- a/entity/BUILD.bazel +++ b/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"], ) diff --git a/entity/optimized_target.go b/entity/optimized_target.go index 36308ff9..30194dfd 100644 --- a/entity/optimized_target.go +++ b/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. diff --git a/entity/optimized_target_test.go b/entity/optimized_target_test.go new file mode 100644 index 00000000..e6fb0a97 --- /dev/null +++ b/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()) + }) + } +} diff --git a/internal/streaming/BUILD.bazel b/internal/streaming/BUILD.bazel new file mode 100644 index 00000000..304c2769 --- /dev/null +++ b/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", + ], +) diff --git a/internal/streaming/streaming.go b/internal/streaming/streaming.go new file mode 100644 index 00000000..84203a8d --- /dev/null +++ b/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 +} diff --git a/internal/streaming/streaming_test.go b/internal/streaming/streaming_test.go new file mode 100644 index 00000000..3e7636fb --- /dev/null +++ b/internal/streaming/streaming_test.go @@ -0,0 +1,159 @@ +package streaming + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/tango/entity" + pb "github.com/uber/tango/tangopb" +) + +func TestSplitBySize(t *testing.T) { + t.Parallel() + + targets := make([]*entity.OptimizedTarget, 25) + for i := range targets { + targets[i] = &entity.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 := &entity.OptimizedTarget{ID: 1, Hash: strings.Repeat("a", 1000)} + small := &entity.OptimizedTarget{ID: 2} + maxBytes := small.Size() + + groups, err := SplitBySize([]*entity.OptimizedTarget{oversized, small}, maxBytes) + require.NoError(t, err) + require.Len(t, groups, 2) + assert.Equal(t, []*entity.OptimizedTarget{oversized}, groups[0]) + assert.Equal(t, []*entity.OptimizedTarget{small}, groups[1]) +} + +func TestSplitBySize_EmptyInputReturnsOneEmptyGroup(t *testing.T) { + t.Parallel() + + groups, err := SplitBySize([]*entity.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) +} + +func TestSplitTargetGraph(t *testing.T) { + t.Parallel() + + numTargets := 50 + targets := make([]entity.OptimizedTarget, numTargets) + for i := range targets { + targets[i] = entity.OptimizedTarget{ID: int32(i + 1), Hash: "ab", RuleType: 1} + } + meta := &entity.Metadata{ + TargetIDMapping: map[int32]string{1: "//pkg:a"}, + RuleTypeMapping: map[int32]string{1: "go_library"}, + } + + targetSize := targets[0].Size() + + tests := []struct { + name string + maxMessageBytes int + wantTargetChunks int + }{ + {name: "25 per chunk", maxMessageBytes: targetSize * 25, wantTargetChunks: 2}, + {name: "10 per chunk", maxMessageBytes: targetSize * 10, wantTargetChunks: 5}, + {name: "all in one", maxMessageBytes: targetSize * 100, wantTargetChunks: 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + chunks, err := SplitTargetGraph(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) + }) + } +} diff --git a/internal/streaming/wire/BUILD.bazel b/internal/streaming/wire/BUILD.bazel new file mode 100644 index 00000000..11345314 --- /dev/null +++ b/internal/streaming/wire/BUILD.bazel @@ -0,0 +1,8 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "wire", + srcs = ["wire.go"], + importpath = "github.com/uber/tango/internal/streaming/wire", + visibility = ["//:__subpackages__"], +) diff --git a/internal/streaming/wire/wire.go b/internal/streaming/wire/wire.go new file mode 100644 index 00000000..f452c7a4 --- /dev/null +++ b/internal/streaming/wire/wire.go @@ -0,0 +1,11 @@ +package wire + +// VarintSize returns the number of bytes needed to encode x as a protobuf varint. +func VarintSize(x uint64) int { + n := 1 + for x >= 0x80 { + x >>= 7 + n++ + } + return n +}