feat: support speedscope format

pkg/convert/speedscope/json.go

@@ -0,0 +1,65 @@
+package speedscope
+
+// Description of Speedscope JSON
+// See spec: https://github.com/jlfwong/speedscope/blob/main/src/lib/file-format-spec.ts
+
+const (
+	schema = "https://www.speedscope.app/file-format-schema.json"
+
+	profileEvented = "evented"
+	profileSampled = "sampled"
+
+	unitNone         = "none"
+	unitNanoseconds  = "nanoseconds"
+	unitMicroseconds = "microseconds"
+	unitMilliseconds = "milliseconds"
+	unitSeconds      = "seconds"
+	unitBytes        = "bytes"
+
+	eventOpen  = "O"
+	eventClose = "C"
+)
+
+type speedscopeFile struct {
+	Schema             string `json:"$schema"`
+	Shared             shared
+	Profiles           []profile
+	Name               string
+	ActiveProfileIndex float64
+	Exporter           string
+}
+
+type shared struct {
+	Frames []frame
+}
+
+type frame struct {
+	Name string
+	File string
+	Line float64
+	Col  float64
+}
+
+type profile struct {
+	Type       string
+	Name       string
+	Unit       string
+	StartValue float64
+	EndValue   float64
+
+	// Evented profile
+	Events []event
+
+	// Sample profile
+	Samples []sample
+	Weights []float64
+}
+
+type event struct {
+	Type  string
+	At    float64
+	Frame float64
+}
+
+// Indexes into Frames
+type sample []float64

pkg/convert/speedscope/parser.go

