diff --git a/config/config.go b/config/config.go index 77ca4939..083695dc 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..9b06e84f 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/controller/BUILD.bazel b/controller/BUILD.bazel index 9f25e63d..743c9a66 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -17,12 +17,12 @@ go_library( deps = [ "//config", "//core/cachekey", - "//core/common", "//core/errors", "//core/storage", "//entity", "//internal/mapper", "//internal/mapper/idmapper", + "//internal/streaming", "//internal/targetdiff", "//internal/url", "//observability/metrics", @@ -47,7 +47,7 @@ go_test( ], embed = [":controller"], deps = [ - "//core/common", + "//config", "//core/storage", "//core/storage/storagemock", "//entity", @@ -55,7 +55,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 2ff4ce9b..b6e1e41d 100644 --- a/controller/controller.go +++ b/controller/controller.go @@ -19,7 +19,6 @@ import ( "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/observability/metrics" "github.com/uber/tango/orchestrator" @@ -31,21 +30,19 @@ 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"` } type controller struct { - logger *zap.Logger - storage storage.Storage - orchestrator orchestrator.Orchestrator - emitter *metrics.Emitter - targetChunkSize int - changedTargetChunkSize int - metadataMapChunkSize int + logger *zap.Logger + storage storage.Storage + orchestrator orchestrator.Orchestrator + emitter *metrics.Emitter + maxMessageBytes int // appCtx is the application lifetime; cancel it on process shutdown. // Used by linkRequestCtx and any fire-and-forget goroutines so they @@ -57,27 +54,17 @@ type controller struct { // shutdown to abort background work. func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { emitter := metrics.New(p.Scope).SubScope("controller") - 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 = config.DefaultMaxMessageBytes } return &controller{ - logger: p.Logger, - storage: p.Storage, - orchestrator: p.Orchestrator, - emitter: emitter, - targetChunkSize: targetChunkSize, - changedTargetChunkSize: changedTargetChunkSize, - metadataMapChunkSize: metadataMapChunkSize, - appCtx: appCtx, + logger: p.Logger, + storage: p.Storage, + orchestrator: p.Orchestrator, + emitter: emitter, + maxMessageBytes: maxMessageBytes, + appCtx: appCtx, } } diff --git a/controller/getchangedtargets.go b/controller/getchangedtargets.go index 3e454218..087f66aa 100644 --- a/controller/getchangedtargets.go +++ b/controller/getchangedtargets.go @@ -22,12 +22,12 @@ import ( "time" "github.com/uber/tango/core/cachekey" - "github.com/uber/tango/core/common" tangoerrors "github.com/uber/tango/core/errors" "github.com/uber/tango/core/storage" "github.com/uber/tango/entity" "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" "github.com/uber/tango/internal/url" "github.com/uber/tango/observability/metrics" @@ -37,12 +37,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 @@ -158,7 +158,7 @@ func (c *controller) serveChangedTargetsFromCache(ctx context.Context, e *metric // 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 { @@ -166,7 +166,7 @@ func (c *controller) serveChangedTargetsFromCache(ctx context.Context, e *metric // Client gave up while we were draining the cache. Surface as a user-cancelled error. return false, fmt.Errorf("cache reader: %w", err) } - var resp *pb.GetChangedTargetsResponse + var resp entity.GetChangedTargetsResponse resp, readErr = cachedReader.Read() if readErr == io.EOF { readErr = nil @@ -206,7 +206,7 @@ func (c *controller) serveChangedTargetsFromCache(ctx context.Context, e *metric // 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, e *metrics.Emitter, logger *zap.Logger, request *pb.GetChangedTargetsRequest) ([]*pb.GetTargetGraphResponse, []*pb.GetTargetGraphResponse, error) { +func (c *controller) fetchTargetGraphs(ctx context.Context, e *metrics.Emitter, 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 @@ -219,7 +219,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, e *metrics.Emitter, 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)) @@ -256,8 +256,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, e *metrics.Emitter, } defer func() { _ = 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 { @@ -276,7 +275,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, e *metrics.Emitter, // 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 { @@ -323,11 +322,11 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, e *metrics.Emitter, 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 } @@ -335,7 +334,7 @@ func (c *controller) fetchTargetGraphs(ctx context.Context, e *metrics.Emitter, // 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 @@ -372,7 +371,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, e *metrics.Emitter, logger *zap.Logger, firstGraph, secondGraph []*pb.GetTargetGraphResponse, maxDist int32) (_ []*pb.GetChangedTargetsResponse, retErr error) { +func (c *controller) compareTargetGraphs(ctx context.Context, e *metrics.Emitter, logger *zap.Logger, firstGraph, secondGraph []entity.GetTargetGraphResponse, maxDist int32) (_ []entity.GetChangedTargetsResponse, retErr error) { op := metrics.Begin(e, opCompareTargetGraphs, slowDurationBuckets) defer func() { op.Complete(retErr) }() logger.Info("compareTargetGraphs: Computing differences between target graphs") @@ -429,9 +428,9 @@ func (c *controller) compareTargetGraphs(ctx context.Context, e *metrics.Emitter // 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), @@ -439,41 +438,27 @@ func (c *controller) compareTargetGraphs(ctx context.Context, e *metrics.Emitter }) } - // 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], - }, - }, - }) + changedGroups, err := streaming.SplitBySize(changed, c.maxMessageBytes) + if err != nil { + return nil, err } - if len(results) == 0 { - results = append(results, &pb.GetChangedTargetsResponse{ - Item: &pb.GetChangedTargetsResponse_ChangedTargets{ - ChangedTargets: &pb.ChangedTargets{}, - }, - }) + results := make([]entity.GetChangedTargetsResponse, 0, len(changedGroups)) + for _, g := range changedGroups { + results = append(results, entity.GetChangedTargetsResponse{ChangedTargets: g}) } - 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 _, m := range metaGroups { + results = append(results, entity.GetChangedTargetsResponse{Metadata: m}) } logger.Info("GetChangedTargets: Target graphs compared") return results, nil @@ -485,10 +470,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), @@ -498,26 +483,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 } } @@ -528,12 +511,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 @@ -548,12 +531,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 != "" { @@ -561,7 +544,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 != "" { @@ -569,7 +552,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 != "" { @@ -605,14 +588,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, @@ -644,31 +627,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 entity equivalent. +func toChangeType(ct targetdiff.ChangeType) entity.ChangeType { switch ct { case targetdiff.ChangeTypeNew: - return pb.CHANGE_TYPE_NEW + return entity.ChangeTypeNew case targetdiff.ChangeTypeDeleted: - return pb.CHANGE_TYPE_DELETED + return entity.ChangeTypeDeleted case targetdiff.ChangeTypeChanged: - return pb.CHANGE_TYPE_CHANGED + return entity.ChangeTypeChanged default: - return pb.CHANGE_TYPE_INVALID + return entity.ChangeTypeInvalid } } // 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 0c203c80..a62fa354 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/storage" @@ -135,18 +135,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.emitter, zap.NewNop(), []*pb.GetTargetGraphResponse{firstGraph}, []*pb.GetTargetGraphResponse{secondGraph}, -1) + response, err := c.compareTargetGraphs(t.Context(), c.emitter, zap.NewNop(), []entity.GetTargetGraphResponse{firstGraph}, []entity.GetTargetGraphResponse{secondGraph}, -1) require.NoError(t, err) require.NotNil(t, response) } @@ -166,21 +158,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) @@ -292,17 +275,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. @@ -349,46 +333,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() @@ -473,20 +445,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() @@ -574,180 +540,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.emitter, 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, entity.ChangeTypeNew, 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.emitter, 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, entity.ChangeTypeChanged, aCT.ChangeType) + require.Equal(t, entity.ChangeTypeChanged, 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"}, }, }, } @@ -756,262 +693,224 @@ func TestCompareTargetGraphs_ChangedRuleUnreachableFromAnySeed(t *testing.T) { // no upstream explanation becomes a distance-0 seed itself. res, err := c.compareTargetGraphs(t.Context(), c.emitter, 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, entity.ChangeTypeChanged, 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.emitter, 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, entity.ChangeTypeChanged, 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.emitter, 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, entity.ChangeTypeChanged, 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.emitter, 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, entity.ChangeTypeChanged, 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: entity.ChangeTypeChanged}, }, }, { - Item: &pb.GetChangedTargetsResponse_Metadata{Metadata: meta}, + Metadata: &entity.Metadata{TargetIDMapping: map[int32]string{1: "//app:T"}}, }, } @@ -1027,18 +926,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{}, }, } @@ -1054,23 +951,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: entity.ChangeTypeChanged}, + {Distance: 2, ChangeType: entity.ChangeTypeChanged}, + }, + }) + enc.Encode(entity.GetChangedTargetsResponse{Metadata: &entity.Metadata{}}) cachedBytes := buf.Bytes() storagemock := storagemock.NewMockStorage(ctrl) @@ -1116,93 +1005,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.emitter, 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, entity.ChangeTypeChanged, 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) { @@ -1212,128 +1092,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.emitter, 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.emitter, 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, entity.ChangeTypeDeleted, 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) { @@ -1341,16 +1202,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: entity.ChangeTypeDeleted}, + {Distance: 1, ChangeType: entity.ChangeTypeChanged}, + {Distance: 5, ChangeType: entity.ChangeTypeChanged}, }, }, } @@ -1375,21 +1232,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"}, @@ -1418,14 +1260,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( @@ -1456,13 +1300,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) @@ -1499,16 +1339,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)) @@ -1529,7 +1367,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)) @@ -1547,7 +1385,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)) @@ -1574,7 +1412,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, @@ -1585,8 +1423,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 0988ff69..ffcfd7ac 100644 --- a/controller/gettargetgraph.go +++ b/controller/gettargetgraph.go @@ -66,7 +66,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) logger.Info("GetTargetGraph: Done streaming", @@ -79,7 +79,8 @@ func (c *controller) GetTargetGraph(request *pb.GetTargetGraphRequest, stream pb if err != nil { return fmt.Errorf("graph reader read: %w", err) } - toSend := applyOptimizedTargetsOutputConfigToChunk(graphStreamChunk, outputConfig) + protoResp := mapper.GetTargetGraphResponseToProto(&chunk) + toSend := applyOptimizedTargetsOutputConfigToChunk(protoResp, outputConfig) err = stream.Send(toSend) if err != nil { return fmt.Errorf("send graph: %w", err) diff --git a/controller/gettargetgraph_test.go b/controller/gettargetgraph_test.go index e8e72e91..495b8a9f 100644 --- a/controller/gettargetgraph_test.go +++ b/controller/gettargetgraph_test.go @@ -21,11 +21,13 @@ 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/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" @@ -120,10 +122,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), @@ -133,7 +132,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", @@ -190,16 +189,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), @@ -277,10 +267,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( @@ -291,7 +278,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) @@ -309,12 +296,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 18c04c07..d44b9dc6 100644 --- a/controller/testhelper_test.go +++ b/controller/testhelper_test.go @@ -16,19 +16,32 @@ package controller import ( "context" + "testing" - "github.com/uber/tango/core/common" + "github.com/stretchr/testify/require" + "github.com/uber/tango/config" + "github.com/uber/tango/core/storage" + "github.com/uber/tango/entity" "github.com/uber/tango/observability/metrics" "go.uber.org/zap" ) func newTestController(logger *zap.Logger) *controller { return &controller{ - logger: logger, - emitter: metrics.Nop(), - targetChunkSize: common.DefaultTargetChunkSize, - changedTargetChunkSize: common.DefaultChangedTargetChunkSize, - metadataMapChunkSize: common.DefaultMetadataMapChunkSize, - appCtx: context.Background(), + logger: logger, + emitter: metrics.Nop(), + maxMessageBytes: config.DefaultMaxMessageBytes, + 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(t.Context(), st, "test-graph", chunks)) + reader, err := storage.NewGraphReader(t.Context(), st, "test-graph") + require.NoError(t, err) + return reader +} diff --git a/core/common/BUILD.bazel b/core/common/BUILD.bazel deleted file mode 100644 index 7497c3b8..00000000 --- a/core/common/BUILD.bazel +++ /dev/null @@ -1,26 +0,0 @@ -load("@rules_go//go:def.bzl", "go_library", "go_test") - -go_library( - name = "common", - srcs = ["utils.go"], - 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", - ], -) diff --git a/core/common/utils.go b/core/common/utils.go deleted file mode 100644 index 2f88ca4c..00000000 --- a/core/common/utils.go +++ /dev/null @@ -1,257 +0,0 @@ -// 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. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package common - -import ( - "context" - "encoding/hex" - - 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 -) - -// 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 deleted file mode 100644 index da6f0d9f..00000000 --- a/core/common/utils_test.go +++ /dev/null @@ -1,116 +0,0 @@ -// 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. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -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 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/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..15aed932 100644 --- a/core/storage/changedtargetsreader.go +++ b/core/storage/changedtargetsreader.go @@ -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..0c9bd129 100644 --- a/core/storage/reader.go +++ b/core/storage/reader.go @@ -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/entity/changed_targets.go b/entity/changed_targets.go index 0fa33533..5a61848c 100644 --- a/entity/changed_targets.go +++ b/entity/changed_targets.go @@ -1,5 +1,7 @@ package entity +import "github.com/uber/tango/internal/streaming/wire" + // ChangedTarget represents a target that differs between two revisions. type ChangedTarget struct { ChangeType ChangeType `json:"change_type"` @@ -8,6 +10,26 @@ type ChangedTarget struct { Distance int32 `json:"distance"` } +// Size returns an estimate of the protobuf wire size for this changed target. +func (ct ChangedTarget) Size() int { + n := 0 + if ct.ChangeType != 0 { + n += 1 + wire.VarintSize(uint64(ct.ChangeType)) + } + if ct.OldTarget != nil { + inner := ct.OldTarget.Size() + n += 1 + wire.VarintSize(uint64(inner)) + inner + } + if ct.NewTarget != nil { + inner := ct.NewTarget.Size() + n += 1 + wire.VarintSize(uint64(inner)) + inner + } + if ct.Distance != 0 { + n += 1 + wire.VarintSize(uint64(ct.Distance)) + } + return n +} + // 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. diff --git a/entity/optimized_target.go b/entity/optimized_target.go index 30194dfd..dc78dab1 100644 --- a/entity/optimized_target.go +++ b/entity/optimized_target.go @@ -18,7 +18,7 @@ type OptimizedTarget struct { // 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 { +func (t OptimizedTarget) Size() int { n := 0 if t.ID != 0 { n += 1 + wire.VarintSize(uint64(t.ID)) diff --git a/example/cmd/query-bench/BUILD.bazel b/example/cmd/query-bench/BUILD.bazel index ba71ce1c..d5f616c0 100644 --- a/example/cmd/query-bench/BUILD.bazel +++ b/example/cmd/query-bench/BUILD.bazel @@ -6,9 +6,13 @@ go_library( importpath = "github.com/uber/tango/example/cmd/query-bench", visibility = ["//visibility:private"], deps = [ + "//config", "//core/bazel", - "//core/common", "//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 28e4e52b..85c32d28 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -31,9 +31,13 @@ import ( "time" "github.com/gogo/protobuf/jsonpb" + "github.com/uber/tango/config" "github.com/uber/tango/core/bazel" - "github.com/uber/tango/core/common" "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" ) @@ -112,18 +116,38 @@ 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 := tgmapper.ResultToTargetGraph(ctx, targethasherResult) + 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: converting to GetTargetGraphResponse: %w", i+1, err) + return fmt.Errorf("run %d: splitting target graph metadata: %w", i+1, err) + } + for _, m := range metaGroups { + chunks = append(chunks, entity.GetTargetGraphResponse{Metadata: m}) } 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+Split: %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() diff --git a/internal/streaming/streaming.go b/internal/streaming/streaming.go index 84203a8d..596bf3d4 100644 --- a/internal/streaming/streaming.go +++ b/internal/streaming/streaming.go @@ -117,42 +117,3 @@ 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 index 3e7636fb..fd2b80a2 100644 --- a/internal/streaming/streaming_test.go +++ b/internal/streaming/streaming_test.go @@ -113,47 +113,3 @@ func TestSplitMetadata_EmptyBigMapsStillCarrySmallMaps(t *testing.T) { 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/orchestrator/BUILD.bazel b/orchestrator/BUILD.bazel index 9baa3b5d..5e211f9c 100644 --- a/orchestrator/BUILD.bazel +++ b/orchestrator/BUILD.bazel @@ -14,7 +14,6 @@ go_library( "//config", "//core/bazel", "//core/cachekey", - "//core/common", "//core/errors", "//core/git", "//core/repomanager", @@ -22,7 +21,9 @@ go_library( "//core/workspace", "//entity", "//graphrunner", + "//internal/streaming", "//internal/url", + "//mapper", "//observability/metrics", "@com_github_uber_go_tally//:tally", "@org_uber_go_zap//:zap", @@ -50,8 +51,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 e62c2b15..26d744c7 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -27,14 +27,15 @@ import ( "github.com/uber/tango/config" "github.com/uber/tango/core/bazel" "github.com/uber/tango/core/cachekey" - "github.com/uber/tango/core/common" "github.com/uber/tango/core/git" "github.com/uber/tango/core/repomanager" "github.com/uber/tango/core/storage" "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" "go.uber.org/zap" ) @@ -212,15 +213,34 @@ 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: %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("convert target graph to response: %w", err) + 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, responses) + 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 03d7e95c..c5e360f7 100644 --- a/orchestrator/native_orchestrator_test.go +++ b/orchestrator/native_orchestrator_test.go @@ -22,7 +22,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" @@ -37,7 +38,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" ) @@ -48,10 +48,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())), @@ -82,13 +79,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) { @@ -104,9 +99,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) @@ -139,9 +132,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) { @@ -175,10 +168,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) @@ -213,9 +203,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