fix: limit max nodes in symbolization queries (#4629)

* fix: limit max nodes in symbolization queries

Author: marcsanmi

pkg/frontend/readpath/queryfrontend/query_frontend.go

@@ -120,7 +120,7 @@ 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: originalQuery.Tree.GetMaxNodes(),
 			}
 			modifiedQueries[i].Tree = nil
 		}
@@ -299,7 +299,7 @@ func (q *QueryFrontend) processAndSymbolizeProfiles(
 
 		// Convert back to tree if originally a tree
 		if i < len(originalQueries) && originalQueries[i].QueryType == queryv1.QueryType_QUERY_TREE {
-			treeBytes, err := model.TreeFromBackendProfile(&prof, originalQueries[i].Tree.MaxNodes)
+			treeBytes, err := model.TreeFromBackendProfile(&prof, originalQueries[i].Tree.GetMaxNodes())
 			if err != nil {
 				return fmt.Errorf("failed to build tree: %w", err)
 			}

pkg/pyroscope/modules_experimental.go

@@ -369,6 +369,7 @@ func (f *Pyroscope) initQueryBackendClient() (services.Service, error) {
 	c, err := querybackendclient.New(
 		f.Cfg.QueryBackend.Address,
 		f.Cfg.QueryBackend.GRPCClientConfig,
+		f.Cfg.QueryBackend.ClientTimeout,
 	)
 	if err != nil {
 		return nil, err

pkg/querybackend/backend.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/grafana/dskit/grpcclient"
@@ -20,10 +21,12 @@ import (
 type Config struct {
 	Address          string            `yaml:"address"`
 	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config" doc:"description=Configures the gRPC client used to communicate between the query-frontends and the query-schedulers."`
+	ClientTimeout    time.Duration     `yaml:"client_timeout"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Address, "query-backend.address", "localhost:9095", "")
+	f.DurationVar(&cfg.ClientTimeout, "query-backend.client-timeout", 30*time.Second, "Timeout for query-backend client requests.")
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-backend.grpc-client-config", f)
 }
 

pkg/querybackend/client/client.go

@@ -2,6 +2,8 @@ package querybackendclient
 
 import (
 	"context"
+	"fmt"
+	"time"
 
 	"github.com/grafana/dskit/grpcclient"
 	"github.com/grafana/dskit/services"
@@ -15,8 +17,8 @@ type Client struct {
 	grpcClient queryv1.QueryBackendServiceClient
 }
 
-func New(address string, grpcClientConfig grpcclient.Config, dialOpts ...grpc.DialOption) (*Client, error) {
-	conn, err := dial(address, grpcClientConfig, dialOpts...)
+func New(address string, grpcClientConfig grpcclient.Config, timeout time.Duration, dialOpts ...grpc.DialOption) (*Client, error) {
+	conn, err := dial(address, grpcClientConfig, timeout, dialOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -26,14 +28,15 @@ func New(address string, grpcClientConfig grpcclient.Config, dialOpts ...grpc.Di
 	return &c, nil
 }
 
-func dial(address string, grpcClientConfig grpcclient.Config, dialOpts ...grpc.DialOption) (*grpc.ClientConn, error) {
+func dial(address string, grpcClientConfig grpcclient.Config, timeout time.Duration, dialOpts ...grpc.DialOption) (*grpc.ClientConn, error) {
 	options, err := grpcClientConfig.DialOption(nil, nil, nil)
 	if err != nil {
 		return nil, err
 	}
 	// TODO: https://github.com/grpc/grpc-proto/blob/master/grpc/service_config/service_config.proto
+	serviceConfig := fmt.Sprintf(grpcServiceConfigTemplate, timeout.Seconds())
 	options = append(options,
-		grpc.WithDefaultServiceConfig(grpcServiceConfig),
+		grpc.WithDefaultServiceConfig(serviceConfig),
 		grpc.WithMaxCallAttempts(500),
 	)
 	options = append(options, dialOpts...)
@@ -48,7 +51,7 @@ func (b *Client) Invoke(ctx context.Context, req *queryv1.InvokeRequest) (*query
 	return b.grpcClient.Invoke(ctx, req)
 }
 
-const grpcServiceConfig = `{
+const grpcServiceConfigTemplate = `{
     "loadBalancingPolicy":"round_robin",
     "methodConfig": [{
         "name": [{"service": ""}],
@@ -63,6 +66,6 @@ const grpcServiceConfig = `{
               "RESOURCE_EXHAUSTED"
             ]
         },
-        "timeout": "20s"
+        "timeout": "%.0fs"
     }]
 }`

pkg/querybackend/client/client_test.go

@@ -94,7 +94,7 @@ func Test_Concurrency(t *testing.T) {
 	resolver.Register(&multiResolverBuilder{targets: addresses})
 	backendAddress := "multi:///"
 
-	cl, err := New(backendAddress, grpcClientCfg, dialOpt)
+	cl, err := New(backendAddress, grpcClientCfg, 30*time.Second, dialOpt)
 	require.NoError(t, err)
 
 	for i := 0; i < nServers; i++ {