-
Notifications
You must be signed in to change notification settings - Fork 581
/
filter_profiles_bidi.go
183 lines (166 loc) · 5.8 KB
/
filter_profiles_bidi.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
package phlaredb
import (
"context"
"io"
"github.com/bufbuild/connect-go"
"github.com/opentracing/opentracing-go"
otlog "github.com/opentracing/opentracing-go/log"
"github.com/pkg/errors"
"github.com/prometheus/common/model"
ingestv1 "github.com/grafana/pyroscope/api/gen/proto/go/ingester/v1"
typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
"github.com/grafana/pyroscope/pkg/iter"
phlaremodel "github.com/grafana/pyroscope/pkg/model"
)
type BidiServerMerge[Res any, Req any] interface {
Send(Res) error
Receive() (Req, error)
}
type labelWithIndex struct {
phlaremodel.Labels
index int
}
type ProfileWithIndex struct {
Profile
Index int
}
type indexedProfileIterator struct {
iter.Iterator[Profile]
querierIndex int
}
func (pqi *indexedProfileIterator) At() ProfileWithIndex {
return ProfileWithIndex{
Profile: pqi.Iterator.At(),
Index: pqi.querierIndex,
}
}
type filterRequest interface {
*ingestv1.MergeProfilesStacktracesRequest |
*ingestv1.MergeProfilesLabelsRequest |
*ingestv1.MergeProfilesPprofRequest |
*ingestv1.MergeSpanProfileRequest
}
type filterResponse interface {
*ingestv1.MergeProfilesStacktracesResponse |
*ingestv1.MergeProfilesLabelsResponse |
*ingestv1.MergeProfilesPprofResponse |
*ingestv1.MergeSpanProfileResponse
}
// filterProfiles merges and dedupe profiles from different iterators and allow filtering via a bidi stream.
func filterProfiles[B BidiServerMerge[Res, Req], Res filterResponse, Req filterRequest](
ctx context.Context, profiles []iter.Iterator[Profile], batchProfileSize int, stream B,
) ([][]Profile, error) {
sp, _ := opentracing.StartSpanFromContext(ctx, "filterProfiles")
defer sp.Finish()
selection := make([][]Profile, len(profiles))
selectProfileResult := &ingestv1.ProfileSets{
Profiles: make([]*ingestv1.SeriesProfile, 0, batchProfileSize),
LabelsSets: make([]*typesv1.Labels, 0, batchProfileSize),
}
its := make([]iter.Iterator[ProfileWithIndex], len(profiles))
for i, iter := range profiles {
iter := iter
its[i] = &indexedProfileIterator{
Iterator: iter,
querierIndex: i,
}
}
if err := iter.ReadBatch(ctx, iter.NewMergeIterator(ProfileWithIndex{
Profile: maxBlockProfile,
Index: 0,
}, true, its...), batchProfileSize, func(ctx context.Context, batch []ProfileWithIndex) error {
sp.LogFields(
otlog.Int("batch_len", len(batch)),
otlog.Int("batch_requested_size", batchProfileSize),
)
seriesByFP := map[model.Fingerprint]labelWithIndex{}
selectProfileResult.LabelsSets = selectProfileResult.LabelsSets[:0]
selectProfileResult.Profiles = selectProfileResult.Profiles[:0]
for _, profile := range batch {
var ok bool
var lblsIdx labelWithIndex
lblsIdx, ok = seriesByFP[profile.Fingerprint()]
if !ok {
lblsIdx = labelWithIndex{
Labels: profile.Labels(),
index: len(selectProfileResult.LabelsSets),
}
seriesByFP[profile.Fingerprint()] = lblsIdx
selectProfileResult.LabelsSets = append(selectProfileResult.LabelsSets, &typesv1.Labels{Labels: profile.Labels()})
}
selectProfileResult.Profiles = append(selectProfileResult.Profiles, &ingestv1.SeriesProfile{
LabelIndex: int32(lblsIdx.index),
Timestamp: int64(profile.Timestamp()),
})
}
sp.LogFields(otlog.String("msg", "sending batch to client"))
var err error
switch s := BidiServerMerge[Res, Req](stream).(type) {
case BidiServerMerge[*ingestv1.MergeProfilesStacktracesResponse, *ingestv1.MergeProfilesStacktracesRequest]:
err = s.Send(&ingestv1.MergeProfilesStacktracesResponse{
SelectedProfiles: selectProfileResult,
})
case BidiServerMerge[*ingestv1.MergeProfilesLabelsResponse, *ingestv1.MergeProfilesLabelsRequest]:
err = s.Send(&ingestv1.MergeProfilesLabelsResponse{
SelectedProfiles: selectProfileResult,
})
case BidiServerMerge[*ingestv1.MergeProfilesPprofResponse, *ingestv1.MergeProfilesPprofRequest]:
err = s.Send(&ingestv1.MergeProfilesPprofResponse{
SelectedProfiles: selectProfileResult,
})
case BidiServerMerge[*ingestv1.MergeSpanProfileResponse, *ingestv1.MergeSpanProfileRequest]:
err = s.Send(&ingestv1.MergeSpanProfileResponse{
SelectedProfiles: selectProfileResult,
})
}
// read a batch of profiles and sends it.
if err != nil {
if errors.Is(err, io.EOF) {
return connect.NewError(connect.CodeCanceled, errors.New("client closed stream"))
}
return err
}
sp.LogFields(otlog.String("msg", "batch sent to client"))
sp.LogFields(otlog.String("msg", "reading selection from client"))
// handle response for the batch.
var selected []bool
switch s := BidiServerMerge[Res, Req](stream).(type) {
case BidiServerMerge[*ingestv1.MergeProfilesStacktracesResponse, *ingestv1.MergeProfilesStacktracesRequest]:
selectionResponse, err := s.Receive()
if err == nil {
selected = selectionResponse.Profiles
}
case BidiServerMerge[*ingestv1.MergeProfilesLabelsResponse, *ingestv1.MergeProfilesLabelsRequest]:
selectionResponse, err := s.Receive()
if err == nil {
selected = selectionResponse.Profiles
}
case BidiServerMerge[*ingestv1.MergeProfilesPprofResponse, *ingestv1.MergeProfilesPprofRequest]:
selectionResponse, err := s.Receive()
if err == nil {
selected = selectionResponse.Profiles
}
case BidiServerMerge[*ingestv1.MergeSpanProfileResponse, *ingestv1.MergeSpanProfileRequest]:
selectionResponse, err := s.Receive()
if err == nil {
selected = selectionResponse.Profiles
}
}
if err != nil {
if errors.Is(err, io.EOF) {
return connect.NewError(connect.CodeCanceled, errors.New("client closed stream"))
}
return err
}
sp.LogFields(otlog.String("msg", "selection received"))
for i, k := range selected {
if k {
selection[batch[i].Index] = append(selection[batch[i].Index], batch[i].Profile)
}
}
return nil
}); err != nil {
return nil, err
}
return selection, nil
}