github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/lightning/common/conn.go (about)

     1  // Copyright 2021 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package common
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  
    20  	"github.com/pingcap/errors"
    21  	"go.uber.org/zap"
    22  	"google.golang.org/grpc"
    23  
    24  	"github.com/pingcap/br/pkg/lightning/log"
    25  )
    26  
    27  // connPool is a lazy pool of gRPC channels.
    28  // When `Get` called, it lazily allocates new connection if connection not full.
    29  // If it's full, then it will return allocated channels round-robin.
    30  type ConnPool struct {
    31  	mu sync.Mutex
    32  
    33  	conns   []*grpc.ClientConn
    34  	next    int
    35  	cap     int
    36  	newConn func(ctx context.Context) (*grpc.ClientConn, error)
    37  }
    38  
    39  func (p *ConnPool) TakeConns() (conns []*grpc.ClientConn) {
    40  	p.mu.Lock()
    41  	defer p.mu.Unlock()
    42  	p.conns, conns = nil, p.conns
    43  	p.next = 0
    44  	return conns
    45  }
    46  
    47  // Close closes the conn pool.
    48  func (p *ConnPool) Close() {
    49  	for _, c := range p.TakeConns() {
    50  		if err := c.Close(); err != nil {
    51  			log.L().Warn("failed to close clientConn", zap.String("target", c.Target()), log.ShortError(err))
    52  		}
    53  	}
    54  }
    55  
    56  // get tries to get an existing connection from the pool, or make a new one if the pool not full.
    57  func (p *ConnPool) get(ctx context.Context) (*grpc.ClientConn, error) {
    58  	p.mu.Lock()
    59  	defer p.mu.Unlock()
    60  	if len(p.conns) < p.cap {
    61  		c, err := p.newConn(ctx)
    62  		if err != nil {
    63  			return nil, errors.Trace(err)
    64  		}
    65  		p.conns = append(p.conns, c)
    66  		return c, nil
    67  	}
    68  
    69  	conn := p.conns[p.next]
    70  	p.next = (p.next + 1) % p.cap
    71  	return conn, nil
    72  }
    73  
    74  // newConnPool creates a new connPool by the specified conn factory function and capacity.
    75  func NewConnPool(cap int, newConn func(ctx context.Context) (*grpc.ClientConn, error)) *ConnPool {
    76  	return &ConnPool{
    77  		cap:     cap,
    78  		conns:   make([]*grpc.ClientConn, 0, cap),
    79  		newConn: newConn,
    80  
    81  		mu: sync.Mutex{},
    82  	}
    83  }
    84  
    85  type GRPCConns struct {
    86  	mu    sync.Mutex
    87  	conns map[uint64]*ConnPool
    88  }
    89  
    90  func (conns *GRPCConns) Close() {
    91  	conns.mu.Lock()
    92  	defer conns.mu.Unlock()
    93  
    94  	for _, cp := range conns.conns {
    95  		cp.Close()
    96  	}
    97  }
    98  
    99  func (conns *GRPCConns) GetGrpcConn(ctx context.Context, storeID uint64, tcpConcurrency int, newConn func(ctx context.Context) (*grpc.ClientConn, error)) (*grpc.ClientConn, error) {
   100  	conns.mu.Lock()
   101  	defer conns.mu.Unlock()
   102  	if _, ok := conns.conns[storeID]; !ok {
   103  		conns.conns[storeID] = NewConnPool(tcpConcurrency, newConn)
   104  	}
   105  	return conns.conns[storeID].get(ctx)
   106  }
   107  
   108  func NewGRPCConns() GRPCConns {
   109  	cons := GRPCConns{conns: make(map[uint64]*ConnPool)}
   110  	return cons
   111  }