github.com/lingyao2333/mo-zero@v1.4.1/zrpc/proxy.go (about)

     1  package zrpc
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/syncx"
     8  	"github.com/lingyao2333/mo-zero/zrpc/internal"
     9  	"github.com/lingyao2333/mo-zero/zrpc/internal/auth"
    10  	"google.golang.org/grpc"
    11  )
    12  
    13  // A RpcProxy is a rpc proxy.
    14  type RpcProxy struct {
    15  	backend      string
    16  	clients      map[string]Client
    17  	options      []internal.ClientOption
    18  	singleFlight syncx.SingleFlight
    19  	lock         sync.Mutex
    20  }
    21  
    22  // NewProxy returns a RpcProxy.
    23  func NewProxy(backend string, opts ...internal.ClientOption) *RpcProxy {
    24  	return &RpcProxy{
    25  		backend:      backend,
    26  		clients:      make(map[string]Client),
    27  		options:      opts,
    28  		singleFlight: syncx.NewSingleFlight(),
    29  	}
    30  }
    31  
    32  // TakeConn returns a grpc.ClientConn.
    33  func (p *RpcProxy) TakeConn(ctx context.Context) (*grpc.ClientConn, error) {
    34  	cred := auth.ParseCredential(ctx)
    35  	key := cred.App + "/" + cred.Token
    36  	val, err := p.singleFlight.Do(key, func() (interface{}, error) {
    37  		p.lock.Lock()
    38  		client, ok := p.clients[key]
    39  		p.lock.Unlock()
    40  		if ok {
    41  			return client, nil
    42  		}
    43  
    44  		opts := append(p.options, WithDialOption(grpc.WithPerRPCCredentials(&auth.Credential{
    45  			App:   cred.App,
    46  			Token: cred.Token,
    47  		})))
    48  		client, err := NewClientWithTarget(p.backend, opts...)
    49  		if err != nil {
    50  			return nil, err
    51  		}
    52  
    53  		p.lock.Lock()
    54  		p.clients[key] = client
    55  		p.lock.Unlock()
    56  		return client, nil
    57  	})
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  
    62  	return val.(Client).Conn(), nil
    63  }