github.com/pachyderm/pachyderm@v1.13.4/src/client/pkg/grpcutil/dialer.go (about)

     1  package grpcutil
     2  
     3  import (
     4  	"strings"
     5  	"sync"
     6  
     7  	"github.com/pachyderm/pachyderm/src/client/pkg/tracing"
     8  	"google.golang.org/grpc"
     9  )
    10  
    11  // Dialer defines a grpc.ClientConn connection dialer.
    12  type Dialer interface {
    13  	Dial(address string) (*grpc.ClientConn, error)
    14  	CloseConns() error
    15  }
    16  
    17  // NewDialer creates a Dialer.
    18  func NewDialer(opts ...grpc.DialOption) Dialer {
    19  	return newDialer(opts...)
    20  }
    21  
    22  type dialer struct {
    23  	opts []grpc.DialOption
    24  	// A map from addresses to connections
    25  	connMap map[string]*grpc.ClientConn
    26  	lock    sync.Mutex
    27  }
    28  
    29  func newDialer(opts ...grpc.DialOption) *dialer {
    30  	return &dialer{
    31  		opts:    opts,
    32  		connMap: make(map[string]*grpc.ClientConn),
    33  	}
    34  }
    35  
    36  func (d *dialer) Dial(addr string) (*grpc.ClientConn, error) {
    37  	d.lock.Lock()
    38  	defer d.lock.Unlock()
    39  	if conn, ok := d.connMap[addr]; ok {
    40  		return conn, nil
    41  	}
    42  	opts := append(d.opts,
    43  		grpc.WithUnaryInterceptor(tracing.UnaryClientInterceptor()),
    44  		grpc.WithStreamInterceptor(tracing.StreamClientInterceptor()),
    45  	)
    46  	daddr := addr
    47  	if !strings.HasPrefix(daddr, "dns:///") {
    48  		daddr = "dns:///" + daddr
    49  	}
    50  	conn, err := grpc.Dial(daddr, opts...)
    51  	if err != nil {
    52  		return nil, err
    53  	}
    54  	d.connMap[addr] = conn
    55  	return conn, nil
    56  }
    57  
    58  func (d *dialer) CloseConns() error {
    59  	d.lock.Lock()
    60  	defer d.lock.Unlock()
    61  	for addr, conn := range d.connMap {
    62  		if err := conn.Close(); err != nil {
    63  			return err
    64  		}
    65  		delete(d.connMap, addr)
    66  	}
    67  	return nil
    68  }