github.com/grafana/pyroscope@v1.18.0/pkg/clientpool/ingester_client_pool.go (about)

     1  package clientpool
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"io"
     7  	"time"
     8  
     9  	"connectrpc.com/connect"
    10  	"github.com/go-kit/log"
    11  	"github.com/grafana/dskit/ring"
    12  	ring_client "github.com/grafana/dskit/ring/client"
    13  	"github.com/prometheus/client_golang/prometheus"
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/credentials/insecure"
    16  	"google.golang.org/grpc/health/grpc_health_v1"
    17  
    18  	"github.com/grafana/pyroscope/api/gen/proto/go/ingester/v1/ingesterv1connect"
    19  	"github.com/grafana/pyroscope/pkg/util"
    20  )
    21  
    22  // PoolConfig is config for creating a Pool.
    23  type PoolConfig struct {
    24  	ClientCleanupPeriod  time.Duration `yaml:"client_cleanup_period"`
    25  	HealthCheckIngesters bool          `yaml:"health_check_ingesters"`
    26  	RemoteTimeout        time.Duration `yaml:"remote_timeout"`
    27  }
    28  
    29  // RegisterFlagsWithPrefix adds the flags required to config this to the given FlagSet.
    30  func (cfg *PoolConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
    31  	f.DurationVar(&cfg.ClientCleanupPeriod, prefix+".client-cleanup-period", 15*time.Second, "How frequently to clean up clients for ingesters that have gone away.")
    32  	f.BoolVar(&cfg.HealthCheckIngesters, prefix+".health-check-ingesters", true, "Run a health check on each ingester client during periodic cleanup.")
    33  	f.DurationVar(&cfg.RemoteTimeout, prefix+".health-check-timeout", 5*time.Second, "Timeout for ingester client healthcheck RPCs.")
    34  }
    35  
    36  func NewIngesterPool(cfg PoolConfig, ring ring.ReadRing, factory ring_client.PoolFactory, clientsMetric prometheus.Gauge, logger log.Logger, options ...connect.ClientOption) *ring_client.Pool {
    37  	if factory == nil {
    38  		factory = newIngesterPoolFactory(options...)
    39  	}
    40  	poolCfg := ring_client.PoolConfig{
    41  		CheckInterval:      cfg.ClientCleanupPeriod,
    42  		HealthCheckEnabled: cfg.HealthCheckIngesters,
    43  		HealthCheckTimeout: cfg.RemoteTimeout,
    44  	}
    45  
    46  	return ring_client.NewPool("ingester", poolCfg, ring_client.NewRingServiceDiscovery(ring), factory, clientsMetric, logger)
    47  }
    48  
    49  type ingesterPoolFactory struct {
    50  	options []connect.ClientOption
    51  }
    52  
    53  func newIngesterPoolFactory(options ...connect.ClientOption) ring_client.PoolFactory {
    54  	return &ingesterPoolFactory{options: options}
    55  }
    56  
    57  func (f *ingesterPoolFactory) FromInstance(inst ring.InstanceDesc) (ring_client.PoolClient, error) {
    58  	conn, err := grpc.Dial(inst.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport(), util.WithBaggageTransport())
    64  	return &ingesterPoolClient{
    65  		IngesterServiceClient: ingesterv1connect.NewIngesterServiceClient(httpClient, "http://"+inst.Addr, f.options...),
    66  		HealthClient:          grpc_health_v1.NewHealthClient(conn),
    67  		Closer:                conn,
    68  	}, nil
    69  }
    70  
    71  type ingesterPoolClient struct {
    72  	ingesterv1connect.IngesterServiceClient
    73  	grpc_health_v1.HealthClient
    74  	io.Closer
    75  }
    76  
    77  func (c *ingesterPoolClient) MergeProfilesStacktraces(ctx context.Context) BidiClientMergeProfilesStacktraces {
    78  	return c.IngesterServiceClient.MergeProfilesStacktraces(ctx)
    79  }
    80  
    81  func (c *ingesterPoolClient) MergeProfilesLabels(ctx context.Context) BidiClientMergeProfilesLabels {
    82  	return c.IngesterServiceClient.MergeProfilesLabels(ctx)
    83  }
    84  
    85  func (c *ingesterPoolClient) MergeProfilesPprof(ctx context.Context) BidiClientMergeProfilesPprof {
    86  	return c.IngesterServiceClient.MergeProfilesPprof(ctx)
    87  }
    88  
    89  func (c *ingesterPoolClient) MergeSpanProfile(ctx context.Context) BidiClientMergeSpanProfile {
    90  	return c.IngesterServiceClient.MergeSpanProfile(ctx)
    91  }