Add option to filter prometheus labels

pkg/cli/cli.go

@@ -44,7 +44,7 @@ func generateRootCmd(cfg *config.Config) *ffcli.Command {
 		rootFlagSet      = flag.NewFlagSet("pyroscope", flag.ExitOnError)
 	)
 
-	serverSortedFlags := PopulateFlagSet(&cfg.Server, serverFlagSet)
+	serverSortedFlags := PopulateFlagSet(&cfg.Server, serverFlagSet, WithSkip("metric-export-rules"))
 	agentSortedFlags := PopulateFlagSet(&cfg.Agent, agentFlagSet, WithSkip("targets"))
 	convertSortedFlags := PopulateFlagSet(&cfg.Convert, convertFlagSet)
 	execSortedFlags := PopulateFlagSet(&cfg.Exec, execFlagSet, WithSkip("pid"))
@@ -61,7 +61,7 @@ func generateRootCmd(cfg *config.Config) *ffcli.Command {
 
 	serverCmd := &ffcli.Command{
 		UsageFunc:  serverSortedFlags.printUsage,
-		Options:    options,
+		Options:    append(options, ff.WithIgnoreUndefined(true)),
 		Name:       "server",
 		ShortUsage: "pyroscope server [flags]",
 		ShortHelp:  "starts pyroscope server. This is the database + web-based user interface",

pkg/cli/flags.go

@@ -287,7 +287,7 @@ func PopulateFlagSet(obj interface{}, flagSet *flag.FlagSet, opts ...FlagOption)
 			}
 			flagSet.UintVar(val, nameVal, defaultVal, descVal)
 		default:
-			logrus.Fatalf("type %s is not supported", field.Type)
+			continue
 		}
 	}
 	return NewSortedFlags(obj, flagSet, deprecatedFields)

pkg/cli/flags_test.go

@@ -118,6 +118,83 @@ var _ = Describe("flags", func() {
 				Expect(cfg.Foo).To(Equal("test-val-4"))
 			})
 
