github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/client/connpool/conn_pool_ring.go (about)

     1  package connpool
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/grafana/dskit/services"
     7  	"google.golang.org/grpc"
     8  	"google.golang.org/grpc/health/grpc_health_v1"
     9  
    10  	"github.com/grafana/dskit/ring"
    11  	ring_client "github.com/grafana/dskit/ring/client"
    12  
    13  	"github.com/grafana/pyroscope/pkg/util/health"
    14  )
    15  
    16  // NOTE(kolesnikovae): This is a tiny wrapper for ring_client.Pool
    17  // that is not tailored for the specific use case of the segment
    18  // writer client, and it should be refactored out.
    19  
    20  type ConnPool interface {
    21  	GetConnFor(addr string) (grpc.ClientConnInterface, error)
    22  	services.Service
    23  }
    24  
    25  type Pool struct{ *ring_client.Pool }
    26  
    27  func (p *Pool) GetConnFor(addr string) (grpc.ClientConnInterface, error) {
    28  	c, err := p.GetClientFor(addr)
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  	return c.(grpc.ClientConnInterface), nil
    33  }
    34  
    35  type ConnFactory struct {
    36  	options func(ring.InstanceDesc) []grpc.DialOption
    37  }
    38  
    39  func NewConnPoolFactory(options func(ring.InstanceDesc) []grpc.DialOption) ring_client.PoolFactory {
    40  	return &ConnFactory{options: options}
    41  }
    42  
    43  func (f *ConnFactory) FromInstance(inst ring.InstanceDesc) (ring_client.PoolClient, error) {
    44  	conn, err := grpc.NewClient(inst.Addr, f.options(inst)...)
    45  	if err != nil {
    46  		return nil, err
    47  	}
    48  	return &poolConn{
    49  		ClientConnInterface: conn,
    50  		HealthClient:        health.NoOpClient,
    51  		Closer:              conn,
    52  	}, nil
    53  }
    54  
    55  type poolConn struct {
    56  	grpc.ClientConnInterface
    57  	grpc_health_v1.HealthClient
    58  	io.Closer
    59  }