diff --git a/core/common/utils.go b/core/common/utils.go index 90df8021..4821608c 100644 --- a/core/common/utils.go +++ b/core/common/utils.go @@ -137,8 +137,11 @@ func HashRequestOptions(opts *tangopb.RequestOptions) string { // for typical target rates. const cancelCheckInterval = 4096 -// ResultToGetTargetGraphResponse converts a Result to a GetTargetGraphResponse -func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result) ([]*tangopb.GetTargetGraphResponse, error) { +// 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. @@ -224,14 +227,14 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res attrStrValIDToVal := attrStrValMapper.Invert() // chunk targets into multiple messages for streaming - responses := chunkTargets(optimizedTargets, DefaultTargetChunkSize) + responses := chunkTargets(optimizedTargets, targetChunkSize) for _, meta := range ChunkMetadata( targetIDToName, ruleTypeIDToName, tagIDToName, attrNameIDToName, attrStrValIDToVal, - DefaultMetadataMapChunkSize, + metadataMapChunkSize, ) { responses = append(responses, &tangopb.GetTargetGraphResponse{ Item: &tangopb.GetTargetGraphResponse_Metadata{Metadata: meta}, diff --git a/core/common/utils_test.go b/core/common/utils_test.go index 4172da07..3ac8846a 100644 --- a/core/common/utils_test.go +++ b/core/common/utils_test.go @@ -161,13 +161,13 @@ func TestGetComparedTargetsCachePath(t *testing.T) { func TestChunkTargets(t *testing.T) { t.Parallel() - // Create 250 targets, chunk by 100 → expect 3 chunks (100, 100, 50) - targets := make([]*pb.OptimizedTarget, 250) + // 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, 100) + responses := chunkTargets(targets, 10) require.Len(t, responses, 3) @@ -180,14 +180,13 @@ func TestChunkTargets(t *testing.T) { total++ } } - assert.Equal(t, 250, total) + assert.Equal(t, 25, total) } func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { t.Parallel() - // 500 targets with DefaultTargetChunkSize=250 → 2 target chunks + 1 metadata = 3 responses - numTargets := 500 + numTargets := 50 result := targethasher.Result{ TargetNames: make([]string, numTargets), Targets: make(map[string]*targethasher.Target, numTargets), @@ -198,16 +197,53 @@ func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { result.Targets[name] = &targethasher.Target{Name: name, Hash: []byte{0}, RuleType: "go_library"} } - responses, err := ResultToGetTargetGraphResponse(context.Background(), result) - require.NoError(t, err) - - // 2 target chunks + 1 metadata chunk (500 targets well under DefaultMetadataMapChunkSize) - require.Len(t, responses, 3) + 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 _, resp := range responses[:2] { - _, ok := resp.Item.(*pb.GetTargetGraphResponse_Targets) - assert.True(t, ok, "expected Targets chunk") + 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) + }) } - _, ok := responses[2].Item.(*pb.GetTargetGraphResponse_Metadata) - assert.True(t, ok, "last response should be Metadata") } diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 3071b018..28e4e52b 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -112,7 +112,10 @@ 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) // also print the time to convert to GetTargetGraphResponse format + response, err := common.ResultToGetTargetGraphResponse( + ctx, targethasherResult, + common.DefaultTargetChunkSize, common.DefaultMetadataMapChunkSize, + ) if err != nil { return fmt.Errorf("run %d: converting to GetTargetGraphResponse: %w", i+1, err) } diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index dda336e2..794cda19 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -135,9 +135,9 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, param GetTarget } } }() - err = ws.Checkout(ctx, param.Req.BuildDescription.Remote, param.Req.BuildDescription.BaseSha) + err = ws.Checkout(ctx, remote, param.Req.BuildDescription.BaseSha) if err != nil { - return nil, fmt.Errorf("checkout %s@%s: %w", param.Req.BuildDescription.Remote, param.Req.BuildDescription.BaseSha, err) + return nil, fmt.Errorf("checkout %s@%s: %w", remote, param.Req.BuildDescription.BaseSha, err) } logger.Infow("GetTargetGraph: Checked out base revision") @@ -166,7 +166,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, param GetTarget if err != nil { return nil, fmt.Errorf("compute treehash: %w", err) } - treehashPath := common.GetGraphByTreeHash(param.Req.BuildDescription.Remote, treehash, param.Req.BuildDescription.GetStrategy(), param.Req.GetRequestOptions()) + treehashPath := common.GetGraphByTreeHash(remote, treehash, param.Req.BuildDescription.GetStrategy(), param.Req.GetRequestOptions()) if !param.BypassCache { graphReader, err := storage.NewGraphReader(ctx, b.storage, treehashPath) if err == nil { @@ -206,7 +206,10 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, param GetTarget if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - responses, err := common.ResultToGetTargetGraphResponse(ctx, result) + responses, err := common.ResultToGetTargetGraphResponse(ctx, result, + b.config.Service.Chunking.TargetChunkSize, + b.config.Service.Chunking.MetadataMapChunkSize, + ) if err != nil { return nil, fmt.Errorf("convert target graph to response: %w", err) } diff --git a/orchestrator/testdata/config.yaml b/orchestrator/testdata/config.yaml index f420f74f..34887e08 100644 --- a/orchestrator/testdata/config.yaml +++ b/orchestrator/testdata/config.yaml @@ -11,3 +11,7 @@ repository: service: worker_pool_size: 3 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000