github.com/grafana/pyroscope@v1.18.0/pkg/clientpool/store_gateway_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  	"github.com/grafana/pyroscope/api/gen/proto/go/storegateway/v1/storegatewayv1connect"
    18  	"github.com/grafana/pyroscope/pkg/util"
    19  )
    20  
    21  func NewStoreGatewayPool(ring ring.ReadRing, factory ring_client.PoolFactory, clientsMetric prometheus.Gauge, logger log.Logger, options ...connect.ClientOption) *ring_client.Pool {
    22  	if factory == nil {
    23  		factory = newStoreGatewayPoolFactory(options...)
    24  	}
    25  	poolCfg := ring_client.PoolConfig{
    26  		CheckInterval:      10 * time.Second,
    27  		HealthCheckEnabled: true,
    28  		HealthCheckTimeout: 10 * time.Second,
    29  	}
    30  
    31  	return ring_client.NewPool("store-gateway", poolCfg, ring_client.NewRingServiceDiscovery(ring), factory, clientsMetric, logger)
    32  }
    33  
    34  type storeGatewayPoolFactory struct {
    35  	options []connect.ClientOption
    36  }
    37  
    38  func newStoreGatewayPoolFactory(options ...connect.ClientOption) ring_client.PoolFactory {
    39  	return &storeGatewayPoolFactory{options: options}
    40  }
    41  
    42  func (f *storeGatewayPoolFactory) FromInstance(inst ring.InstanceDesc) (ring_client.PoolClient, error) {
    43  	conn, err := grpc.Dial(inst.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	httpClient := util.InstrumentedDefaultHTTPClient(util.WithTracingTransport(), util.WithBaggageTransport())
    49  	return &storeGatewayPoolClient{
    50  		StoreGatewayServiceClient: storegatewayv1connect.NewStoreGatewayServiceClient(httpClient, "http://"+inst.Addr, f.options...),
    51  		HealthClient:              grpc_health_v1.NewHealthClient(conn),
    52  		Closer:                    conn,
    53  	}, nil
    54  
    55  }
    56  
    57  type storeGatewayPoolClient struct {
    58  	storegatewayv1connect.StoreGatewayServiceClient
    59  	grpc_health_v1.HealthClient
    60  	io.Closer
    61  }
    62  
    63  func (c *storeGatewayPoolClient) MergeProfilesStacktraces(ctx context.Context) BidiClientMergeProfilesStacktraces {
    64  	return c.StoreGatewayServiceClient.MergeProfilesStacktraces(ctx)
    65  }
    66  
    67  func (c *storeGatewayPoolClient) MergeProfilesLabels(ctx context.Context) BidiClientMergeProfilesLabels {
    68  	return c.StoreGatewayServiceClient.MergeProfilesLabels(ctx)
    69  }
    70  
    71  func (c *storeGatewayPoolClient) MergeProfilesPprof(ctx context.Context) BidiClientMergeProfilesPprof {
    72  	return c.StoreGatewayServiceClient.MergeProfilesPprof(ctx)
    73  }
    74  
    75  func (c *storeGatewayPoolClient) MergeSpanProfile(ctx context.Context) BidiClientMergeSpanProfile {
    76  	return c.StoreGatewayServiceClient.MergeSpanProfile(ctx)
    77  }