add support for Pyroscope

mise.lock

@@ -1,10 +1,10 @@
 [tools.go]
-version = "1.24.1"
+version = "1.24.2"
 backend = "core:go"
 
 [tools.golangci-lint]
-version = "2.0.2"
+version = "2.1.1"
 backend = "aqua:golangci/golangci-lint"
 
 [tools.golangci-lint.checksums]
-"golangci-lint-2.0.2-linux-amd64.tar.gz" = "sha256:89cc8a7810dc63b9a37900da03e37c3601caf46d42265d774e0f1a5d883d53e2"
+"golangci-lint-2.1.1-linux-amd64.tar.gz" = "sha256:7167df345d0146d662b12ba068306f843d0eba408f9dc8f4d3ebb239786e87da"

observability/observability_endpoint.go

@@ -14,7 +14,7 @@ type Endpoint interface {
 	SearchTags(context.Context, map[string]string) ([]byte, error)
 
 	// Metrics
-	RunPromQL(context.Context, string) ([]byte, error)
+	RunPromQL(string) ([]byte, error)
 
 	Start(context.Context) error
 	Stop(context.Context) error

testhelpers/kubernetes/kubernetes.go

@@ -121,5 +121,9 @@ func start(model *Kubernetes, ports remote.PortsConfig, testName string, run fun
 	if err != nil {
 		return err
 	}
+	err = portForward(ports.PyroscopeHttpPort, 4040)
+	if err != nil {
+		return err
+	}
 	return nil
 }

testhelpers/prometheus/responses/models.go

@@ -1,6 +1,6 @@
 package responses
 
-type QueryResult struct {
+type PrometheusQueryResult struct {
 	Status string `json:"status"`
 	Data   Data   `json:"data"`
 }

testhelpers/prometheus/responses/response_helpers.go

@@ -6,7 +6,7 @@ import (
 )
 
 func ParseQueryOutput(body []byte) ([]Result, error) {
-	qr := QueryResult{}
+	qr := PrometheusQueryResult{}
 	if err := json.Unmarshal(body, &qr); err != nil {
 		return nil, fmt.Errorf("decoding Prometheus response: %w", err)
 	}

testhelpers/remote/remote_observability_endpoint.go

@@ -24,6 +24,7 @@ type PortsConfig struct {
 	MimirHTTPPort      int
 	PrometheusHTTPPort int
 	LokiHttpPort       int
+	PyroscopeHttpPort  int
 }
 
 type Endpoint struct {
@@ -101,8 +102,7 @@ func (e *Endpoint) GetTraceByID(ctx context.Context, id string) ([]byte, error)
 		return nil, ctx.Err()
 	}
 
-	url := fmt.Sprintf("http://localhost:%d/api/traces/%s", e.ports.TempoHTTPPort, id)
-	return e.makeGetRequest(url)
+	return e.makeGetRequest(fmt.Sprintf("http://localhost:%d/api/traces/%s", e.ports.TempoHTTPPort, id))
 }
 
 func (e *Endpoint) SearchTempo(ctx context.Context, query string) ([]byte, error) {
@@ -128,12 +128,10 @@ func (e *Endpoint) SearchTags(ctx context.Context, tags map[string]string) ([]by
 		tb.WriteString(url.QueryEscape(s))
 	}
 
-	url := fmt.Sprintf("http://localhost:%d/api/search?tags=%s", e.ports.TempoHTTPPort, tb.String())
-
-	return e.makeGetRequest(url)
+	return e.makeGetRequest(fmt.Sprintf("http://localhost:%d/api/search?tags=%s", e.ports.TempoHTTPPort, tb.String()))
 }
 
-func (e *Endpoint) RunPromQL(ctx context.Context, promQL string) ([]byte, error) {
+func (e *Endpoint) RunPromQL(promQL string) ([]byte, error) {
 	var u string
 	if e.ports.MimirHTTPPort != 0 {
 		u = fmt.Sprintf("http://localhost:%d/prometheus/api/v1/query?query=%s", e.ports.MimirHTTPPort, url.PathEscape(promQL))
@@ -184,6 +182,30 @@ func (e *Endpoint) SearchLoki(query string) ([]byte, error) {
 	return body, nil
 }
 
+func (e *Endpoint) SearchPyroscope(query string) ([]byte, error) {
+	if e.ports.PyroscopeHttpPort == 0 {
+		return nil, fmt.Errorf("to search Pyroscope you must configure a PyroscopeHttpPort")
+	}
+
+	u := fmt.Sprintf("http://localhost:%d/pyroscope/render?from=from=now-1m&query=%s", e.ports.PyroscopeHttpPort, url.PathEscape(query))
+
+	resp, err := http.Get(u)
+	if err != nil {
+		return nil, fmt.Errorf("querying pyroscope: %w", err)
+	}
+
+	defer func(Body io.ReadCloser) {
+		_ = Body.Close()
+	}(resp.Body)
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("can't read response body: %w", err)
+	}
+
+	return body, nil
+}
+
 func (e *Endpoint) Start(ctx context.Context) error {
 	return e.start(ctx)
 }

yaml/docker-compose-docker-lgtm-template.yml

@@ -7,3 +7,4 @@ services:
       - "{{ .PrometheusHTTPPort }}:9090"
       - "{{ .TempoHTTPPort }}:3200"
       - "{{ .LokiHTTPPort }}:3100"
+      - "{{ .PyroscopeHttpPort }}:4040"

yaml/generator.go

@@ -43,6 +43,7 @@ func (c *TestCase) generateDockerComposeFile() []byte {
 	vars["PrometheusHTTPPort"] = c.PortConfig.PrometheusHTTPPort
 	vars["LokiHTTPPort"] = c.PortConfig.LokiHTTPPort
 	vars["TempoHTTPPort"] = c.PortConfig.TempoHTTPPort
+	vars["PyroscopeHttpPort"] = c.PortConfig.PyroscopeHttpPort
 	vars["LgtmVersion"] = c.LgtmVersion
 
 	env := os.Environ()
@@ -98,7 +99,11 @@ func (c *TestCase) generateDockerComposeFile() []byte {
 	args := []string{"compose", "-f", f.Name(), "config"}
 	cmd := exec.Command("docker", args...)
 	cmd.Env = env
+	cmd.Stderr = os.Stderr
 	content, err := cmd.Output()
+	if err != nil {
+		slog.Error("failed to run docker compose", "error", err)
+	}
 	gomega.Expect(err).ToNot(gomega.HaveOccurred())
 	return content
 }

yaml/logs.go

@@ -6,7 +6,7 @@ import (
 	"log/slog"
 )
 
-type QueryResponse struct {
+type LokiQueryResponse struct {
 	Status string `json:"status"`
 	Data   struct {
 		Result []struct {
@@ -28,7 +28,7 @@ func AssertLokiResponse(b []byte, l ExpectedLogs, r *runner) {
 	g := r.gomegaInst
 	g.Expect(len(b)).Should(gomega.BeNumerically(">", 0), "expected loki response to be non-empty")
 
-	response := QueryResponse{}
+	response := LokiQueryResponse{}
 	err := json.Unmarshal(b, &response)
 	if err != nil {
 		slog.Info("error unmarshalling loki", "response", string(b))

yaml/metrics.go

@@ -1,7 +1,6 @@
 package yaml
 
 import (
-	"context"
 	"strconv"
 	"strings"
 
@@ -20,8 +19,7 @@ func replaceVariables(promQL string) string {
 
 func AssertProm(r *runner, promQL string, value string) {
 	promQL = replaceVariables(promQL)
-	ctx := context.Background()
-	b, err := r.endpoint.RunPromQL(ctx, promQL)
+	b, err := r.endpoint.RunPromQL(promQL)
 	r.LogQueryResult("promQL query %v response %v err=%v\n", promQL, string(b), err)
 	g := r.gomegaInst
 	g.Expect(err).ToNot(gomega.HaveOccurred())

yaml/model.go

@@ -43,6 +43,15 @@ type ExpectedLogs struct {
 	NoExtraAttributes bool              `yaml:"no-extra-attributes"`
 }
 
+type Flamebearers struct {
+	Contains string `yaml:"contains"`
+}
+
+type ExpectedProfiles struct {
+	Query        string       `yaml:"query"`
+	Flamebearers Flamebearers `yaml:"flamebearers"`
+}
+
 type ExpectedTraces struct {
 	TraceQL string         `yaml:"traceql"`
 	Spans   []ExpectedSpan `yaml:"spans"`
@@ -53,11 +62,12 @@ type CustomCheck struct {
 }
 
 type Expected struct {
-	ComposeLogs  []string          `yaml:"compose-logs"`
-	Logs         []ExpectedLogs    `yaml:"logs"`
-	Traces       []ExpectedTraces  `yaml:"traces"`
-	Metrics      []ExpectedMetrics `yaml:"metrics"`
-	CustomChecks []CustomCheck     `yaml:"custom-checks"`
+	ComposeLogs  []string           `yaml:"compose-logs"`
+	Logs         []ExpectedLogs     `yaml:"logs"`
+	Traces       []ExpectedTraces   `yaml:"traces"`
+	Metrics      []ExpectedMetrics  `yaml:"metrics"`
+	Profiles     []ExpectedProfiles `yaml:"profiles"`
+	CustomChecks []CustomCheck      `yaml:"custom-checks"`
 }
 
 type DockerCompose struct {
@@ -85,6 +95,7 @@ func (d *TestCaseDefinition) Merge(other TestCaseDefinition) {
 	d.Expected.Logs = append(d.Expected.Logs, other.Expected.Logs...)
 	d.Expected.Traces = append(d.Expected.Traces, other.Expected.Traces...)
 	d.Expected.Metrics = append(d.Expected.Metrics, other.Expected.Metrics...)
+	d.Expected.Profiles = append(d.Expected.Profiles, other.Expected.Profiles...)
 	d.Expected.CustomChecks = append(d.Expected.CustomChecks, other.Expected.CustomChecks...)
 	if d.DockerCompose == nil {
 		d.DockerCompose = other.DockerCompose
@@ -98,6 +109,7 @@ type PortConfig struct {
 	PrometheusHTTPPort int
 	LokiHTTPPort       int
 	TempoHTTPPort      int
+	PyroscopeHttpPort  int
 }
 
 type TestCase struct {
@@ -130,7 +142,7 @@ func (c *TestCase) validateAndSetVariables() {
 	}
 	validateInput(c.Definition.Input)
 	expected := c.Definition.Expected
-	gomega.Expect(len(expected.Metrics) == 0 && len(expected.Traces) == 0 && len(expected.Logs) == 0).To(gomega.BeFalse())
+	gomega.Expect(len(expected.Metrics) == 0 && len(expected.Traces) == 0 && len(expected.Logs) == 0 && len(expected.Profiles) == 0).To(gomega.BeFalse())
 
 	for _, c := range expected.CustomChecks {
 		gomega.Expect(c.Script).ToNot(gomega.BeEmpty(), "script is empty in "+string(c.Script))
@@ -160,6 +172,11 @@ func (c *TestCase) validateAndSetVariables() {
 			}
 		}
 	}
+	for _, p := range expected.Profiles {
+		out, _ := yaml.Marshal(p)
+		gomega.Expect(p.Query).ToNot(gomega.BeEmpty(), "query is empty in "+string(out))
+		gomega.Expect(p.Flamebearers.Contains).ToNot(gomega.BeEmpty(), "Flamebearers.contains is empty in "+string(out))
+	}
 
 	if c.PortConfig == nil {
 		// We're in non-parallel mode, so we can static ports here.
@@ -169,6 +186,7 @@ func (c *TestCase) validateAndSetVariables() {
 			PrometheusHTTPPort: 9090,
 			LokiHTTPPort:       3100,
 			TempoHTTPPort:      3200,
+			PyroscopeHttpPort:  4040,
 		}
 	}
 
@@ -177,6 +195,7 @@ func (c *TestCase) validateAndSetVariables() {
 		"prometheus", c.PortConfig.PrometheusHTTPPort,
 		"loki", c.PortConfig.LokiHTTPPort,
 		"tempo", c.PortConfig.TempoHTTPPort,
+		"pyroscope", c.PortConfig.PyroscopeHttpPort,
 		"application", c.PortConfig.ApplicationPort)
 }
 

yaml/profiles.go

@@ -0,0 +1,36 @@
+package yaml
+
+import (
+	"encoding/json"
+	"log/slog"
+
+	"github.com/onsi/gomega"
+)
+
+type PyroscopeQueryResponse struct {
+	Flamebearer struct {
+		Names []string `json:"names"`
+	} `json:"flamebearer"`
+}
+
+func AssertPyroscope(r *runner, p ExpectedProfiles) {
+	b, err := r.endpoint.SearchPyroscope(p.Query)
+	r.LogQueryResult("query %v response %v err=%v\n", p.Query, string(b), err)
+	g := r.gomegaInst
+	g.Expect(err).ToNot(gomega.HaveOccurred())
+	assertPyroscopeResponse(b, p, r)
+}
+
+func assertPyroscopeResponse(b []byte, p ExpectedProfiles, r *runner) {
+	g := r.gomegaInst
+	g.Expect(len(b)).Should(gomega.BeNumerically(">", 0), "expected pyroscope response to be non-empty")
+
+	response := PyroscopeQueryResponse{}
+	err := json.Unmarshal(b, &response)
+	if err != nil {
+		slog.Info("error unmarshalling pyroscope", "response", string(b))
+	}
+
+	g.Expect(err).ToNot(gomega.HaveOccurred())
+	g.Expect(response.Flamebearer.Names).To(gomega.ContainElement(gomega.ContainSubstring(p.Flamebearers.Contains)))
+}