gitee.com/liuxuezhan/go-micro-v1.18.0@v1.0.0/client/grpc/grpc_pool.go (about)

     1  package grpc
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"google.golang.org/grpc"
     8  )
     9  
    10  type pool struct {
    11  	size int
    12  	ttl  int64
    13  
    14  	sync.Mutex
    15  	conns map[string][]*poolConn
    16  }
    17  
    18  type poolConn struct {
    19  	*grpc.ClientConn
    20  	created int64
    21  }
    22  
    23  func newPool(size int, ttl time.Duration) *pool {
    24  	return &pool{
    25  		size:  size,
    26  		ttl:   int64(ttl.Seconds()),
    27  		conns: make(map[string][]*poolConn),
    28  	}
    29  }
    30  
    31  func (p *pool) getConn(addr string, opts ...grpc.DialOption) (*poolConn, error) {
    32  	p.Lock()
    33  	conns := p.conns[addr]
    34  	now := time.Now().Unix()
    35  
    36  	// while we have conns check age and then return one
    37  	// otherwise we'll create a new conn
    38  	for len(conns) > 0 {
    39  		conn := conns[len(conns)-1]
    40  		conns = conns[:len(conns)-1]
    41  		p.conns[addr] = conns
    42  
    43  		// if conn is old kill it and move on
    44  		if d := now - conn.created; d > p.ttl {
    45  			conn.ClientConn.Close()
    46  			continue
    47  		}
    48  
    49  		// we got a good conn, lets unlock and return it
    50  		p.Unlock()
    51  
    52  		return conn, nil
    53  	}
    54  
    55  	p.Unlock()
    56  
    57  	// create new conn
    58  	cc, err := grpc.Dial(addr, opts...)
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	return &poolConn{cc, time.Now().Unix()}, nil
    64  }
    65  
    66  func (p *pool) release(addr string, conn *poolConn, err error) {
    67  	// don't store the conn if it has errored
    68  	if err != nil {
    69  		conn.ClientConn.Close()
    70  		return
    71  	}
    72  
    73  	// otherwise put it back for reuse
    74  	p.Lock()
    75  	conns := p.conns[addr]
    76  	if len(conns) >= p.size {
    77  		p.Unlock()
    78  		conn.ClientConn.Close()
    79  		return
    80  	}
    81  	p.conns[addr] = append(conns, conn)
    82  	p.Unlock()
    83  }