github.com/grafana/pyroscope@v1.18.0/pkg/querybackend/backend.go (about)

     1  package querybackend
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/go-kit/log"
    10  	"github.com/grafana/dskit/grpcclient"
    11  	"github.com/grafana/dskit/services"
    12  	"github.com/opentracing/opentracing-go"
    13  	"github.com/prometheus/client_golang/prometheus"
    14  	"golang.org/x/sync/errgroup"
    15  
    16  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    17  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    18  	"github.com/grafana/pyroscope/pkg/util"
    19  )
    20  
    21  type Config struct {
    22  	Address          string            `yaml:"address"`
    23  	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config" doc:"description=Configures the gRPC client used to communicate between the query-frontends and the query-schedulers."`
    24  	ClientTimeout    time.Duration     `yaml:"client_timeout"`
    25  }
    26  
    27  func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
    28  	f.StringVar(&cfg.Address, "query-backend.address", "localhost:9095", "")
    29  	f.DurationVar(&cfg.ClientTimeout, "query-backend.client-timeout", 30*time.Second, "Timeout for query-backend client requests.")
    30  	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-backend.grpc-client-config", f)
    31  }
    32  
    33  func (cfg *Config) Validate() error {
    34  	if cfg.Address == "" {
    35  		return fmt.Errorf("query-backend.address is required")
    36  	}
    37  	return cfg.GRPCClientConfig.Validate()
    38  }
    39  
    40  type QueryHandler interface {
    41  	Invoke(context.Context, *queryv1.InvokeRequest) (*queryv1.InvokeResponse, error)
    42  }
    43  
    44  type QueryBackend struct {
    45  	service services.Service
    46  	queryv1.QueryBackendServiceServer
    47  
    48  	config Config
    49  	logger log.Logger
    50  	reg    prometheus.Registerer
    51  
    52  	backendClient QueryHandler
    53  	blockReader   QueryHandler
    54  }
    55  
    56  func New(
    57  	config Config,
    58  	logger log.Logger,
    59  	reg prometheus.Registerer,
    60  	backendClient QueryHandler,
    61  	blockReader QueryHandler,
    62  ) (*QueryBackend, error) {
    63  	q := QueryBackend{
    64  		config:        config,
    65  		logger:        logger,
    66  		reg:           reg,
    67  		backendClient: backendClient,
    68  		blockReader:   blockReader,
    69  	}
    70  	q.service = services.NewIdleService(q.starting, q.stopping)
    71  	return &q, nil
    72  }
    73  
    74  func (q *QueryBackend) Service() services.Service      { return q.service }
    75  func (q *QueryBackend) starting(context.Context) error { return nil }
    76  func (q *QueryBackend) stopping(error) error           { return nil }
    77  
    78  func (q *QueryBackend) Invoke(
    79  	ctx context.Context,
    80  	req *queryv1.InvokeRequest,
    81  ) (*queryv1.InvokeResponse, error) {
    82  	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryBackend.Invoke")
    83  	defer span.Finish()
    84  
    85  	switch r := req.QueryPlan.Root; r.Type {
    86  	case queryv1.QueryNode_MERGE:
    87  		return q.merge(ctx, req, r.Children)
    88  	case queryv1.QueryNode_READ:
    89  		return q.read(ctx, req, r.Blocks)
    90  	default:
    91  		panic("query plan: unknown node type")
    92  	}
    93  }
    94  
    95  func (q *QueryBackend) merge(
    96  	ctx context.Context,
    97  	request *queryv1.InvokeRequest,
    98  	children []*queryv1.QueryNode,
    99  ) (*queryv1.InvokeResponse, error) {
   100  	request.QueryPlan = nil
   101  	m := newAggregator(request)
   102  	g, ctx := errgroup.WithContext(ctx)
   103  	for _, child := range children {
   104  		req := request.CloneVT()
   105  		req.QueryPlan = &queryv1.QueryPlan{
   106  			Root: child,
   107  		}
   108  		g.Go(util.RecoverPanic(func() error {
   109  			// TODO: Speculative retry.
   110  			return m.aggregateResponse(q.backendClient.Invoke(ctx, req))
   111  		}))
   112  	}
   113  	if err := g.Wait(); err != nil {
   114  		return nil, err
   115  	}
   116  	return m.response()
   117  }
   118  
   119  func (q *QueryBackend) read(
   120  	ctx context.Context,
   121  	request *queryv1.InvokeRequest,
   122  	blocks []*metastorev1.BlockMeta,
   123  ) (*queryv1.InvokeResponse, error) {
   124  	request.QueryPlan = &queryv1.QueryPlan{
   125  		Root: &queryv1.QueryNode{
   126  			Blocks: blocks,
   127  		},
   128  	}
   129  	return q.blockReader.Invoke(ctx, request)
   130  }