@@ -0,0 +1,190 @@
+package speedscope
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
+	"github.com/pyroscope-io/pyroscope/pkg/storage"
+	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
+	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
+	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
+)
+
+// RawProfile implements ingestion.RawProfile for Speedscope format
+type RawProfile struct {
+	RawData []byte
+}
+
+// Parse parses a profile
+func (p *RawProfile) Parse(ctx context.Context, putter storage.Putter, _ storage.MetricsExporter, md ingestion.Metadata) error {
+	profiles, err := parseAll(p.RawData, md)
+	if err != nil {
+		return err
+	}
+
+	for _, putInput := range profiles {
+		err = putter.Put(ctx, putInput)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func parseAll(rawData []byte, md ingestion.Metadata) ([]*storage.PutInput, error) {
+	file := speedscopeFile{}
+	err := json.Unmarshal(rawData, &file)
+	if err != nil {
+		return nil, err
+	}
+	if file.Schema != schema {
+		return nil, fmt.Errorf("Unknown schema: %s", file.Schema)
+	}
+
+	results := make([]*storage.PutInput, 0, len(file.Profiles))
+	// Not a pointer, we _want_ to copy on call
+	input := storage.PutInput{
+		StartTime:  md.StartTime,
+		EndTime:    md.EndTime,
+		SpyName:    md.SpyName,
+		SampleRate: md.SampleRate,
+		Key:        md.Key,
+	}
+
+	for _, prof := range file.Profiles {
+		putInput, err := parseOne(&prof, input, file.Shared.Frames, len(file.Profiles) > 1)
+		if err != nil {
+			return nil, err
+		}
+		results = append(results, putInput)
+	}
+	return results, nil
+}
+
+func parseOne(prof *profile, putInput storage.PutInput, frames []frame, multi bool) (*storage.PutInput, error) {
+	// Fixup some metadata
+	putInput.Units = chooseUnit(prof.Unit)
+	putInput.AggregationType = metadata.SumAggregationType
+	if multi {
+		putInput.Key = chooseKey(putInput.Key, prof.Unit)
+	}
+	// Don't override sampleRate. Sometimes units corresponds to that, but not necessarily.
+
+	var err error
+	tr := tree.New()
+	switch prof.Type {
+	case profileEvented:
+		err = parseEvented(tr, prof, frames)
+	case profileSampled:
+		err = parseSampled(tr, prof, frames)
+	default:
+		return nil, fmt.Errorf("Profile type %s not supported", prof.Type)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	putInput.Val = tr
+	return &putInput, nil
+}
+
+func chooseUnit(unit string) metadata.Units {
+	switch unit {
+	case unitBytes:
+		return metadata.BytesUnits
+	default:
+		return metadata.SamplesUnits
+	}
+}
+
+func chooseKey(orig *segment.Key, unit string) *segment.Key {
+	// This means we'll have duplicate keys if multiple profiles have the same units. Probably ok.
+	name := fmt.Sprintf("%s.%s", orig.AppName(), unit)
+	result := orig.Clone()
+	result.Add("__name__", name)
+	return result
+}
+
+func parseEvented(tr *tree.Tree, prof *profile, frames []frame) error {
+	last := prof.StartValue
+	indexStack := []int{}
+	nameStack := []string{}
+
+	for _, ev := range prof.Events {
+		if ev.At < last {
+			return fmt.Errorf("Events out of order, %f < %f", ev.At, last)
+		}
+		fid := int(ev.Frame)
+		if fid < 0 || fid >= len(frames) {
+			return fmt.Errorf("Invalid frame %d", fid)
+		}
+
+		if ev.Type == eventClose {
+			if len(indexStack) == 0 {
+				return fmt.Errorf("No stack to close at %f", ev.At)
+			}
+			lastIdx := len(indexStack) - 1
+			if indexStack[lastIdx] != fid {
+				return fmt.Errorf("Closing non-open frame %d", fid)
+			}
+
+			// Close this frame
+			tr.InsertStackString(nameStack, uint64(ev.At-last))
+			indexStack = indexStack[:lastIdx]
+			nameStack = nameStack[:lastIdx]
+		} else if ev.Type == eventOpen {
+			// Add any time up til now
+			if len(nameStack) > 0 {
+				tr.InsertStackString(nameStack, uint64(ev.At-last))
+			}
+
+			// Open the frame
+			indexStack = append(indexStack, fid)
+			nameStack = append(nameStack, frames[fid].Name)
+		} else {
+			return fmt.Errorf("Unknown event type %s", ev.Type)
+		}
+
+		last = ev.At
+	}
+
+	return nil
+}
+
+func parseSampled(tr *tree.Tree, prof *profile, frames []frame) error {
+	if len(prof.Samples) != len(prof.Weights) {
+		return fmt.Errorf("Unequal lengths of samples and weights: %d != %d", len(prof.Samples), len(prof.Weights))
+	}
+
+	stack := []string{}
+	for i, samp := range prof.Samples {
+		weight := prof.Weights[i]
+		if weight < 0 {
+			return fmt.Errorf("Negative weight %f", weight)
+		}
+
+		for _, frameID := range samp {
+			fid := int(frameID)
+			if fid < 0 || fid > len(frames) {
+				return fmt.Errorf("Invalid frame %d", fid)
+			}
+			stack = append(stack, frames[fid].Name)
+		}
+		tr.InsertStackString(stack, uint64(weight))
+
+		stack = stack[:0] // clear, but retain memory
+	}
+	return nil
+}
+
+// Bytes returns the raw bytes of the profile
+func (p *RawProfile) Bytes() ([]byte, error) {
+	return p.RawData, nil
+}
+
+// ContentType returns the HTTP ContentType of the profile
+func (*RawProfile) ContentType() string {
+	return "application/json"
+}

pkg/convert/speedscope/speedscope_suite_test.go

@@ -0,0 +1,13 @@
+package speedscope_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestConvert(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Speedscope Suite")
+}

pkg/convert/speedscope/speedscope_test.go

@@ -0,0 +1,75 @@
+package speedscope
+
+import (
+	"context"
+	"os"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
+	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
+	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
+
+	"github.com/pyroscope-io/pyroscope/pkg/storage"
+)
+
+type mockIngester struct{ actual []*storage.PutInput }
+
+func (m *mockIngester) Put(_ context.Context, p *storage.PutInput) error {
+	m.actual = append(m.actual, p)
+	return nil
+}
+
+var _ = Describe("Speedscope", func() {
+	It("Can parse an event-format profile", func() {
+		data, err := os.ReadFile("testdata/simple.speedscope.json")
+		Expect(err).ToNot(HaveOccurred())
+
+		key, err := segment.ParseKey("foo")
+		Expect(err).ToNot(HaveOccurred())
+
+		ingester := new(mockIngester)
+		profile := &RawProfile{RawData: data}
+
+		md := ingestion.Metadata{Key: key}
+		err = profile.Parse(context.Background(), ingester, nil, md)
+		Expect(err).ToNot(HaveOccurred())
+
+		Expect(ingester.actual).To(HaveLen(1))
+		input := ingester.actual[0]
+
+		Expect(input.Units).To(Equal(metadata.SamplesUnits))
+		Expect(input.Key.Normalized()).To(Equal("foo{}"))
+		expectedResult := `a;b 5
+a;b;c 5
+a;b;d 4
+`
+		Expect(input.Val.String()).To(Equal(expectedResult))
+	})
+
+	It("Can parse a sample-format profile", func() {
+		data, err := os.ReadFile("testdata/two-sampled.speedscope.json")
+		Expect(err).ToNot(HaveOccurred())
+
+		key, err := segment.ParseKey("foo{x=y}")
+		Expect(err).ToNot(HaveOccurred())
+
+		ingester := new(mockIngester)
+		profile := &RawProfile{RawData: data}
+
+		md := ingestion.Metadata{Key: key}
+		err = profile.Parse(context.Background(), ingester, nil, md)
+		Expect(err).ToNot(HaveOccurred())
+
+		Expect(ingester.actual).To(HaveLen(2))
+
+		input := ingester.actual[0]
+		Expect(input.Units).To(Equal(metadata.SamplesUnits))
+		Expect(input.Key.Normalized()).To(Equal("foo.seconds{x=y}"))
+		expectedResult := `a;b 5
+a;b;c 5
+a;b;d 4
+`
+		Expect(input.Val.String()).To(Equal(expectedResult))
+	})
+})