From e594bfd75e5cbf9483bbc84ad465c0937479b38d Mon Sep 17 00:00:00 2001
From: Matthias Loibl <mail@matthiasloibl.com>
Date: Thu, 31 Mar 2022 16:04:40 +0200
Subject: [PATCH] pkg/query: Add stripEmptyNodes helper for flamegraphs

This should get rid of all empty (meaning the cumulative value is zero) such that we don't even send metadata along for those nodes.
---
 pkg/query/query.go      | 17 +++++++++++++++++
 pkg/query/query_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/pkg/query/query.go b/pkg/query/query.go
index 6db90228574..bb4eb6271a3 100644
--- a/pkg/query/query.go
+++ b/pkg/query/query.go
@@ -294,6 +294,7 @@ func (q *Query) renderReport(ctx context.Context, p profile.InstantProfile, typ
 		if err != nil {
 			return nil, status.Errorf(codes.Internal, "failed to generate flamegraph: %v", err.Error())
 		}
+		fg = stripEmptyNodes(fg)
 		return &pb.QueryResponse{
 			Report: &pb.QueryResponse_Flamegraph{
 				Flamegraph: fg,
@@ -326,6 +327,22 @@ func (q *Query) renderReport(ctx context.Context, p profile.InstantProfile, typ
 	}
 }
 
+// stripEmptyNodes iterates over all root children and only adds points to the children to the
+// nodes that actually have a non-zero cumulative value to the returned root children.
+// It's more a fix/hack than anything.
+// In the best case we don't even create the empty nodes to begin with when querying the storage.
+func stripEmptyNodes(fg *pb.Flamegraph) *pb.Flamegraph {
+	// Iterate over all root children and remove the once that have cumulative value == 0
+	children := make([]*pb.FlamegraphNode, 0, len(fg.Root.Children))
+	for _, child := range fg.Root.Children {
+		if child.Cumulative != 0 {
+			children = append(children, child)
+		}
+	}
+	fg.Root.Children = children
+	return fg
+}
+
 func (q *Query) findSingle(ctx context.Context, sel []*labels.Matcher, t time.Time) (profile.InstantProfile, error) {
 	requestedTime := timestamp.FromTime(t)
 
diff --git a/pkg/query/query_test.go b/pkg/query/query_test.go
index 98f9da2a076..e08e5dc78bc 100644
--- a/pkg/query/query_test.go
+++ b/pkg/query/query_test.go
@@ -720,3 +720,33 @@ func Test_QueryRange_MultipleLabels_NoMatch(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 0, len(resp.GetSeries()))
 }
+
+func Test_stripEmptyNodes(t *testing.T) {
+	original := &pb.Flamegraph{
+		Root: &pb.FlamegraphRootNode{
+			Cumulative: 7,
+			Children: []*pb.FlamegraphNode{
+				{Cumulative: 0},
+				{Cumulative: 1},
+				{Cumulative: 0},
+				{Cumulative: 2},
+				{Cumulative: 0},
+				{Cumulative: 4},
+				{Cumulative: 0},
+				{Cumulative: 0},
+			},
+		},
+	}
+	expected := &pb.Flamegraph{
+		Root: &pb.FlamegraphRootNode{
+			Cumulative: 7,
+			Children: []*pb.FlamegraphNode{
+				{Cumulative: 1},
+				{Cumulative: 2},
+				{Cumulative: 4},
+			},
+		},
+	}
+
+	require.Equal(t, expected, stripEmptyNodes(original))
+}