refactor code to use chunk interface

cmd/vroom/chunk.go

@@ -261,7 +261,7 @@ func (env *environment) postProfileFromChunkIDs(w http.ResponseWriter, r *http.R
 		}
 	}
 
-	chunks := make([]chunk.SampleChunk, 0, len(requestBody.ChunkIDs))
+	chunks := make([]chunk.Chunk, 0, len(requestBody.ChunkIDs))
 	// read the output of each tasks
 	for i := 0; i < len(requestBody.ChunkIDs); i++ {
 		res := <-results
@@ -305,26 +305,80 @@ func (env *environment) postProfileFromChunkIDs(w http.ResponseWriter, r *http.R
 
 	s = sentry.StartSpan(ctx, "chunks.merge")
 	s.Description = "Merge profile chunks into a single one"
-	chunk, err := chunk.MergeSampleChunks(chunks, requestBody.Start, requestBody.End)
-	s.Finish()
-	if err != nil {
-		hub.CaptureException(err)
-		w.WriteHeader(http.StatusInternalServerError)
+	if len(chunks) == 0 {
+		w.WriteHeader(http.StatusOK)
 		return
 	}
+	var resp []byte
+	// Here we check what type of chunks we're dealing with,
+	// since Android chunks and Sample chunks return completely
+	// different types (Chunk vs Speedscope), hence we can't hide
+	// the implementation behind an interface.
+	//
+	// We check the first chunk type, and use that to assert the
+	// type of all the elements in the slice and then call the
+	// appropriate utility.
+	switch chunks[0].Chunk().(type) {
+	case *chunk.SampleChunk:
+		sampleChunks := make([]chunk.SampleChunk, 0, len(chunks))
+		for _, c := range chunks {
+			sc, ok := c.Chunk().(*chunk.SampleChunk)
+			if !ok {
+				w.WriteHeader(http.StatusBadRequest)
+				return
+			}
+			sampleChunks = append(sampleChunks, *sc)
+		}
+		mergedChunk, err := chunk.MergeSampleChunks(sampleChunks, requestBody.Start, requestBody.End)
+		s.Finish()
+		if err != nil {
+			hub.CaptureException(err)
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+		s = sentry.StartSpan(ctx, "json.marshal")
+		resp, err = json.Marshal(postProfileFromChunkIDsResponse{Chunk: mergedChunk})
+		s.Finish()
+		if err != nil {
+			hub.CaptureException(err)
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
 
-	s = sentry.StartSpan(ctx, "json.marshal")
-	defer s.Finish()
-	b, err := json.Marshal(postProfileFromChunkIDsResponse{Chunk: chunk})
-	if err != nil {
-		hub.CaptureException(err)
-		w.WriteHeader(http.StatusInternalServerError)
+	case *chunk.AndroidChunk:
+		androidChunks := make([]chunk.AndroidChunk, 0, len(chunks))
+		for _, c := range chunks {
+			ac, ok := c.Chunk().(*chunk.AndroidChunk)
+			if !ok {
+				w.WriteHeader(http.StatusBadRequest)
+				return
+			}
+			androidChunks = append(androidChunks, *ac)
+		}
+		sp, err := chunk.SpeedscopeFromAndroidChunks(androidChunks, requestBody.Start, requestBody.End)
+		s.Finish()
+		if err != nil {
+			hub.CaptureException(err)
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+		s = sentry.StartSpan(ctx, "json.marshal")
+		resp, err = json.Marshal(sp)
+		s.Finish()
+		if err != nil {
+			hub.CaptureException(err)
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+	default:
+		// Should never happen.
+		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
 
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(http.StatusOK)
-	_, _ = w.Write(b)
+	_, _ = w.Write(resp)
 }
 
 type (

cmd/vroom/chunk_test.go

@@ -60,6 +60,7 @@ func TestPostAndReadSampleChunk(t *testing.T) {
 		Release:        "1.2",
 		OrganizationID: 1,
 		ProjectID:      1,
+		Version:        "2",
 		Profile: chunk.SampleData{
 			Frames: []frame.Frame{
 				{
@@ -125,7 +126,7 @@ func TestPostAndReadSampleChunk(t *testing.T) {
 
 			// read the chunk with UnmarshalCompressed and make sure that we can unmarshal
 			// the data into the Chunk struct and that it matches the original
-			var c chunk.SampleChunk
+			var c chunk.Chunk
 			err = storageutil.UnmarshalCompressed(
 				context.Background(),
 				test.blobBucket,
@@ -135,7 +136,7 @@ func TestPostAndReadSampleChunk(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
-			if diff := testutil.Diff(chunkData, c); diff != "" {
+			if diff := testutil.Diff(&chunkData, c.Chunk()); diff != "" {
 				t.Fatalf("Result mismatch: got - want +\n%s", diff)
 			}
 		})

internal/chunk/chunk.go

@@ -69,6 +69,10 @@ func (c Chunk) MarshalJSON() ([]byte, error) {
 	return json.Marshal(c.chunk)
 }
 
+func (c Chunk) Chunk() chunkInterface {
+	return c.chunk
+}
+
 func StoragePath(OrganizationID uint64, ProjectID uint64, ProfilerID string, ID string) string {
 	return fmt.Sprintf(
 		"%d/%d/%s/%s",

internal/chunk/sample_readjob.go

@@ -25,7 +25,7 @@ type (
 
 	ReadJobResult struct {
 		Err           error
-		Chunk         SampleChunk
+		Chunk         Chunk
 		TransactionID string
 		ThreadID      *string
 		Start         uint64
@@ -34,7 +34,7 @@ type (
 )
 
 func (job ReadJob) Read() {
-	var chunk SampleChunk
+	var chunk Chunk
 
 	err := storageutil.UnmarshalCompressed(
 		job.Ctx,

internal/flamegraph/flamegraph.go

@@ -426,7 +426,7 @@ func GetFlamegraphFromChunks(
 			}
 			continue
 		}
-		cm := chunkIDToMetadata[result.Chunk.ID]
+		cm := chunkIDToMetadata[result.Chunk.GetID()]
 		for _, interval := range cm.SpanIntervals {
 			callTrees, err := result.Chunk.CallTrees(&interval.ActiveThreadID)
 			if err != nil {
@@ -439,9 +439,9 @@ func GetFlamegraphFromChunks(
 
 			annotate := annotateWithProfileExample(
 				utils.NewExampleFromProfilerChunk(
-					result.Chunk.ProjectID,
-					result.Chunk.ProfilerID,
-					result.Chunk.ID,
+					result.Chunk.GetProjectID(),
+					result.Chunk.GetProfilerID(),
+					result.Chunk.GetID(),
 					result.TransactionID,
 					result.ThreadID,
 					result.Start,

internal/metrics/metrics.go

@@ -261,9 +261,9 @@ func (ma *Aggregator) GetMetricsFromCandidates(
 			}
 
 			resultMetadata = utils.NewExampleFromProfilerChunk(
-				result.Chunk.ProjectID,
-				result.Chunk.ProfilerID,
-				result.Chunk.ID,
+				result.Chunk.GetProjectID(),
+				result.Chunk.GetProfilerID(),
+				result.Chunk.GetID(),
 				result.TransactionID,
 				result.ThreadID,
 				result.Start,