Dependency inject the memory.Allocator for rendering (#3461)

pkg/parca/parca.go

@@ -374,6 +374,7 @@ func Run(ctx context.Context, logger log.Logger, reg *prometheus.Registry, flags
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 
 	ctx, cancel := context.WithCancel(ctx)

pkg/parca/parca_test.go

@@ -257,6 +257,7 @@ func TestConsistency(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 
 	ts := timestamppb.New(timestamp.Time(1608199718549)) // time_nanos of the profile divided by 1e6
@@ -366,6 +367,7 @@ func TestPGOE2e(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 
 	res, err := api.Query(ctx, &querypb.QueryRequest{
@@ -464,6 +466,7 @@ func TestLabels(t *testing.T) {
 			"labels",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 
 	ts := timestamppb.New(timestamp.Time(1677488315039)) // time_nanos of the profile divided by 1e6

pkg/query/columnquery.go

@@ -21,6 +21,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/apache/arrow/go/v13/arrow/memory"
 	"github.com/go-kit/log"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
@@ -54,20 +55,23 @@ type ColumnQueryAPI struct {
 	querier     Querier
 
 	tableConverterPool *sync.Pool
+	mem                memory.Allocator
 }
 
 func NewColumnQueryAPI(
 	logger log.Logger,
 	tracer trace.Tracer,
 	shareClient sharepb.ShareServiceClient,
 	querier Querier,
+	mem memory.Allocator,
 ) *ColumnQueryAPI {
 	return &ColumnQueryAPI{
 		logger:             logger,
 		tracer:             tracer,
 		shareClient:        shareClient,
 		querier:            querier,
 		tableConverterPool: NewTableConverterPool(),
+		mem:                mem,
 	}
 }
 
@@ -237,7 +241,7 @@ func (q *ColumnQueryAPI) renderReport(
 	nodeTrimThreshold float32,
 	filtered int64,
 ) (*pb.QueryResponse, error) {
-	return RenderReport(ctx, q.tracer, p, typ, nodeTrimThreshold, filtered, q.tableConverterPool)
+	return RenderReport(ctx, q.tracer, p, typ, nodeTrimThreshold, filtered, q.tableConverterPool, q.mem)
 }
 
 func RenderReport(
@@ -248,6 +252,7 @@ func RenderReport(
 	nodeTrimThreshold float32,
 	filtered int64,
 	pool *sync.Pool,
+	mem memory.Allocator,
 ) (*pb.QueryResponse, error) {
 	ctx, span := tracer.Start(ctx, "renderReport")
 	span.SetAttributes(attribute.String("reportType", typ.String()))
@@ -287,7 +292,7 @@ func RenderReport(
 		}, nil
 	case pb.QueryRequest_REPORT_TYPE_FLAMEGRAPH_ARROW:
 		// TODO: Make the fields to aggregate by configurable via the API.
-		fa, total, err := GenerateFlamegraphArrow(ctx, tracer, p, []string{FlamegraphFieldFunctionName}, nodeTrimFraction)
+		fa, total, err := GenerateFlamegraphArrow(ctx, mem, tracer, p, []string{FlamegraphFieldFunctionName}, nodeTrimFraction)
 		if err != nil {
 			return nil, status.Errorf(codes.Internal, "failed to generate arrow flamegraph: %v", err.Error())
 		}

pkg/query/columnquery_test.go

@@ -97,6 +97,7 @@ func TestColumnQueryAPIQueryRangeEmpty(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 	_, err = api.QueryRange(ctx, &pb.QueryRangeRequest{
 		Query: `memory:alloc_objects:count:space:bytes{job="default"}`,
@@ -219,6 +220,7 @@ func TestColumnQueryAPIQueryRange(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 	res, err := api.QueryRange(ctx, &pb.QueryRangeRequest{
 		Query: `memory:alloc_objects:count:space:bytes{job="default"}`,
@@ -313,6 +315,7 @@ func TestColumnQueryAPIQuerySingle(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 	ts := timestamppb.New(timestamp.Time(p.TimeNanos / time.Millisecond.Nanoseconds()))
 	res, err := api.Query(ctx, &pb.QueryRequest{
@@ -441,6 +444,7 @@ func TestColumnQueryAPIQueryFgprof(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 
 	res, err := api.QueryRange(ctx, &pb.QueryRangeRequest{
@@ -542,6 +546,7 @@ func TestColumnQueryAPIQueryCumulative(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 
 	// These have been extracted from the profiles above.
@@ -751,6 +756,7 @@ func TestColumnQueryAPIQueryDiff(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 
 	res, err := api.Query(ctx, &pb.QueryRequest{
@@ -932,6 +938,7 @@ func TestColumnQueryAPITypes(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 	res, err := api.ProfileTypes(ctx, &pb.ProfileTypesRequest{})
 	require.NoError(t, err)
@@ -1024,6 +1031,7 @@ func TestColumnQueryAPILabelNames(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 	res, err := api.Labels(ctx, &pb.LabelsRequest{})
 	require.NoError(t, err)
@@ -1108,6 +1116,7 @@ func TestColumnQueryAPILabelValues(t *testing.T) {
 			"stacktraces",
 			metastore,
 		),
+		memory.DefaultAllocator,
 	)
 	res, err := api.Values(ctx, &pb.ValuesRequest{
 		LabelName: "job",
@@ -1136,7 +1145,7 @@ func BenchmarkQuery(b *testing.B) {
 	b.ResetTimer()
 
 	for i := 0; i < b.N; i++ {
-		_, _ = RenderReport(ctx, tracer, sp, pb.QueryRequest_REPORT_TYPE_FLAMEGRAPH_ARROW, 0, 0, NewTableConverterPool())
+		_, _ = RenderReport(ctx, tracer, sp, pb.QueryRequest_REPORT_TYPE_FLAMEGRAPH_ARROW, 0, 0, NewTableConverterPool(), memory.DefaultAllocator)
 	}
 }
 

pkg/query/flamegraph_arrow.go

@@ -55,8 +55,7 @@ const (
 	FlamegraphFieldDiff       = "diff"
 )
 
-func GenerateFlamegraphArrow(ctx context.Context, tracer trace.Tracer, p *profile.Profile, aggregate []string, trimFraction float32) (*queryv1alpha1.FlamegraphArrow, int64, error) {
-	mem := memory.NewGoAllocator()
+func GenerateFlamegraphArrow(ctx context.Context, mem memory.Allocator, tracer trace.Tracer, p *profile.Profile, aggregate []string, trimFraction float32) (*queryv1alpha1.FlamegraphArrow, int64, error) {
 	record, cumulative, height, trimmed, err := generateFlamegraphArrowRecord(ctx, mem, tracer, p, aggregate, trimFraction)
 	if err != nil {
 		return nil, 0, err

pkg/query/query_test.go

@@ -118,6 +118,7 @@ func Benchmark_Query_Merge(b *testing.B) {
 					"stacktraces",
 					metastore,
 				),
+				memory.DefaultAllocator,
 			)
 			b.ResetTimer()
 
@@ -185,6 +186,7 @@ func Benchmark_ProfileTypes(b *testing.B) {
 			"stacktraces",
 			m,
 		),
+		memory.DefaultAllocator,
 	)
 	b.ResetTimer()