+			It("server configuration", func() {
+				exampleFlagSet := flag.NewFlagSet("example flag set", flag.ExitOnError)
+				var cfg config.Server
+				PopulateFlagSet(&cfg, exampleFlagSet)
+
+				exampleCommand := &ffcli.Command{
+					FlagSet: exampleFlagSet,
+					Options: []ff.Option{
+						ff.WithIgnoreUndefined(true),
+						ff.WithConfigFileParser(parser),
+						ff.WithConfigFileFlag("config"),
+					},
+					Exec: func(_ context.Context, args []string) error {
+						return nil
+					},
+				}
+
+				err := exampleCommand.ParseAndRun(context.Background(), []string{
+					"-config", "testdata/server.yml",
+				})
+
+				Expect(err).ToNot(HaveOccurred())
+				Expect(cfg).To(Equal(config.Server{
+					AnalyticsOptOut:          false,
+					Config:                   "testdata/server.yml",
+					LogLevel:                 "info",
+					BadgerLogLevel:           "error",
+					StoragePath:              "/var/lib/pyroscope",
+					APIBindAddr:              ":4040",
+					BaseURL:                  "",
+					CacheEvictThreshold:      0.25,
+					CacheEvictVolume:         0.33,
+					BadgerNoTruncate:         false,
+					DisablePprofEndpoint:     false,
+					MaxNodesSerialization:    2048,
+					MaxNodesRender:           8192,
+					HideApplications:         nil,
+					Retention:                0,
+					SampleRate:               0,
+					OutOfSpaceThreshold:      0,
+					CacheDimensionSize:       0,
+					CacheDictionarySize:      0,
+					CacheSegmentSize:         0,
+					CacheTreeSize:            0,
+					GoogleEnabled:            false,
+					GoogleClientID:           "",
+					GoogleClientSecret:       "",
+					GoogleRedirectURL:        "",
+					GoogleAuthURL:            "https://accounts.google.com/o/oauth2/auth",
+					GoogleTokenURL:           "https://accounts.google.com/o/oauth2/token",
+					GitlabEnabled:            false,
+					GitlabApplicationID:      "",
+					GitlabClientSecret:       "",
+					GitlabRedirectURL:        "",
+					GitlabAuthURL:            "https://gitlab.com/oauth/authorize",
+					GitlabTokenURL:           "https://gitlab.com/oauth/token",
+					GitlabAPIURL:             "https://gitlab.com/api/v4/user",
+					GithubEnabled:            false,
+					GithubClientID:           "",
+					GithubClientSecret:       "",
+					GithubRedirectURL:        "",
+					GithubAuthURL:            "https://github.com/login/oauth/authorize",
+					GithubTokenURL:           "https://github.com/login/oauth/access_token",
+					JWTSecret:                "",
+					LoginMaximumLifetimeDays: 0,
+					MetricExportRules:        nil,
+				}))
+
+				Expect(loadServerConfig(&cfg)).ToNot(HaveOccurred())
+				Expect(cfg.MetricExportRules).To(Equal(config.MetricExportRules{
+					"my_metric_name": {
+						Expr: `app.name{foo=~"bar"}`,
+						Node: "a;b;c",
+					},
+				}))
+			})
+
 			It("agent configuration", func() {
 				exampleFlagSet := flag.NewFlagSet("example flag set", flag.ExitOnError)
 				var cfg config.Agent

pkg/cli/testdata/server.yml

@@ -1,4 +1,4 @@
 metric-export-rules:
-  - name: my_metric_name
+  my_metric_name:
     expr: app.name{foo=~"bar"}
     node: a;b;c

pkg/config/config.go

@@ -117,13 +117,15 @@ type Server struct {
 	JWTSecret                string `deprecated:"true" def:"" desc:"secret used to secure your JWT tokens"`
 	LoginMaximumLifetimeDays int    `deprecated:"true" def:"0" desc:"amount of days after which user will be logged out. 0 means non-expiring."`
 
-	MetricExportRules []MetricExportRule `yaml:"metric-export-rules" deprecated:"true" def:"" desc:"metric export rules"`
+	MetricExportRules MetricExportRules `yaml:"metric-export-rules" def:"" desc:"metric export rules"`
 }
 
+type MetricExportRules map[string]MetricExportRule
+
 type MetricExportRule struct {
-	Name string `def:"" desc:"exported metric name"`
-	Expr string `def:"" desc:"expression in FlameQL syntax to be evaluate against samples"`
-	Node string `def:"total" desc:"tree node filter expression. Should be either 'total' or a valid regexp"`
+	Expr   string   `def:"" desc:"expression in FlameQL syntax to be evaluated against samples"`
+	Node   string   `def:"total" desc:"tree node filter expression. Should be either 'total' or a valid regexp"`
+	Labels []string `def:"" desc:"list of tags to be exported as prometheus labels"`
 }
 
 type Convert struct {

pkg/exporter/exporter.go

@@ -20,36 +20,45 @@ type MetricsExporter struct{ rules []*rule }
 type rule struct {
 	reg prometheus.Registerer
 
-	name string
-	qry  *flameql.Query
+	name   string
+	qry    *flameql.Query
+	labels []string
 	node
 
 	// N.B: CounterVec/MetricVec is not used due to the fact
-	// that label names are not determined.
+	// that label names are not predetermined.
 	sync.RWMutex
 	counters map[uint64]prometheus.Counter
 }
 
 // NewExporter validates configuration and creates a new prometheus MetricsExporter.
-func NewExporter(rules []config.MetricExportRule, reg prometheus.Registerer) (*MetricsExporter, error) {
+func NewExporter(rules config.MetricExportRules, reg prometheus.Registerer) (*MetricsExporter, error) {
 	var e MetricsExporter
-	for _, c := range rules {
-		if !model.IsValidMetricName(model.LabelValue(c.Name)) {
-			return nil, fmt.Errorf("%q is not a valid metric name", c.Name)
+	if rules == nil {
+		return &e, nil
+	}
+	for name, r := range rules {
+		if !model.IsValidMetricName(model.LabelValue(name)) {
+			return nil, fmt.Errorf("%q is not a valid metric name", name)
+		}
+		qry, err := flameql.ParseQuery(r.Expr)
+		if err != nil {
+			return nil, fmt.Errorf("rule %q: invalid expression %q: %w", name, r.Expr, err)
 		}
-		qry, err := flameql.ParseQuery(c.Expr)
+		n, err := newNode(r.Node)
 		if err != nil {
-			return nil, fmt.Errorf("rule %q: invalid expression %q: %w", c.Name, c.Expr, err)
+			return nil, fmt.Errorf("rule %q: invalid node %q: %w", name, r.Node, err)
 		}
-		n, err := newNode(c.Node)
+		g, err := validateTagKeys(r.Labels)
 		if err != nil {
-			return nil, fmt.Errorf("rule %q: invalid node %q: %w", c.Name, c.Node, err)
+			return nil, fmt.Errorf("rule %q: invalid tags to group by %q: %w", name, r.Labels, err)
 		}
 		e.rules = append(e.rules, &rule{
-			name:     c.Name,
+			name:     name,
 			qry:      qry,
 			reg:      reg,
 			node:     n,
+			labels:   g,
 			counters: make(map[uint64]prometheus.Counter),
 		})
 	}
@@ -80,7 +89,7 @@ func (e MetricsExporter) Observe(k *segment.Key, tree *tree.Tree) {
 // eval returns existing counter for the key or creates a new one,
 // if the key satisfies the rule expression.
 func (r *rule) eval(k *segment.Key) (prometheus.Counter, bool) {
-	m, ok := r.matchedLabels(k)
+	m, ok := r.matchLabelNames(k)
 	if !ok {
 		return nil, false
 	}
@@ -93,12 +102,9 @@ func (r *rule) eval(k *segment.Key) (prometheus.Counter, bool) {
 	}
 	r.RUnlock()
 	if match(r.qry, m) {
-		// Remove app name label to avoid
-		// collision with prometheus labels.
-		m = m[1:]
 		c = prometheus.NewCounter(prometheus.CounterOpts{
 			Name:        r.name,
-			ConstLabels: m.labels(),
+			ConstLabels: promLabels(k, r.labels...),
 		})
 		r.reg.MustRegister(c)
 		r.Lock()
@@ -109,8 +115,33 @@ func (r *rule) eval(k *segment.Key) (prometheus.Counter, bool) {
 	return nil, false
 }
 
+func validateTagKeys(tagKeys []string) ([]string, error) {
+	for _, l := range tagKeys {
+		if err := flameql.ValidateTagKey(l); err != nil {
+			return nil, err
+		}
+	}
+	return tagKeys, nil
+}
+
+// promLabels converts key to prometheus.Labels ignoring reserved tag keys.
+func promLabels(key *segment.Key, labels ...string) prometheus.Labels {
+	if len(labels) == 0 {
+		return nil
+	}
+	l := key.Labels()
+	p := make(prometheus.Labels, len(labels))
+	// labels are guarantied to be valid.
+	for _, k := range labels {
+		if v, ok := l[k]; ok {
+			p[k] = v
+		}
+	}
+	return p
+}
+
 // match reports whether the key matches the query.
-func match(qry *flameql.Query, labels matchedLabels) bool {
+func match(qry *flameql.Query, labels labels) bool {
 	for _, m := range qry.Matchers {
 		var ok bool
 		for _, l := range labels {

pkg/exporter/exporter_test.go

@@ -1,6 +1,7 @@
 package exporter
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -82,21 +83,21 @@ func TestRegister(t *testing.T) {
 }
 
 func TestObserve(t *testing.T) {
-	rules := []config.MetricExportRule{
-		{
-			Name: "app_name_cpu_total",
-			Expr: `app.name.cpu{foo="bar"}`,
-			Node: "total",
+	rules := config.MetricExportRules{
+		"app_name_cpu_total": {
+			Expr:   `app.name.cpu{foo="bar"}`,
+			Node:   "total",
+			Labels: []string{"foo"},
 		},
-		{
-			Name: "app_name_cpu_abc",
-			Expr: `app.name.cpu{foo=~"b.*"}`,
-			Node: "^a;b;c$",
+		"app_name_cpu_abc": {
+			Expr:   `app.name.cpu{foo=~"b.*"}`,
+			Node:   "^a;b;c$",
+			Labels: []string{"foo"},
 		},
-		{
-			Name: "app_name_cpu_ab",
-			Expr: `app.name.cpu{foo=~"b.*"}`,
-			Node: "a;b",
+		"app_name_cpu_ab": {
+			Expr:   `app.name.cpu{foo=~"b.*"}`,
+			Node:   "a;b",
+			Labels: []string{"foo"},
 		},
 	}
 
@@ -105,23 +106,85 @@ func TestObserve(t *testing.T) {
 	v := createTree()
 	exporter.Observe(k, v)
 
-	// Hashes are predetermined.
-	total := testutil.ToFloat64(exporter.rules[0].counters[16252301464360304376])
-	if total != 5 {
-		t.Fatalf("total counter must be 5, got %v", total)
+	if total := getRuleCounterValue(exporter, "app_name_cpu_total", k); total != 5 {
+		t.Fatalf("Total counter must be 5, got %v", total)
 	}
 
-	abc := testutil.ToFloat64(exporter.rules[1].counters[16252301464360304376])
-	if abc != 2 {
+	if abc := getRuleCounterValue(exporter, "app_name_cpu_abc", k); abc != 2 {
 		t.Fatalf("a;b;c counter must be 2, got %v", abc)
 	}
 
-	ab := testutil.ToFloat64(exporter.rules[2].counters[16252301464360304376])
-	if ab != 3 {
+	if ab := getRuleCounterValue(exporter, "app_name_cpu_ab", k); ab != 3 {
 		t.Fatalf("a;b counter must be 3, got %v", ab)
 	}
 }
 
+func TestGroupBy(t *testing.T) {
+	const rule = "app_name_cpu_total"
+	rules := config.MetricExportRules{
+		rule: {
+			Expr:   `app.name.cpu{foo=~"bar"}`,
+			Labels: []string{"foo"},
+		},
+	}
+
+	exporter, _ := NewExporter(rules, prometheus.NewRegistry())
+	k1 := observe(exporter, `app.name.cpu{foo=bar_a,bar=a}`)
+	k2 := observe(exporter, `app.name.cpu{foo=bar_a,bar=b}`)
+	k3 := observe(exporter, `app.name.cpu{foo=bar_b,bar=c}`)
+
+	counters := len(exporter.rules[0].counters)
+	if counters != 2 {
+		t.Fatalf("Expected 2 counters, got %v", counters)
+	}
+
+	c1 := getRuleCounter(exporter, rule, k1)
+	c2 := getRuleCounter(exporter, rule, k2)
+	c3 := getRuleCounter(exporter, rule, k3)
+
+	if !reflect.DeepEqual(c1, c2) {
+		t.Fatalf("Expected c1 and c2 is the same counter")
+	}
+
+	if t1 := testutil.ToFloat64(c1); t1 != 10 {
+		t.Fatalf("Total counter for k1 must be 10, got %v", t1)
+	}
+
+	if t2 := testutil.ToFloat64(c2); t2 != 10 {
+		t.Fatalf("Total counter for k2 must be 10, got %v", t2)
+	}
+
+	if t3 := testutil.ToFloat64(c3); t3 != 5 {
+		t.Fatalf("Total counter for k3 must be 5, got %v", t3)
+	}
+}
+
+func getRuleCounter(e *MetricsExporter, name string, k *segment.Key) prometheus.Counter {
+	for _, r := range e.rules {
+		if r.name != name {
+			continue
+		}
+		m, ok := r.matchLabelNames(k)
+		if !ok {
+			continue
+		}
+		if c, ok := r.counters[m.hash()]; ok {
+			return c
+		}
+	}
+	return nil
+}
+
+func getRuleCounterValue(e *MetricsExporter, name string, k *segment.Key) float64 {
+	return testutil.ToFloat64(getRuleCounter(e, name, k))
+}
+
+func observe(e *MetricsExporter, key string) *segment.Key {
+	k, _ := segment.ParseKey(key)
+	e.Observe(k, createTree())
+	return k
+}
+
 func createTree() *tree.Tree {
 	t := tree.New()
 	t.Insert([]byte("a;b"), uint64(1))