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

     1  package clientpool
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"time"
     7  
     8  	"connectrpc.com/connect"
     9  	"github.com/go-kit/log"
    10  	"github.com/grafana/dskit/ring"
    11  	ring_client "github.com/grafana/dskit/ring/client"
    12  	"github.com/prometheus/client_golang/prometheus"
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/credentials/insecure"
    15  	"google.golang.org/grpc/health/grpc_health_v1"
    16  
    17  	ingestv1 "github.com/grafana/pyroscope/api/gen/proto/go/ingester/v1"
    18  	"github.com/grafana/pyroscope/api/gen/proto/go/storegateway/v1/storegatewayv1connect"
    19  	"github.com/grafana/pyroscope/pkg/util"
    20  )
    21  
    22  type BidiClientMergeProfilesStacktraces interface {
    23  	Send(*ingestv1.MergeProfilesStacktracesRequest) error
    24  	Receive() (*ingestv1.MergeProfilesStacktracesResponse, error)
    25  	CloseRequest() error
    26  	CloseResponse() error
    27  }
    28  
    29  type BidiClientMergeProfilesLabels interface {
    30  	Send(*ingestv1.MergeProfilesLabelsRequest) error
    31  	Receive() (*ingestv1.MergeProfilesLabelsResponse, error)
    32  	CloseRequest() error
    33  	CloseResponse() error
    34  }
    35  
    36  type BidiClientMergeProfilesPprof interface {
    37  	Send(*ingestv1.MergeProfilesPprofRequest) error
    38  	Receive() (*ingestv1.MergeProfilesPprofResponse, error)
    39  	CloseRequest() error
    40  	CloseResponse() error
    41  }
    42  
    43  func NewPool(ring ring.ReadRing, factory ring_client.PoolFactory, clientsMetric prometheus.Gauge, logger log.Logger, options ...connect.ClientOption) *ring_client.Pool {
    44  	if factory == nil {
    45  		factory = newPoolFactory()
    46  	}
    47  	poolCfg := ring_client.PoolConfig{
    48  		CheckInterval:      10 * time.Second,
    49  		HealthCheckEnabled: true,
    50  		HealthCheckTimeout: 10 * time.Second,
    51  	}
    52  
    53  	return ring_client.NewPool("store-gateway", poolCfg, ring_client.NewRingServiceDiscovery(ring), factory, clientsMetric, logger)
    54  }
    55  
    56  type poolFactory struct {
    57  	options []connect.ClientOption
    58  }
    59  
    60  func newPoolFactory(options ...connect.ClientOption) ring_client.PoolFactory {
    61  	return &poolFactory{options: options}
    62  }
    63  
    64  func (f *poolFactory) FromInstance(inst ring.InstanceDesc) (ring_client.PoolClient, error) {
    65  	conn, err := grpc.Dial(inst.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
    66  	if err != nil {
    67  		return nil, err
    68  	}
    69  
    70  	httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport(), util.WithBaggageTransport())
    71  	return &storeGatewayPoolClient{
    72  		StoreGatewayServiceClient: storegatewayv1connect.NewStoreGatewayServiceClient(httpClient, "http://"+inst.Addr, f.options...),
    73  		HealthClient:              grpc_health_v1.NewHealthClient(conn),
    74  		Closer:                    conn,
    75  	}, nil
    76  }
    77  
    78  type storeGatewayPoolClient struct {
    79  	storegatewayv1connect.StoreGatewayServiceClient
    80  	grpc_health_v1.HealthClient
    81  	io.Closer
    82  }
    83  
    84  func (c *storeGatewayPoolClient) MergeProfilesStacktraces(ctx context.Context) BidiClientMergeProfilesStacktraces {
    85  	return c.StoreGatewayServiceClient.MergeProfilesStacktraces(ctx)
    86  }
    87  
    88  func (c *storeGatewayPoolClient) MergeProfilesLabels(ctx context.Context) BidiClientMergeProfilesLabels {
    89  	return c.StoreGatewayServiceClient.MergeProfilesLabels(ctx)
    90  }
    91  
    92  func (c *storeGatewayPoolClient) MergeProfilesPprof(ctx context.Context) BidiClientMergeProfilesPprof {
    93  	return c.StoreGatewayServiceClient.MergeProfilesPprof(ctx)
    94  }