Skip to content

Commit 5b32919

Browse files
simonswinesimonswine
authored
Add subcommand to upload profiles to Phlare using PusherService.Push (grafana/phlare#607)
* Refactor client * Code is kept in a separate file * Support for basic auth via flags * Uses environment variables to read some parameters * Add uploads * Add profile uploads to profilecli
1 parent 0d2d248 commit 5b32919

File tree

6 files changed

+190
-48
lines changed

6 files changed

+190
-48
lines changed

‎cmd/profilecli/client.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
"net/http"
6+
7+
"github.com/prometheus/common/version"
8+
"gopkg.in/alecthomas/kingpin.v2"
9+
)
10+
11+
const (
12+
envPrefix = "PROFILECLI_"
13+
)
14+
15+
var userAgentHeader = fmt.Sprintf("phlare/%s", version.Version)
16+
17+
type phlareClient struct {
18+
TenantID string
19+
URL string
20+
BasicAuth struct {
21+
Username string
22+
Password string
23+
}
24+
client *http.Client
25+
}
26+
27+
type authRoundTripper struct {
28+
client *phlareClient
29+
next http.RoundTripper
30+
}
31+
32+
func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
33+
if c := a.client; c != nil {
34+
if c.TenantID != "" {
35+
req.Header.Set("X-Scope-OrgID", c.TenantID)
36+
}
37+
if c.BasicAuth.Username != "" || c.BasicAuth.Password != "" {
38+
req.SetBasicAuth(c.BasicAuth.Username, c.BasicAuth.Password)
39+
}
40+
}
41+
42+
req.Header.Set("User-Agent", userAgentHeader)
43+
return a.next.RoundTrip(req)
44+
}
45+
46+
func (c *phlareClient) httpClient() *http.Client {
47+
if c.client == nil {
48+
c.client = &http.Client{Transport: &authRoundTripper{
49+
client: c,
50+
next: http.DefaultTransport,
51+
}}
52+
}
53+
return c.client
54+
}
55+
56+
type commander interface {
57+
Flag(name, help string) *kingpin.FlagClause
58+
Arg(name, help string) *kingpin.ArgClause
59+
}
60+
61+
func addPhlareClient(cmd commander) *phlareClient {
62+
client := &phlareClient{}
63+
64+
cmd.Flag("url", "URL of the profile store.").Default("http://localhost:4100").Envar(envPrefix + "URL").StringVar(&client.URL)
65+
cmd.Flag("tenant-id", "The tenant ID to be used for the X-Scope-OrgID header.").Default("").Envar(envPrefix + "TENANT_ID").StringVar(&client.TenantID)
66+
cmd.Flag("username", "The username to be used for basic auth.").Default("").Envar(envPrefix + "USERNAME").StringVar(&client.BasicAuth.Username)
67+
cmd.Flag("password", "The password to be used for basic auth.").Default("").Envar(envPrefix + "PASSWORD").StringVar(&client.BasicAuth.Password)
68+
return client
69+
}

