Add support for typed internal pprof transfers

Author: aleks-p

api/gen/proto/go/query/v1/query.pb.go

@@ -1662,6 +1662,14 @@
         }
       }
     },
+    "v1PprofOutputMode": {
+      "type": "string",
+      "enum": [
+        "PPROF_BYTES",
+        "PPROF_TYPED"
+      ],
+      "default": "PPROF_BYTES"
+    },
     "v1PprofQuery": {
       "type": "object",
       "properties": {
@@ -1670,7 +1678,10 @@
           "format": "int64"
         },
         "stackTraceSelector": {
-          "$ref": "#/definitions/v1StackTraceSelector",
+          "$ref": "#/definitions/v1StackTraceSelector"
+        },
+        "outputMode": {
+          "$ref": "#/definitions/v1PprofOutputMode",
           "description": "TODO(kolesnikovae): Go PGO options."
         }
       }
@@ -1684,6 +1695,9 @@
         "pprof": {
           "type": "string",
           "format": "byte"
+        },
+        "typedPprof": {
+          "$ref": "#/definitions/googlev1Profile"
         }
       }
     },

api/gen/proto/go/query/v1/query_vtproto.pb.go

@@ -182,9 +182,16 @@ message PprofQuery {
   int64 max_nodes = 1;
   optional types.v1.StackTraceSelector stack_trace_selector = 2;
   // TODO(kolesnikovae): Go PGO options.
+  PprofOutputMode output_mode = 3;
+}
+
+enum PprofOutputMode {
+  PPROF_BYTES = 0;
+  PPROF_TYPED = 1;
 }
 
 message PprofReport {
   PprofQuery query = 1;
   bytes pprof = 2;
+  google.v1.Profile typed_pprof = 3;
 }

api/openapiv2/gen/phlare.swagger.json

@@ -573,6 +573,8 @@ Usage of ./pyroscope:
     	Timeout for ingester client healthcheck RPCs. (default 5s)
   -querier.id string
     	Querier ID, sent to the query-frontend to identify requests from the same querier. Defaults to hostname.
+  -querier.internal-pprof-output-mode string
+    	The output mode for pprof queries, for transfers between internal components. Valid values are 'bytes' and 'typed'. (default "bytes")
   -querier.max-concurrent int
     	The maximum number of concurrent queries allowed. (default 4)
   -querier.max-flamegraph-nodes-default int

api/query/v1/query.proto

@@ -157,6 +157,8 @@ Usage of ./pyroscope:
     	Run a health check on each ingester client during periodic cleanup. (default true)
   -querier.health-check-timeout duration
     	Timeout for ingester client healthcheck RPCs. (default 5s)
+  -querier.internal-pprof-output-mode string
+    	The output mode for pprof queries, for transfers between internal components. Valid values are 'bytes' and 'typed'. (default "bytes")
   -querier.max-flamegraph-nodes-default int
     	Maximum number of flame graph nodes by default. 0 to disable. (default 8192)
   -querier.max-flamegraph-nodes-max int

cmd/pyroscope/help-all.txt.tmpl

@@ -2316,6 +2316,11 @@ distributor_usage_groups:
 # CLI flag: -querier.sanitize-on-merge
 [query_sanitize_on_merge: <boolean> | default = true]
 
+# The output mode for pprof queries, for transfers between internal components.
+# Valid values are 'bytes' and 'typed'.
+# CLI flag: -querier.internal-pprof-output-mode
+[query_internal_pprof_output_mode: <string> | default = "bytes"]
+
 # Delete blocks containing samples older than the specified retention period. 0
 # to disable.
 # CLI flag: -compactor.blocks-retention-period

cmd/pyroscope/help.txt.tmpl

@@ -114,6 +114,7 @@ type Limits interface {
 	QueryAnalysisEnabled(string) bool
 	SymbolizerEnabled(string) bool
 	QuerySanitizeOnMerge(string) bool
+	QueryInternalPprofOutputMode(string) string
 	validation.FlameGraphLimits
 }
 

docs/sources/configure-server/reference-configuration-parameters/index.md

@@ -27,6 +27,10 @@ func (m *mockLimits) QuerySanitizeOnMerge(_ string) bool {
 	return true
 }
 
+func (m *mockLimits) QueryInternalPprofOutputMode(_ string) string {
+	return "bytes"
+}
+
 func (m *mockLimits) MaxQueryParallelism(_ string) int {
 	return 100
 }

pkg/frontend/frontend.go

@@ -108,7 +108,8 @@ func (q *QueryFrontend) Query(
 		if shouldSymbolize && originalQuery.QueryType == queryv1.QueryType_QUERY_TREE {
 			modifiedQueries[i].QueryType = queryv1.QueryType_QUERY_PPROF
 			modifiedQueries[i].Pprof = &queryv1.PprofQuery{
-				MaxNodes: 0,
+				MaxNodes:   0,
+				OutputMode: queryv1.PprofOutputMode_PPROF_BYTES,
 			}
 			modifiedQueries[i].Tree = nil
 		}
@@ -239,12 +240,17 @@ func (q *QueryFrontend) processAndSymbolizeProfiles(
 		}
 
 		var prof googlev1.Profile
-		if err := pprof.Unmarshal(r.Pprof.Pprof, &prof); err != nil {
-			return fmt.Errorf("failed to unmarshal profile: %w", err)
-		}
-
-		if err := q.symbolizer.SymbolizePprof(ctx, &prof); err != nil {
-			return fmt.Errorf("failed to symbolize profile: %w", err)
+		if r.Pprof.TypedPprof != nil {
+			if err := q.symbolizer.SymbolizePprof(ctx, r.Pprof.TypedPprof); err != nil {
+				return fmt.Errorf("failed to symbolize profile: %w", err)
+			}
+		} else {
+			if err := pprof.Unmarshal(r.Pprof.Pprof, &prof); err != nil {
+				return fmt.Errorf("failed to unmarshal profile: %w", err)
+			}
+			if err := q.symbolizer.SymbolizePprof(ctx, &prof); err != nil {
+				return fmt.Errorf("failed to symbolize profile: %w", err)
+			}
 		}
 
 		// Convert back to tree if originally a tree
@@ -256,7 +262,7 @@ func (q *QueryFrontend) processAndSymbolizeProfiles(
 			r.Tree = &queryv1.TreeReport{Tree: treeBytes}
 			r.ReportType = queryv1.ReportType_REPORT_TREE
 			r.Pprof = nil
-		} else {
+		} else if len(prof.Sample) > 0 {
 			symbolizedBytes, err := pprof.Marshal(&prof, true)
 			if err != nil {
 				return fmt.Errorf("failed to marshal symbolized profile: %w", err)