github.com/pachyderm/pachyderm@v1.13.4/src/client/pkg/shard/router.go (about)

     1  package shard
     2  
     3  import (
     4  	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
     5  	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
     6  	"google.golang.org/grpc"
     7  )
     8  
     9  type router struct {
    10  	sharder      Sharder
    11  	dialer       grpcutil.Dialer
    12  	localAddress string
    13  }
    14  
    15  func newRouter(
    16  	sharder Sharder,
    17  	dialer grpcutil.Dialer,
    18  	localAddress string,
    19  ) *router {
    20  	return &router{
    21  		sharder,
    22  		dialer,
    23  		localAddress,
    24  	}
    25  }
    26  
    27  func (r *router) GetShards(version int64) (map[uint64]bool, error) {
    28  	shardToAddress, err := r.sharder.GetShardToAddress(version)
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  	result := make(map[uint64]bool)
    33  	for shard, address := range shardToAddress {
    34  		if address == r.localAddress {
    35  			result[shard] = true
    36  		}
    37  	}
    38  	return result, nil
    39  }
    40  
    41  func (r *router) GetAllShards(version int64) (map[uint64]bool, error) {
    42  	shardToAddress, err := r.sharder.GetShardToAddress(version)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  	result := make(map[uint64]bool)
    47  	for shard, address := range shardToAddress {
    48  		if address == r.localAddress {
    49  			result[shard] = true
    50  		}
    51  	}
    52  	return result, nil
    53  }
    54  
    55  func (r *router) GetClientConn(shard uint64, version int64) (*grpc.ClientConn, error) {
    56  	address, ok, err := r.sharder.GetAddress(shard, version)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	if !ok {
    61  		return nil, errors.Errorf("no master found for %d", shard)
    62  	}
    63  	return r.dialer.Dial(address)
    64  }
    65  
    66  func (r *router) GetAllClientConns(version int64) ([]*grpc.ClientConn, error) {
    67  	addresses, err := r.getAllAddresses(version)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  	var clientConns []*grpc.ClientConn
    72  	for address := range addresses {
    73  		clientConn, err := r.dialer.Dial(address)
    74  		if err != nil {
    75  			return nil, err
    76  		}
    77  		clientConns = append(clientConns, clientConn)
    78  	}
    79  	return clientConns, nil
    80  }
    81  
    82  func (r *router) getAllAddresses(version int64) (map[string]bool, error) {
    83  	result := make(map[string]bool)
    84  	shardToAddress, err := r.sharder.GetShardToAddress(version)
    85  	if err != nil {
    86  		return nil, err
    87  	}
    88  	for _, address := range shardToAddress {
    89  		result[address] = true
    90  	}
    91  	return result, nil
    92  }
    93  
    94  func (r *router) CloseClientConns() error {
    95  	return r.dialer.CloseConns()
    96  }