‎cmd/profilecli/main.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,15 @@ func main() {
5353
queryOutput := queryCmd.Flag("output", "How to output the result, examples: console, raw, pprof=./my.pprof").Default("console").String()
5454
queryMergeCmd := queryCmd.Command("merge", "Request merged profile.")
5555

56+
uploadCmd := app.Command("upload", "Upload profile(s).")
57+
uploadParams := addUploadParams(uploadCmd)
58+
5659
// parse command line arguments
5760
parsedCmd := kingpin.MustParse(app.Parse(os.Args[1:]))
5861

5962
// enable verbose logging if requested
6063
if !cfg.verbose {
61-
logger = level.NewFilter(logger, level.AllowWarn())
64+
logger = level.NewFilter(logger, level.AllowInfo())
6265
}
6366

6467
switch parsedCmd {
@@ -74,6 +77,10 @@ func main() {
7477
if err := queryMerge(ctx, queryParams, *queryOutput); err != nil {
7578
os.Exit(checkError(err))
7679
}
80+
case uploadCmd.FullCommand():
81+
if err := upload(ctx, uploadParams); err != nil {
82+
os.Exit(checkError(err))
83+
}
7784
default:
7885
level.Error(logger).Log("msg", "unknown command", "cmd", parsedCmd)
7986
}

‎cmd/profilecli/my.pprof

38.2 KB
Binary file not shown.

‎cmd/profilecli/query.go

Lines changed: 1 addition & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"context"
66
"fmt"
77
"io"
8-
"net/http"
98
"os"
109
"strings"
1110
"time"
@@ -19,8 +18,6 @@ import (
1918
"github.com/mattn/go-isatty"
2019
"github.com/pkg/errors"
2120
"github.com/prometheus/common/model"
22-
"github.com/prometheus/common/version"
23-
"gopkg.in/alecthomas/kingpin.v2"
2421

2522
querierv1 "github.com/grafana/phlare/api/gen/proto/go/querier/v1"
2623
"github.com/grafana/phlare/api/gen/proto/go/querier/v1/querierv1connect"
@@ -32,8 +29,6 @@ const (
3229
outputPprof = "pprof="
3330
)
3431

35-
var userAgentHeader = fmt.Sprintf("phlare/%s", version.Version)
36-
3732
func parseTime(s string) (time.Time, error) {
3833
if s == "" {
3934
return time.Time{}, fmt.Errorf("empty time")
@@ -68,35 +63,6 @@ func parseRelativeTime(s string) (time.Duration, error) {
6863
return time.Duration(d), nil
6964
}
7065

71-
type phlareClient struct {
72-
TenantID string
73-
URL string
74-
client *http.Client
75-
}
76-
77-
type authRoundTripper struct {
78-
tenantID string
79-
next http.RoundTripper
80-
}
81-
82-
func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
83-
if a.tenantID != "" {
84-
req.Header.Set("X-Scope-OrgID", a.tenantID)
85-
}
86-
req.Header.Set("User-Agent", userAgentHeader)
87-
return a.next.RoundTrip(req)
88-
}
89-
90-
func (c *phlareClient) httpClient() *http.Client {
91-
if c.client == nil {
92-
c.client = &http.Client{Transport: &authRoundTripper{
93-
tenantID: c.TenantID,
94-
next: http.DefaultTransport,
95-
}}
96-
}
97-
return c.client
98-
}
99-
10066
func (c *phlareClient) queryClient() querierv1connect.QuerierServiceClient {
10167
return querierv1connect.NewQuerierServiceClient(
10268
c.httpClient(),
@@ -129,19 +95,7 @@ func (p *queryParams) parseFromTo() (from time.Time, to time.Time, err error) {
12995
return from, to, nil
13096
}
13197

132-
type flagger interface {
133-
Flag(name, help string) *kingpin.FlagClause
134-
}
135-
136-
func addPhlareClient(queryCmd flagger) *phlareClient {
137-
client := &phlareClient{}
138-
139-
queryCmd.Flag("url", "URL of the profile store.").Default("http://localhost:4100").StringVar(&client.URL)
140-
queryCmd.Flag("tenant-id", "The tenant ID to be used for the X-Scope-OrgID header.").Default("").StringVar(&client.TenantID)
141-
return client
142-
}
143-
144-
func addQueryParams(queryCmd flagger) *queryParams {
98+
func addQueryParams(queryCmd commander) *queryParams {
14599
var (
146100
params = &queryParams{}
147101
)

‎cmd/profilecli/upload.go

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package main
2+
3+
import (
4+
"context"
5+
"os"
6+
7+
"github.com/bufbuild/connect-go"
8+
"github.com/go-kit/log/level"
9+
"github.com/google/uuid"
10+
11+
pushv1 "github.com/grafana/phlare/api/gen/proto/go/push/v1"
12+
"github.com/grafana/phlare/api/gen/proto/go/push/v1/pushv1connect"
13+
"github.com/grafana/phlare/pkg/model"
14+
"github.com/grafana/phlare/pkg/pprof"
15+
)
16+
17+
func (c *phlareClient) pusherClient() pushv1connect.PusherServiceClient {
18+
return pushv1connect.NewPusherServiceClient(
19+
c.httpClient(),
20+
c.URL,
21+
)
22+
}
23+
24+
type uploadParams struct {
25+
*phlareClient
26+
paths []string
27+
extraLabels map[string]string
28+
}
29+
30+
func addUploadParams(cmd commander) *uploadParams {
31+
var (
32+
params = &uploadParams{
33+
extraLabels: map[string]string{},
34+
}
35+
)
36+
params.phlareClient = addPhlareClient(cmd)
37+
38+
cmd.Arg("path", "Path(s) to profile(s) to upload").Required().ExistingFilesVar(&params.paths)
39+
cmd.Flag("extra-labels", "Add additional labels to the profile(s)").Default("job=profilecli-upload").StringMapVar(&params.extraLabels)
40+
return params
41+
}
42+
43+
func upload(ctx context.Context, params *uploadParams) (err error) {
44+
pc := params.phlareClient.pusherClient()
45+
46+
lblStrings := make([]string, 0, len(params.extraLabels)*2)
47+
for key, value := range params.extraLabels {
48+
lblStrings = append(lblStrings, key, value)
49+
}
50+
51+
var (
52+
lbl = model.LabelsFromStrings(lblStrings...)
53+
series = make([]*pushv1.RawProfileSeries, len(params.paths))
54+
lblBuilder = model.NewLabelsBuilder(lbl)
55+
)
56+
for idx, path := range params.paths {
57+
lblBuilder.Reset(lbl)
58+
59+
data, err := os.ReadFile(path)
60+
if err != nil {
61+
return nil
62+
}
63+
64+
profile, err := pprof.RawFromBytes(data)
65+
if err != nil {
66+
return err
67+
}
68+
69+
// detect name if no name has been set
70+
if lbl.Get(model.LabelNameProfileName) == "" {
71+
name := "unknown"
72+
for _, t := range profile.Profile.SampleType {
73+
if sid := int(t.Type); sid < len(profile.StringTable) {
74+
if s := profile.StringTable[sid]; s == "cpu" {
75+
name = "process_cpu"
76+
break
77+
} else if s == "alloc_space" || s == "inuse_space" {
78+
name = "memory"
79+
break
80+
} else {
81+
level.Debug(logger).Log("msg", "unspecific/unknown profile sample type", "profile", s)
82+
}
83+
}
84+
}
85+
lblBuilder.Set(model.LabelNameProfileName, name)
86+
}
87+
88+
series[idx] = &pushv1.RawProfileSeries{
89+
Labels: lblBuilder.Labels(),
90+
Samples: []*pushv1.RawSample{{
91+
ID: uuid.New().String(),
92+
RawProfile: data,
93+
}},
94+
}
95+
}
96+
97+
_, err = pc.Push(ctx, connect.NewRequest(&pushv1.PushRequest{
98+
Series: series,
99+
}))
100+
101+
if err != nil {
102+
return err
103+
}
104+
105+
for idx := range series {
106+
level.Info(logger).Log("msg", "successfully uploaded profile", "id", series[idx].Samples[0].ID, "labels", model.Labels(series[idx].Labels).ToPrometheusLabels().String(), "path", params.paths[idx])
107+
}
108+
109+
return nil
110+
}

‎pkg/model/labels.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"strings"
88

99
"github.com/cespare/xxhash/v2"
10+
pmodel "github.com/prometheus/common/model"
1011
"github.com/prometheus/prometheus/model/labels"
1112
"github.com/prometheus/prometheus/promql/parser"
1213

@@ -22,6 +23,7 @@ const (
2223
LabelNamePeriodType = "__period_type__"
2324
LabelNamePeriodUnit = "__period_unit__"
2425
LabelNameDelta = "__delta__"
26+
LabelNameProfileName = pmodel.MetricNameLabel
2527

2628
labelSep = '\xfe'
2729
)

0 commit comments

Comments
 (0)