feat(chunk): Add attachments to profile chunks

internal/chunk/android.go

@@ -97,6 +97,11 @@ func (c AndroidChunk) GetProjectID() uint64 {
 	return c.ProjectID
 }
 
+// Attachments are only supported for sample chunks.
+func (c AndroidChunk) GetAttachments() []Attachment {
+	return nil
+}
+
 func (c AndroidChunk) GetReceived() float64 {
 	return c.Received
 }

internal/chunk/chunk.go

@@ -18,6 +18,7 @@ type (
 		GetPlatform() platform.Platform
 		GetProfilerID() string
 		GetProjectID() uint64
+		GetAttachments() []Attachment
 		GetReceived() float64
 		GetRelease() string
 		GetRetentionDays() int
@@ -107,6 +108,10 @@ func (c Chunk) GetProjectID() uint64 {
 	return c.chunk.GetProjectID()
 }
 
+func (c Chunk) GetAttachments() []Attachment {
+	return c.chunk.GetAttachments()
+}
+
 func (c Chunk) GetReceived() float64 {
 	return c.chunk.GetReceived()
 }

internal/chunk/chunk_test.go

@@ -0,0 +1,122 @@
+package chunk
+
+import (
+	"bytes"
+	"encoding/json"
+	"testing"
+
+	"github.com/getsentry/vroom/internal/platform"
+	"github.com/getsentry/vroom/internal/testutil"
+)
+
+func TestAttachments(t *testing.T) {
+	payload := `{
+		"version": "2",
+		"chunk_id": "0432a0a4c25f4697bf9f0a2fcbe6a814",
+		"attachments": [
+			{
+				"name": "raw_profile",
+				"content_type": "application/x-perfetto",
+				"stored_id": "aef123345"
+			}
+		]
+	}`
+
+	var c Chunk
+	if err := json.Unmarshal([]byte(payload), &c); err != nil {
+		t.Fatal(err)
+	}
+	sc, ok := c.Chunk().(*SampleChunk)
+	if !ok {
+		t.Fatalf("expected *SampleChunk, got %T", c.Chunk())
+	}
+	want := []Attachment{
+		{Name: "raw_profile", ContentType: "application/x-perfetto", StoredID: "aef123345"},
+	}
+	if diff := testutil.Diff(sc.Attachments, want); diff != "" {
+		t.Fatalf("Result mismatch: got - want +\n%s", diff)
+	}
+	if diff := testutil.Diff(c.GetAttachments(), want); diff != "" {
+		t.Fatalf("Result mismatch: got - want +\n%s", diff)
+	}
+
+	b, err := json.Marshal(c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Contains(b, []byte(`"attachments":[{"name":"raw_profile","content_type":"application/x-perfetto","stored_id":"aef123345"}]`)) {
+		t.Errorf("expected serialized chunk to contain the attachments: %s", b)
+	}
+}
+
+func TestAttachmentsOmittedWhenEmpty(t *testing.T) {
+	payload := `{"version": "2", "chunk_id": "0432a0a4c25f4697bf9f0a2fcbe6a814"}`
+
+	var c Chunk
+	if err := json.Unmarshal([]byte(payload), &c); err != nil {
+		t.Fatal(err)
+	}
+	b, err := json.Marshal(c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if bytes.Contains(b, []byte(`"attachments"`)) {
+		t.Errorf("expected attachments to be omitted: %s", b)
+	}
+}
+
+// A chunk with platform=android but a version set uses the sample v2 format
+// and must not be treated as a legacy android chunk.
+func TestUnmarshalAndroidPlatformWithVersionAsSampleChunk(t *testing.T) {
+	payload := `{
+		"version": "2",
+		"platform": "android",
+		"chunk_id": "0432a0a4c25f4697bf9f0a2fcbe6a814"
+	}`
+
+	var c Chunk
+	if err := json.Unmarshal([]byte(payload), &c); err != nil {
+		t.Fatal(err)
+	}
+	sc, ok := c.Chunk().(*SampleChunk)
+	if !ok {
+		t.Fatalf("expected *SampleChunk, got %T", c.Chunk())
+	}
+	if sc.Platform != platform.Android {
+		t.Errorf("expected platform %q, got %q", platform.Android, sc.Platform)
+	}
+}
+
+// Attachments are only supported for sample chunks:
+// the field is dropped on android chunks.
+func TestAttachmentsIgnoredOnAndroidChunks(t *testing.T) {
+	payload := `{
+		"chunk_id": "0432a0a4c25f4697bf9f0a2fcbe6a814",
+		"attachments": [
+			{
+				"name": "raw_profile",
+				"content_type": "application/x-perfetto",
+				"stored_id": "aef123345"
+			}
+		]
+	}`
+
+	var c Chunk
+	if err := json.Unmarshal([]byte(payload), &c); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := c.Chunk().(*AndroidChunk); !ok {
+		t.Fatalf("expected *AndroidChunk, got %T", c.Chunk())
+	}
+	if got := c.GetAttachments(); len(got) != 0 {
+		t.Errorf("expected GetAttachments to return an empty list, got %v", got)
+	}
+
+	b, err := json.Marshal(c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if bytes.Contains(b, []byte(`"attachments"`)) {
+		t.Errorf("expected attachments to be omitted: %s", b)
+	}
+}

internal/chunk/sample.go

@@ -46,6 +46,18 @@ type (
 		Measurements json.RawMessage `json:"measurements"`
 
 		Options options.Options `json:"options,omitempty"`
+
+		// Attachments lists the files related to this chunk (e.g. a raw
+		// profile) stored in the object store. On merged chunks it contains
+		// the attachments of all the chunks that were merged
+		// (see MergeSampleChunks).
+		Attachments []Attachment `json:"attachments,omitempty"`
+	}
+
+	Attachment struct {
+		Name        string `json:"name"`
+		ContentType string `json:"content_type,omitempty"`
+		StoredID    string `json:"stored_id"`
 	}
 
 	SampleData struct {
@@ -226,6 +238,10 @@ func (c SampleChunk) GetProjectID() uint64 {
 	return c.ProjectID
 }
 
+func (c SampleChunk) GetAttachments() []Attachment {
+	return c.Attachments
+}
+
 func (c SampleChunk) GetReceived() float64 {
 	return c.Received
 }

internal/chunk/sample_utils.go

@@ -101,6 +101,19 @@ func MergeSampleChunks(chunks []SampleChunk, startTS, endTS uint64) (SampleChunk
 		chunk.Measurements = jsonRawMesaurement
 	}
 
+	// Collect the attachments of all the merged chunks.
+	// The merged chunk is based on the first one, so reset its list
+	// before appending to avoid duplicating the first chunk's attachments.
+	chunk.Attachments = nil
+	for _, c := range chunks {
+		for _, attachment := range c.Attachments {
+			if attachment.StoredID == "" {
+				continue
+			}
+			chunk.Attachments = append(chunk.Attachments, attachment)
+		}
+	}
+
 	return chunk, nil
 }
 

internal/chunk/sample_utils_test.go

@@ -101,3 +101,75 @@ func TestMergeSampleChunks(t *testing.T) {
 		})
 	}
 }
+
+func TestMergeSampleChunksAttachments(t *testing.T) {
+	chunks := []SampleChunk{
+		{
+			Attachments: []Attachment{
+				{Name: "raw_profile", ContentType: "application/x-perfetto", StoredID: "raw_profile_b"},
+				// Attachments without a stored ID are skipped.
+				{Name: "raw_profile", ContentType: "application/x-perfetto", StoredID: ""},
+			},
+			Profile: SampleData{
+				Frames: []frame.Frame{
+					{Function: "b"},
+				},
+				Samples: []Sample{
+					{StackID: 0, Timestamp: 3.0},
+					{StackID: 0, Timestamp: 4.0},
+				},
+				Stacks: [][]int{
+					{0},
+				},
+			},
+		},
+		// chunk without attachments
+		{
+			Profile: SampleData{
+				Frames: []frame.Frame{
+					{Function: "c"},
+				},
+				Samples: []Sample{
+					{StackID: 0, Timestamp: 5.0},
+					{StackID: 0, Timestamp: 6.0},
+				},
+				Stacks: [][]int{
+					{0},
+				},
+			},
+		},
+		{
+			Attachments: []Attachment{
+				{Name: "raw_profile", ContentType: "application/x-perfetto", StoredID: "raw_profile_a"},
+			},
+			Profile: SampleData{
+				Frames: []frame.Frame{
+					{Function: "a"},
+				},
+				Samples: []Sample{
+					{StackID: 0, Timestamp: 1.0},
+					{StackID: 0, Timestamp: 2.0},
+				},
+				Stacks: [][]int{
+					{0},
+				},
+			},
+		},
+	}
+
+	merged, err := MergeSampleChunks(chunks, 0, uint64(10e9))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Attachments follow the merged chunk order. Note: the chunk sort in
+	// MergeSampleChunks is only well-defined for non-overlapping chunks,
+	// so no strict chronological order is guaranteed.
+	want := []Attachment{
+		{Name: "raw_profile", ContentType: "application/x-perfetto", StoredID: "raw_profile_a"},
+		{Name: "raw_profile", ContentType: "application/x-perfetto", StoredID: "raw_profile_b"},
+	}
+	if diff := testutil.Diff(merged.Attachments, want); diff != "" {
+		t.Fatalf("Result mismatch: got - want +\n%s", diff)
+	}
+}