fix: Handle Speedscope sample types in ingestion (#4568)

Author: jake-kramer

pkg/ingester/pyroscope/ingest_adapter.go

@@ -283,6 +283,14 @@ func convertMetadata(pi *storage.PutInput) (metricName, stType, stUnit, app stri
 		metricName = "exceptions"
 		stType = stTypeSamples
 		stUnit = stUnitCount
+	case "seconds", "nanoseconds", "microseconds", "milliseconds":
+		metricName = metricWall
+		stType = stTypeSamples
+		stUnit = stUnitCount
+	case "bytes":
+		metricName = metricMemory
+		stType = stTypeSamples
+		stUnit = stUnitBytes
 	default:
 		err = fmt.Errorf("unknown profile type: %s", stType)
 	}

pkg/og/convert/speedscope/testdata/two-sampled-bytes.speedscope.json

@@ -0,0 +1,46 @@
+{
+  "exporter": "[email protected]",
+  "$schema": "https://www.speedscope.app/file-format-schema.json",
+  "name": "Two Samples",
+  "activeProfileIndex": 1,
+  "profiles": [
+    {
+      "type": "sampled",
+      "name": "one",
+      "unit": "bytes",
+      "startValue": 0,
+      "endValue": 14,
+      "samples": [
+        [0, 1, 2],
+        [0, 1, 2],
+        [0, 1, 3],
+        [0, 1, 2],
+        [0, 1]
+      ],
+      "weights": [1, 1, 4, 3, 5]
+    },
+    {
+      "type": "sampled",
+      "name": "two",
+      "unit": "bytes",
+      "startValue": 0,
+      "endValue": 14,
+      "samples": [
+        [0, 1, 2],
+        [0, 1, 2],
+        [0, 1, 3],
+        [0, 1, 2],
+        [0, 1]
+      ],
+      "weights": [1, 1, 4, 3, 5]
+    }
+  ],
+  "shared": {
+    "frames": [
+      { "name": "a" },
+      { "name": "b" },
+      { "name": "c" },
+      { "name": "d" }
+    ]
+  }
+}

pkg/test/integration/helper.go

@@ -244,7 +244,7 @@ func (p *PyroscopeTest) NewRequestBuilder(t *testing.T) *RequestBuilder {
 }
 
 func (p *PyroscopeTest) TempAppName() string {
-	return fmt.Sprintf("pprof.integration.%d",
+	return fmt.Sprintf("pprof-integration-%d",
 		rand.Uint64())
 }
 
@@ -362,6 +362,18 @@ func (b *RequestBuilder) IngestJFRRequestBody(jfr []byte, labels []byte) *http.R
 	return req
 }
 
+func (b *RequestBuilder) IngestSpeedscopeRequest(speedscopePath string) *http.Request {
+	speedscopeData, err := os.ReadFile(speedscopePath)
+	require.NoError(b.t, err)
+
+	url := b.url + "/ingest?name=" + b.AppName + "&format=speedscope"
+	req, err := http.NewRequest("POST", url, bytes.NewReader(speedscopeData))
+	require.NoError(b.t, err)
+	req.Header.Set("Content-Type", "application/json")
+
+	return req
+}
+
 func (b *RequestBuilder) Render(metric string) *flamebearer.FlamebearerProfile {
 	queryURL := b.url + "/pyroscope/render?query=" + createRenderQuery(metric, b.AppName) + "&from=946656000&until=now&format=collapsed"
 	fmt.Println(queryURL)

pkg/test/integration/ingest_speedscope_test.go

@@ -0,0 +1,80 @@
+package integration
+
+import (
+	"testing"
+
+	"connectrpc.com/connect"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
+)
+
+type speedscopeTestDataStruct struct {
+	name            string
+	speedscopeFile  string
+	expectStatus    int
+	expectedMetrics []expectedMetric
+}
+
+const (
+	testdataDirSpeedscope = repoRoot + "pkg/og/convert/speedscope/testdata"
+)
+
+var (
+	speedscopeTestData = []speedscopeTestDataStruct{
+		{
+			name:           "single profile evented",
+			speedscopeFile: testdataDirSpeedscope + "/simple.speedscope.json",
+			expectStatus:   200,
+			expectedMetrics: []expectedMetric{
+				// The difference between the metric name here and in the other test is a quirk in
+				// how the speedscope parsing logic. Only multi profile uploads will
+				// append the unit to the metric name which is parsed differently downstream.
+				{"process_cpu:cpu:nanoseconds:cpu:nanoseconds", nil, 0},
+			},
+		},
+		{
+			name:           "multi profile sampled",
+			speedscopeFile: testdataDirSpeedscope + "/two-sampled.speedscope.json",
+			expectStatus:   200,
+			expectedMetrics: []expectedMetric{
+				{"wall:wall:nanoseconds:cpu:nanoseconds", nil, 0},
+			},
+		},
+		{
+			name:           "multi profile samples bytes units",
+			speedscopeFile: testdataDirSpeedscope + "/two-sampled-bytes.speedscope.json",
+			expectStatus:   200,
+			expectedMetrics: []expectedMetric{
+				{"memory:samples:bytes::", nil, 0},
+			},
+		},
+	}
+)
+
+func TestIngestSpeedscope(t *testing.T) {
+	EachPyroscopeTest(t, func(p *PyroscopeTest, t *testing.T) {
+		for _, td := range speedscopeTestData {
+			t.Run(td.name, func(t *testing.T) {
+				rb := p.NewRequestBuilder(t)
+				req := rb.IngestSpeedscopeRequest(td.speedscopeFile)
+				p.Ingest(t, req, td.expectStatus)
+
+				if td.expectStatus == 200 {
+					for _, metric := range td.expectedMetrics {
+						rb.Render(metric.name)
+						profile := rb.SelectMergeProfile(metric.name, metric.query)
+						assertSpeedscopeProfile(t, profile)
+					}
+				}
+			})
+		}
+	})
+}
+
+func assertSpeedscopeProfile(t *testing.T, resp *connect.Response[profilev1.Profile]) {
+	assert.Equal(t, 1, len(resp.Msg.SampleType), "SampleType should be set")
+	require.Greater(t, len(resp.Msg.Sample), 0, "Profile should contain samples")
+	assert.Greater(t, resp.Msg.Sample[0].Value[0], int64(0), "Sample value should be positive")
+}