From 6c67b876b6869d82c70b2f9d362dc875ef7e4ba5 Mon Sep 17 00:00:00 2001 From: Marc Sanmiquel Date: Mon, 1 Dec 2025 17:38:00 +0100 Subject: [PATCH] feat: Add profile_id_selector for individual profile retrieval --- .../gen/querier/v1/querier.openapi.yaml | 10 + .../gen/query/v1/query.openapi.yaml | 12 +- api/gen/proto/go/querier/v1/querier.pb.go | 18 +- .../proto/go/querier/v1/querier_vtproto.pb.go | 61 ++++++ api/gen/proto/go/query/v1/query.pb.go | 28 ++- api/gen/proto/go/query/v1/query_vtproto.pb.go | 122 ++++++++++++ api/openapiv2/gen/phlare.swagger.json | 14 +- api/querier/v1/querier.proto | 2 + api/query/v1/query.proto | 2 + docs/sources/reference-server-api/index.md | 7 + .../query_select_merge_profile.go | 1 + .../query_select_merge_stacktraces.go | 1 + pkg/querier/querier.go | 3 +- pkg/querybackend/block_reader_test.go | 175 ++++++++++++++++++ pkg/querybackend/query_pprof.go | 11 +- pkg/querybackend/query_tree.go | 11 +- 16 files changed, 463 insertions(+), 15 deletions(-) diff --git a/api/connect-openapi/gen/querier/v1/querier.openapi.yaml b/api/connect-openapi/gen/querier/v1/querier.openapi.yaml index 7dfdf8ea07..2f5f6dca95 100644 --- a/api/connect-openapi/gen/querier/v1/querier.openapi.yaml +++ b/api/connect-openapi/gen/querier/v1/querier.openapi.yaml @@ -1209,6 +1209,16 @@ components: description: Select stack traces that match the provided selector. nullable: true $ref: '#/components/schemas/types.v1.StackTraceSelector' + profileIdSelector: + type: array + examples: + - - 7c9e6679-7425-40de-944b-e07fc1f90ae7 + items: + type: string + examples: + - - 7c9e6679-7425-40de-944b-e07fc1f90ae7 + title: profile_id_selector + description: List of Profile UUIDs to query title: SelectMergeProfileRequest additionalProperties: false querier.v1.SelectMergeSpanProfileRequest: diff --git a/api/connect-openapi/gen/query/v1/query.openapi.yaml b/api/connect-openapi/gen/query/v1/query.openapi.yaml index 7cfffaaec0..b209aa0934 100644 --- a/api/connect-openapi/gen/query/v1/query.openapi.yaml +++ b/api/connect-openapi/gen/query/v1/query.openapi.yaml @@ -461,9 +461,14 @@ components: format: int64 stackTraceSelector: title: stack_trace_selector - description: 'TODO(kolesnikovae): Go PGO options.' nullable: true $ref: '#/components/schemas/types.v1.StackTraceSelector' + profileIdSelector: + type: array + items: + type: string + title: profile_id_selector + description: 'TODO(kolesnikovae): Go PGO options.' title: PprofQuery additionalProperties: false query.v1.PprofReport: @@ -714,6 +719,11 @@ components: title: stack_trace_selector nullable: true $ref: '#/components/schemas/types.v1.StackTraceSelector' + profileIdSelector: + type: array + items: + type: string + title: profile_id_selector title: TreeQuery additionalProperties: false query.v1.TreeReport: diff --git a/api/gen/proto/go/querier/v1/querier.pb.go b/api/gen/proto/go/querier/v1/querier.pb.go index 4ba89f377c..df01abad6f 100644 --- a/api/gen/proto/go/querier/v1/querier.pb.go +++ b/api/gen/proto/go/querier/v1/querier.pb.go @@ -921,8 +921,10 @@ type SelectMergeProfileRequest struct { MaxNodes *int64 `protobuf:"varint,5,opt,name=max_nodes,json=maxNodes,proto3,oneof" json:"max_nodes,omitempty"` // Select stack traces that match the provided selector. StackTraceSelector *v1.StackTraceSelector `protobuf:"bytes,6,opt,name=stack_trace_selector,json=stackTraceSelector,proto3,oneof" json:"stack_trace_selector,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // List of Profile UUIDs to query + ProfileIdSelector []string `protobuf:"bytes,7,rep,name=profile_id_selector,json=profileIdSelector,proto3" json:"profile_id_selector,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *SelectMergeProfileRequest) Reset() { @@ -997,6 +999,13 @@ func (x *SelectMergeProfileRequest) GetStackTraceSelector() *v1.StackTraceSelect return nil } +func (x *SelectMergeProfileRequest) GetProfileIdSelector() []string { + if x != nil { + return x.ProfileIdSelector + } + return nil +} + type SelectSeriesRequest struct { state protoimpl.MessageState `protogen:"open.v1"` // Profile Type ID string in the form @@ -1522,14 +1531,15 @@ const file_querier_v1_querier_proto_rawDesc = "" + "rightTicks\x18\x06 \x01(\x03R\n" + "rightTicks\"\x1f\n" + "\x05Level\x12\x16\n" + - "\x06values\x18\x01 \x03(\x03R\x06values\"\xb4\x03\n" + + "\x06values\x18\x01 \x03(\x03R\x06values\"\x95\x04\n" + "\x19SelectMergeProfileRequest\x12Y\n" + "\x0eprofile_typeID\x18\x01 \x01(\tB2\xbaG/:-\x12+process_cpu:cpu:nanoseconds:cpu:nanosecondsR\rprofileTypeID\x12J\n" + "\x0elabel_selector\x18\x02 \x01(\tB#\xbaG :\x1e\x12\x1c'{namespace=\"my-namespace\"}'R\rlabelSelector\x12*\n" + "\x05start\x18\x03 \x01(\x03B\x14\xbaG\x11:\x0f\x12\r1676282400000R\x05start\x12&\n" + "\x03end\x18\x04 \x01(\x03B\x14\xbaG\x11:\x0f\x12\r1676289600000R\x03end\x12 \n" + "\tmax_nodes\x18\x05 \x01(\x03H\x00R\bmaxNodes\x88\x01\x01\x12S\n" + - "\x14stack_trace_selector\x18\x06 \x01(\v2\x1c.types.v1.StackTraceSelectorH\x01R\x12stackTraceSelector\x88\x01\x01B\f\n" + + "\x14stack_trace_selector\x18\x06 \x01(\v2\x1c.types.v1.StackTraceSelectorH\x01R\x12stackTraceSelector\x88\x01\x01\x12_\n" + + "\x13profile_id_selector\x18\a \x03(\tB/\xbaG,:*\x12(['7c9e6679-7425-40de-944b-e07fc1f90ae7']R\x11profileIdSelectorB\f\n" + "\n" + "_max_nodesB\x17\n" + "\x15_stack_trace_selector\"\xfb\x04\n" + diff --git a/api/gen/proto/go/querier/v1/querier_vtproto.pb.go b/api/gen/proto/go/querier/v1/querier_vtproto.pb.go index 2e0837e93c..5cc0866b97 100644 --- a/api/gen/proto/go/querier/v1/querier_vtproto.pb.go +++ b/api/gen/proto/go/querier/v1/querier_vtproto.pb.go @@ -376,6 +376,11 @@ func (m *SelectMergeProfileRequest) CloneVT() *SelectMergeProfileRequest { r.StackTraceSelector = proto.Clone(rhs).(*v1.StackTraceSelector) } } + if rhs := m.ProfileIdSelector; rhs != nil { + tmpContainer := make([]string, len(rhs)) + copy(tmpContainer, rhs) + r.ProfileIdSelector = tmpContainer + } if len(m.unknownFields) > 0 { r.unknownFields = make([]byte, len(m.unknownFields)) copy(r.unknownFields, m.unknownFields) @@ -1016,6 +1021,15 @@ func (this *SelectMergeProfileRequest) EqualVT(that *SelectMergeProfileRequest) } else if !proto.Equal(this.StackTraceSelector, that.StackTraceSelector) { return false } + if len(this.ProfileIdSelector) != len(that.ProfileIdSelector) { + return false + } + for i, vx := range this.ProfileIdSelector { + vy := that.ProfileIdSelector[i] + if vx != vy { + return false + } + } return string(this.unknownFields) == string(that.unknownFields) } @@ -2548,6 +2562,15 @@ func (m *SelectMergeProfileRequest) MarshalToSizedBufferVT(dAtA []byte) (int, er i -= len(m.unknownFields) copy(dAtA[i:], m.unknownFields) } + if len(m.ProfileIdSelector) > 0 { + for iNdEx := len(m.ProfileIdSelector) - 1; iNdEx >= 0; iNdEx-- { + i -= len(m.ProfileIdSelector[iNdEx]) + copy(dAtA[i:], m.ProfileIdSelector[iNdEx]) + i = protohelpers.EncodeVarint(dAtA, i, uint64(len(m.ProfileIdSelector[iNdEx]))) + i-- + dAtA[i] = 0x3a + } + } if m.StackTraceSelector != nil { if vtmsg, ok := interface{}(m.StackTraceSelector).(interface { MarshalToSizedBufferVT([]byte) (int, error) @@ -3356,6 +3379,12 @@ func (m *SelectMergeProfileRequest) SizeVT() (n int) { } n += 1 + l + protohelpers.SizeOfVarint(uint64(l)) } + if len(m.ProfileIdSelector) > 0 { + for _, s := range m.ProfileIdSelector { + l = len(s) + n += 1 + l + protohelpers.SizeOfVarint(uint64(l)) + } + } n += len(m.unknownFields) return n } @@ -5573,6 +5602,38 @@ func (m *SelectMergeProfileRequest) UnmarshalVT(dAtA []byte) error { } } iNdEx = postIndex + case 7: + if wireType != 2 { + return fmt.Errorf("proto: wrong wireType = %d for field ProfileIdSelector", wireType) + } + var stringLen uint64 + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return protohelpers.ErrIntOverflow + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + stringLen |= uint64(b&0x7F) << shift + if b < 0x80 { + break + } + } + intStringLen := int(stringLen) + if intStringLen < 0 { + return protohelpers.ErrInvalidLength + } + postIndex := iNdEx + intStringLen + if postIndex < 0 { + return protohelpers.ErrInvalidLength + } + if postIndex > l { + return io.ErrUnexpectedEOF + } + m.ProfileIdSelector = append(m.ProfileIdSelector, string(dAtA[iNdEx:postIndex])) + iNdEx = postIndex default: iNdEx = preIndex skippy, err := protohelpers.Skip(dAtA[iNdEx:]) diff --git a/api/gen/proto/go/query/v1/query.pb.go b/api/gen/proto/go/query/v1/query.pb.go index a055e32c50..40771289aa 100644 --- a/api/gen/proto/go/query/v1/query.pb.go +++ b/api/gen/proto/go/query/v1/query.pb.go @@ -1245,6 +1245,7 @@ type TreeQuery struct { MaxNodes int64 `protobuf:"varint,1,opt,name=max_nodes,json=maxNodes,proto3" json:"max_nodes,omitempty"` SpanSelector []string `protobuf:"bytes,2,rep,name=span_selector,json=spanSelector,proto3" json:"span_selector,omitempty"` StackTraceSelector *v11.StackTraceSelector `protobuf:"bytes,3,opt,name=stack_trace_selector,json=stackTraceSelector,proto3,oneof" json:"stack_trace_selector,omitempty"` + ProfileIdSelector []string `protobuf:"bytes,4,rep,name=profile_id_selector,json=profileIdSelector,proto3" json:"profile_id_selector,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -1300,6 +1301,13 @@ func (x *TreeQuery) GetStackTraceSelector() *v11.StackTraceSelector { return nil } +func (x *TreeQuery) GetProfileIdSelector() []string { + if x != nil { + return x.ProfileIdSelector + } + return nil +} + type TreeReport struct { state protoimpl.MessageState `protogen:"open.v1"` Query *TreeQuery `protobuf:"bytes,1,opt,name=query,proto3" json:"query,omitempty"` @@ -1355,7 +1363,8 @@ func (x *TreeReport) GetTree() []byte { type PprofQuery struct { state protoimpl.MessageState `protogen:"open.v1"` MaxNodes int64 `protobuf:"varint,1,opt,name=max_nodes,json=maxNodes,proto3" json:"max_nodes,omitempty"` - StackTraceSelector *v11.StackTraceSelector `protobuf:"bytes,2,opt,name=stack_trace_selector,json=stackTraceSelector,proto3,oneof" json:"stack_trace_selector,omitempty"` // TODO(kolesnikovae): Go PGO options. + StackTraceSelector *v11.StackTraceSelector `protobuf:"bytes,2,opt,name=stack_trace_selector,json=stackTraceSelector,proto3,oneof" json:"stack_trace_selector,omitempty"` + ProfileIdSelector []string `protobuf:"bytes,3,rep,name=profile_id_selector,json=profileIdSelector,proto3" json:"profile_id_selector,omitempty"` // TODO(kolesnikovae): Go PGO options. unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -1404,6 +1413,13 @@ func (x *PprofQuery) GetStackTraceSelector() *v11.StackTraceSelector { return nil } +func (x *PprofQuery) GetProfileIdSelector() []string { + if x != nil { + return x.ProfileIdSelector + } + return nil +} + type PprofReport struct { state protoimpl.MessageState `protogen:"open.v1"` Query *PprofQuery `protobuf:"bytes,1,opt,name=query,proto3" json:"query,omitempty"` @@ -1544,20 +1560,22 @@ const file_query_v1_query_proto_rawDesc = "" + "\x10TimeSeriesReport\x12/\n" + "\x05query\x18\x01 \x01(\v2\x19.query.v1.TimeSeriesQueryR\x05query\x121\n" + "\vtime_series\x18\x02 \x03(\v2\x10.types.v1.SeriesR\n" + - "timeSeries\"\xbb\x01\n" + + "timeSeries\"\xeb\x01\n" + "\tTreeQuery\x12\x1b\n" + "\tmax_nodes\x18\x01 \x01(\x03R\bmaxNodes\x12#\n" + "\rspan_selector\x18\x02 \x03(\tR\fspanSelector\x12S\n" + - "\x14stack_trace_selector\x18\x03 \x01(\v2\x1c.types.v1.StackTraceSelectorH\x00R\x12stackTraceSelector\x88\x01\x01B\x17\n" + + "\x14stack_trace_selector\x18\x03 \x01(\v2\x1c.types.v1.StackTraceSelectorH\x00R\x12stackTraceSelector\x88\x01\x01\x12.\n" + + "\x13profile_id_selector\x18\x04 \x03(\tR\x11profileIdSelectorB\x17\n" + "\x15_stack_trace_selector\"K\n" + "\n" + "TreeReport\x12)\n" + "\x05query\x18\x01 \x01(\v2\x13.query.v1.TreeQueryR\x05query\x12\x12\n" + - "\x04tree\x18\x02 \x01(\fR\x04tree\"\x97\x01\n" + + "\x04tree\x18\x02 \x01(\fR\x04tree\"\xc7\x01\n" + "\n" + "PprofQuery\x12\x1b\n" + "\tmax_nodes\x18\x01 \x01(\x03R\bmaxNodes\x12S\n" + - "\x14stack_trace_selector\x18\x02 \x01(\v2\x1c.types.v1.StackTraceSelectorH\x00R\x12stackTraceSelector\x88\x01\x01B\x17\n" + + "\x14stack_trace_selector\x18\x02 \x01(\v2\x1c.types.v1.StackTraceSelectorH\x00R\x12stackTraceSelector\x88\x01\x01\x12.\n" + + "\x13profile_id_selector\x18\x03 \x03(\tR\x11profileIdSelectorB\x17\n" + "\x15_stack_trace_selector\"O\n" + "\vPprofReport\x12*\n" + "\x05query\x18\x01 \x01(\v2\x14.query.v1.PprofQueryR\x05query\x12\x14\n" + diff --git a/api/gen/proto/go/query/v1/query_vtproto.pb.go b/api/gen/proto/go/query/v1/query_vtproto.pb.go index 2186c9c91b..897fd167c5 100644 --- a/api/gen/proto/go/query/v1/query_vtproto.pb.go +++ b/api/gen/proto/go/query/v1/query_vtproto.pb.go @@ -463,6 +463,11 @@ func (m *TreeQuery) CloneVT() *TreeQuery { r.StackTraceSelector = proto.Clone(rhs).(*v11.StackTraceSelector) } } + if rhs := m.ProfileIdSelector; rhs != nil { + tmpContainer := make([]string, len(rhs)) + copy(tmpContainer, rhs) + r.ProfileIdSelector = tmpContainer + } if len(m.unknownFields) > 0 { r.unknownFields = make([]byte, len(m.unknownFields)) copy(r.unknownFields, m.unknownFields) @@ -511,6 +516,11 @@ func (m *PprofQuery) CloneVT() *PprofQuery { r.StackTraceSelector = proto.Clone(rhs).(*v11.StackTraceSelector) } } + if rhs := m.ProfileIdSelector; rhs != nil { + tmpContainer := make([]string, len(rhs)) + copy(tmpContainer, rhs) + r.ProfileIdSelector = tmpContainer + } if len(m.unknownFields) > 0 { r.unknownFields = make([]byte, len(m.unknownFields)) copy(r.unknownFields, m.unknownFields) @@ -1157,6 +1167,15 @@ func (this *TreeQuery) EqualVT(that *TreeQuery) bool { } else if !proto.Equal(this.StackTraceSelector, that.StackTraceSelector) { return false } + if len(this.ProfileIdSelector) != len(that.ProfileIdSelector) { + return false + } + for i, vx := range this.ProfileIdSelector { + vy := that.ProfileIdSelector[i] + if vx != vy { + return false + } + } return string(this.unknownFields) == string(that.unknownFields) } @@ -1207,6 +1226,15 @@ func (this *PprofQuery) EqualVT(that *PprofQuery) bool { } else if !proto.Equal(this.StackTraceSelector, that.StackTraceSelector) { return false } + if len(this.ProfileIdSelector) != len(that.ProfileIdSelector) { + return false + } + for i, vx := range this.ProfileIdSelector { + vy := that.ProfileIdSelector[i] + if vx != vy { + return false + } + } return string(this.unknownFields) == string(that.unknownFields) } @@ -2510,6 +2538,15 @@ func (m *TreeQuery) MarshalToSizedBufferVT(dAtA []byte) (int, error) { i -= len(m.unknownFields) copy(dAtA[i:], m.unknownFields) } + if len(m.ProfileIdSelector) > 0 { + for iNdEx := len(m.ProfileIdSelector) - 1; iNdEx >= 0; iNdEx-- { + i -= len(m.ProfileIdSelector[iNdEx]) + copy(dAtA[i:], m.ProfileIdSelector[iNdEx]) + i = protohelpers.EncodeVarint(dAtA, i, uint64(len(m.ProfileIdSelector[iNdEx]))) + i-- + dAtA[i] = 0x22 + } + } if m.StackTraceSelector != nil { if vtmsg, ok := interface{}(m.StackTraceSelector).(interface { MarshalToSizedBufferVT([]byte) (int, error) @@ -2629,6 +2666,15 @@ func (m *PprofQuery) MarshalToSizedBufferVT(dAtA []byte) (int, error) { i -= len(m.unknownFields) copy(dAtA[i:], m.unknownFields) } + if len(m.ProfileIdSelector) > 0 { + for iNdEx := len(m.ProfileIdSelector) - 1; iNdEx >= 0; iNdEx-- { + i -= len(m.ProfileIdSelector[iNdEx]) + copy(dAtA[i:], m.ProfileIdSelector[iNdEx]) + i = protohelpers.EncodeVarint(dAtA, i, uint64(len(m.ProfileIdSelector[iNdEx]))) + i-- + dAtA[i] = 0x1a + } + } if m.StackTraceSelector != nil { if vtmsg, ok := interface{}(m.StackTraceSelector).(interface { MarshalToSizedBufferVT([]byte) (int, error) @@ -3139,6 +3185,12 @@ func (m *TreeQuery) SizeVT() (n int) { } n += 1 + l + protohelpers.SizeOfVarint(uint64(l)) } + if len(m.ProfileIdSelector) > 0 { + for _, s := range m.ProfileIdSelector { + l = len(s) + n += 1 + l + protohelpers.SizeOfVarint(uint64(l)) + } + } n += len(m.unknownFields) return n } @@ -3180,6 +3232,12 @@ func (m *PprofQuery) SizeVT() (n int) { } n += 1 + l + protohelpers.SizeOfVarint(uint64(l)) } + if len(m.ProfileIdSelector) > 0 { + for _, s := range m.ProfileIdSelector { + l = len(s) + n += 1 + l + protohelpers.SizeOfVarint(uint64(l)) + } + } n += len(m.unknownFields) return n } @@ -5754,6 +5812,38 @@ func (m *TreeQuery) UnmarshalVT(dAtA []byte) error { } } iNdEx = postIndex + case 4: + if wireType != 2 { + return fmt.Errorf("proto: wrong wireType = %d for field ProfileIdSelector", wireType) + } + var stringLen uint64 + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return protohelpers.ErrIntOverflow + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + stringLen |= uint64(b&0x7F) << shift + if b < 0x80 { + break + } + } + intStringLen := int(stringLen) + if intStringLen < 0 { + return protohelpers.ErrInvalidLength + } + postIndex := iNdEx + intStringLen + if postIndex < 0 { + return protohelpers.ErrInvalidLength + } + if postIndex > l { + return io.ErrUnexpectedEOF + } + m.ProfileIdSelector = append(m.ProfileIdSelector, string(dAtA[iNdEx:postIndex])) + iNdEx = postIndex default: iNdEx = preIndex skippy, err := protohelpers.Skip(dAtA[iNdEx:]) @@ -5989,6 +6079,38 @@ func (m *PprofQuery) UnmarshalVT(dAtA []byte) error { } } iNdEx = postIndex + case 3: + if wireType != 2 { + return fmt.Errorf("proto: wrong wireType = %d for field ProfileIdSelector", wireType) + } + var stringLen uint64 + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return protohelpers.ErrIntOverflow + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + stringLen |= uint64(b&0x7F) << shift + if b < 0x80 { + break + } + } + intStringLen := int(stringLen) + if intStringLen < 0 { + return protohelpers.ErrInvalidLength + } + postIndex := iNdEx + intStringLen + if postIndex < 0 { + return protohelpers.ErrInvalidLength + } + if postIndex > l { + return io.ErrUnexpectedEOF + } + m.ProfileIdSelector = append(m.ProfileIdSelector, string(dAtA[iNdEx:postIndex])) + iNdEx = postIndex default: iNdEx = preIndex skippy, err := protohelpers.Skip(dAtA[iNdEx:]) diff --git a/api/openapiv2/gen/phlare.swagger.json b/api/openapiv2/gen/phlare.swagger.json index bae1ad7d9d..4f8e36fbe9 100644 --- a/api/openapiv2/gen/phlare.swagger.json +++ b/api/openapiv2/gen/phlare.swagger.json @@ -1733,7 +1733,13 @@ "format": "int64" }, "stackTraceSelector": { - "$ref": "#/definitions/v1StackTraceSelector", + "$ref": "#/definitions/v1StackTraceSelector" + }, + "profileIdSelector": { + "type": "array", + "items": { + "type": "string" + }, "description": "TODO(kolesnikovae): Go PGO options." } } @@ -2518,6 +2524,12 @@ }, "stackTraceSelector": { "$ref": "#/definitions/v1StackTraceSelector" + }, + "profileIdSelector": { + "type": "array", + "items": { + "type": "string" + } } } }, diff --git a/api/querier/v1/querier.proto b/api/querier/v1/querier.proto index a27d01248e..a8126c5ae9 100644 --- a/api/querier/v1/querier.proto +++ b/api/querier/v1/querier.proto @@ -199,6 +199,8 @@ message SelectMergeProfileRequest { optional int64 max_nodes = 5; // Select stack traces that match the provided selector. optional types.v1.StackTraceSelector stack_trace_selector = 6; + // List of Profile UUIDs to query + repeated string profile_id_selector = 7 [(gnostic.openapi.v3.property).example = {yaml: "['7c9e6679-7425-40de-944b-e07fc1f90ae7']"}]; } message SelectSeriesRequest { diff --git a/api/query/v1/query.proto b/api/query/v1/query.proto index ce30905b41..0e867ac5af 100644 --- a/api/query/v1/query.proto +++ b/api/query/v1/query.proto @@ -173,6 +173,7 @@ message TreeQuery { int64 max_nodes = 1; repeated string span_selector = 2; optional types.v1.StackTraceSelector stack_trace_selector = 3; + repeated string profile_id_selector = 4; } message TreeReport { @@ -183,6 +184,7 @@ message TreeReport { message PprofQuery { int64 max_nodes = 1; optional types.v1.StackTraceSelector stack_trace_selector = 2; + repeated string profile_id_selector = 3; // TODO(kolesnikovae): Go PGO options. } diff --git a/docs/sources/reference-server-api/index.md b/docs/sources/reference-server-api/index.md index 23bfa03ae8..383e574000 100644 --- a/docs/sources/reference-server-api/index.md +++ b/docs/sources/reference-server-api/index.md @@ -326,6 +326,7 @@ A request body with the following fields is required: |`end` | Milliseconds since epoch. | `1676289600000` | |`labelSelector` | Label selector string | `{namespace="my-namespace"}` | |`maxNodes` | Limit the nodes returned to only show the node with the max_node's biggest total | | +|`profileIdSelector` | List of Profile UUIDs to query | `["7c9e6679-7425-40de-944b-e07fc1f90ae7"]` | |`profileTypeID` | Profile Type ID string in the form ::::. | `process_cpu:cpu:nanoseconds:cpu:nanoseconds` | |`stackTraceSelector.callSite[].name` | | | |`stackTraceSelector.goPgo.aggregateCallees` | Aggregate callees causes the leaf location line number to be ignored, thus aggregating all callee samples (but not callers). | | @@ -338,6 +339,9 @@ curl \ -d '{ "end": '$(date +%s)000', "labelSelector": "{namespace=\"my-namespace\"}", + "profileIdSelector": [ + "7c9e6679-7425-40de-944b-e07fc1f90ae7" + ], "profileTypeID": "process_cpu:cpu:nanoseconds:cpu:nanoseconds", "start": '$(expr $(date +%s) - 3600 )000' }' \ @@ -350,6 +354,9 @@ import datetime body = { "end": int(datetime.datetime.now().timestamp() * 1000), "labelSelector": "{namespace=\"my-namespace\"}", + "profileIdSelector": [ + "7c9e6679-7425-40de-944b-e07fc1f90ae7" + ], "profileTypeID": "process_cpu:cpu:nanoseconds:cpu:nanoseconds", "start": int((datetime.datetime.now()- datetime.timedelta(hours = 1)).timestamp() * 1000) } diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_profile.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_profile.go index 552b5068f9..aa6c626c04 100644 --- a/pkg/frontend/readpath/queryfrontend/query_select_merge_profile.go +++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_profile.go @@ -52,6 +52,7 @@ func (q *QueryFrontend) SelectMergeProfile( Pprof: &queryv1.PprofQuery{ MaxNodes: c.Msg.GetMaxNodes(), StackTraceSelector: c.Msg.StackTraceSelector, + ProfileIdSelector: c.Msg.ProfileIdSelector, }, }}, }) diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces.go index 4f4ed822b0..e950704452 100644 --- a/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces.go +++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces.go @@ -73,6 +73,7 @@ func (q *QueryFrontend) selectMergeStacktracesTree( Tree: &queryv1.TreeQuery{ MaxNodes: maxNodes, StackTraceSelector: c.Msg.StackTraceSelector, + ProfileIdSelector: c.Msg.ProfileIdSelector, }, }}, }) diff --git a/pkg/querier/querier.go b/pkg/querier/querier.go index 627547454a..373ad6ad01 100644 --- a/pkg/querier/querier.go +++ b/pkg/querier/querier.go @@ -664,9 +664,8 @@ func (q *Querier) SelectMergeStacktraces(ctx context.Context, req *connect.Reque req.Msg.MaxNodes = &mn } - // TODO: remove when profile_id_selector is implemented if len(req.Msg.ProfileIdSelector) > 0 { - return nil, connect.NewError(connect.CodeUnimplemented, errors.New("profile_id_selector is not yet implemented")) + return nil, connect.NewError(connect.CodeUnimplemented, errors.New("profile_id_selector is only supported with the v2 query backend")) } t, err := q.selectTree(ctx, req.Msg) diff --git a/pkg/querybackend/block_reader_test.go b/pkg/querybackend/block_reader_test.go index 1f76614ca6..260aa8c8f5 100644 --- a/pkg/querybackend/block_reader_test.go +++ b/pkg/querybackend/block_reader_test.go @@ -15,6 +15,7 @@ import ( profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1" metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1" queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1" + typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1" "github.com/grafana/pyroscope/pkg/block" "github.com/grafana/pyroscope/pkg/block/metadata" phlaremodel "github.com/grafana/pyroscope/pkg/model" @@ -327,3 +328,177 @@ func (s *testSuite) Test_QueryTree_All_Tenant_Isolation() { s.Require().NotNil(resp) s.Require().Len(resp.Reports, 0) } + +func (s *testSuite) Test_ProfileIDSelector() { + // Get a real profile ID for valid test case + validProfileID := s.getProfileIDFromExemplars(s.T()) + + // Load baseline fixture for tree comparison + baselineTree, err := os.ReadFile("testdata/fixtures/tree_16.txt") + s.Require().NoError(err) + + // Get baseline tree for comparison + allTreeResp, err := s.reader.Invoke(s.ctx, &queryv1.InvokeRequest{ + EndTime: time.Now().UnixMilli(), + LabelSelector: "{}", + QueryPlan: s.plan, + Query: []*queryv1.Query{{ + QueryType: queryv1.QueryType_QUERY_TREE, + Tree: &queryv1.TreeQuery{MaxNodes: 16}, + }}, + Tenant: s.tenant, + }) + s.Require().NoError(err) + allTree, err := phlaremodel.UnmarshalTree(allTreeResp.Reports[0].Tree.Tree) + s.Require().NoError(err) + + // Get baseline pprof for comparison + allPprofResp, err := s.reader.Invoke(s.ctx, &queryv1.InvokeRequest{ + StartTime: startTime.UnixMilli(), + EndTime: startTime.Add(5 * time.Minute).UnixMilli(), + LabelSelector: "{}", + QueryPlan: s.plan, + Query: []*queryv1.Query{{ + QueryType: queryv1.QueryType_QUERY_PPROF, + Pprof: &queryv1.PprofQuery{}, + }}, + Tenant: s.tenant, + }) + s.Require().NoError(err) + var allProfile profilev1.Profile + err = pprof.Unmarshal(allPprofResp.Reports[0].Pprof.Pprof, &allProfile) + s.Require().NoError(err) + + tests := []struct { + queryType queryv1.QueryType + name string + profileIDSelector []string + wantErr bool + expectBaseline bool + expectFiltered bool + expectEmpty bool + }{ + // Tree query tests + {queryv1.QueryType_QUERY_TREE, "tree/invalid UUID returns error", []string{"invalid-uuid"}, true, false, false, false}, + {queryv1.QueryType_QUERY_TREE, "tree/empty selector returns baseline", []string{}, false, true, false, false}, + {queryv1.QueryType_QUERY_TREE, "tree/nil selector returns baseline", nil, false, true, false, false}, + {queryv1.QueryType_QUERY_TREE, "tree/non-existent UUID returns empty result", []string{"00000000-0000-0000-0000-000000000000"}, false, false, false, true}, + {queryv1.QueryType_QUERY_TREE, "tree/valid UUID filters to single profile", []string{validProfileID}, false, false, true, false}, + + // Pprof query tests + {queryv1.QueryType_QUERY_PPROF, "pprof/invalid UUID returns error", []string{"not-a-uuid"}, true, false, false, false}, + {queryv1.QueryType_QUERY_PPROF, "pprof/empty selector returns baseline", []string{}, false, true, false, false}, + {queryv1.QueryType_QUERY_PPROF, "pprof/nil selector returns baseline", nil, false, true, false, false}, + {queryv1.QueryType_QUERY_PPROF, "pprof/non-existent UUID returns empty result", []string{"00000000-0000-0000-0000-000000000000"}, false, false, false, true}, + {queryv1.QueryType_QUERY_PPROF, "pprof/valid UUID filters to single profile", []string{validProfileID}, false, false, true, false}, + } + + for _, tt := range tests { + s.Run(tt.name, func() { + var query *queryv1.Query + var reqStartTime, reqEndTime int64 + + if tt.queryType == queryv1.QueryType_QUERY_TREE { + reqEndTime = time.Now().UnixMilli() + query = &queryv1.Query{ + QueryType: queryv1.QueryType_QUERY_TREE, + Tree: &queryv1.TreeQuery{ + MaxNodes: 16, + ProfileIdSelector: tt.profileIDSelector, + }, + } + } else { + reqStartTime = startTime.UnixMilli() + reqEndTime = startTime.Add(5 * time.Minute).UnixMilli() + query = &queryv1.Query{ + QueryType: queryv1.QueryType_QUERY_PPROF, + Pprof: &queryv1.PprofQuery{ + ProfileIdSelector: tt.profileIDSelector, + }, + } + } + + resp, err := s.reader.Invoke(s.ctx, &queryv1.InvokeRequest{ + StartTime: reqStartTime, + EndTime: reqEndTime, + LabelSelector: "{}", + QueryPlan: s.plan, + Query: []*queryv1.Query{query}, + Tenant: s.tenant, + }) + + if tt.wantErr { + s.Require().Error(err) + s.Require().Nil(resp) + return + } + + s.Require().NoError(err) + s.Require().NotNil(resp) + s.Require().Len(resp.Reports, 1) + + if tt.queryType == queryv1.QueryType_QUERY_TREE { + tree, err := phlaremodel.UnmarshalTree(resp.Reports[0].Tree.Tree) + s.Require().NoError(err) + + if tt.expectBaseline { + s.Assert().Equal(string(baselineTree), tree.String()) + } + if tt.expectEmpty { + s.Assert().Zero(tree.Total()) + } + if tt.expectFiltered { + s.Assert().Less(tree.Total(), allTree.Total()) + s.Assert().NotZero(tree.Total()) + } + } else { + var profile profilev1.Profile + err = pprof.Unmarshal(resp.Reports[0].Pprof.Pprof, &profile) + s.Require().NoError(err) + + if tt.expectBaseline { + s.Assert().Equal(len(allProfile.Sample), len(profile.Sample)) + } + if tt.expectEmpty { + s.Assert().Zero(len(profile.Sample)) + } + if tt.expectFiltered { + s.Assert().Less(len(profile.Sample), len(allProfile.Sample)) + s.Assert().NotZero(len(profile.Sample)) + } + } + }) + } +} + +func (s *testSuite) getProfileIDFromExemplars(t *testing.T) string { + t.Helper() + + resp, err := s.reader.Invoke(s.ctx, &queryv1.InvokeRequest{ + StartTime: startTime.UnixMilli(), + EndTime: startTime.Add(5 * time.Minute).UnixMilli(), + LabelSelector: "{}", + QueryPlan: s.plan, + Query: []*queryv1.Query{{ + QueryType: queryv1.QueryType_QUERY_TIME_SERIES, + TimeSeries: &queryv1.TimeSeriesQuery{ + Step: 30.0, + ExemplarType: typesv1.ExemplarType_EXEMPLAR_TYPE_INDIVIDUAL, + }, + }}, + Tenant: s.tenant, + }) + s.Require().NoError(err) + s.Require().NotNil(resp) + + // Find first exemplar with a profile ID + for _, serie := range resp.Reports[0].TimeSeries.TimeSeries { + for _, point := range serie.Points { + if len(point.Exemplars) > 0 && point.Exemplars[0].ProfileId != "" { + return point.Exemplars[0].ProfileId + } + } + } + s.Require().FailNow("no profile ID found in exemplars") + return "" +} diff --git a/pkg/querybackend/query_pprof.go b/pkg/querybackend/query_pprof.go index e1526f3b96..5b2167e353 100644 --- a/pkg/querybackend/query_pprof.go +++ b/pkg/querybackend/query_pprof.go @@ -31,7 +31,16 @@ func init() { } func queryPprof(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) { - entries, err := profileEntryIterator(q) + var profileOpts []profileIteratorOption + if len(query.Pprof.ProfileIdSelector) > 0 { + opt, err := withProfileIDSelector(query.Pprof.ProfileIdSelector...) + if err != nil { + return nil, err + } + profileOpts = append(profileOpts, opt) + } + + entries, err := profileEntryIterator(q, profileOpts...) if err != nil { return nil, err } diff --git a/pkg/querybackend/query_tree.go b/pkg/querybackend/query_tree.go index d311f32941..50b44b1232 100644 --- a/pkg/querybackend/query_tree.go +++ b/pkg/querybackend/query_tree.go @@ -29,7 +29,16 @@ func init() { } func queryTree(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) { - entries, err := profileEntryIterator(q) + var profileOpts []profileIteratorOption + if len(query.Tree.ProfileIdSelector) > 0 { + opt, err := withProfileIDSelector(query.Tree.ProfileIdSelector...) + if err != nil { + return nil, err + } + profileOpts = append(profileOpts, opt) + } + + entries, err := profileEntryIterator(q, profileOpts...) if err != nil { return nil, err }