github.com/annwntech/go-micro/v2@v2.9.5/util/pool/default.go (about)

     1  package pool
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/google/uuid"
     8  	"github.com/annwntech/go-micro/v2/transport"
     9  )
    10  
    11  type pool struct {
    12  	size int
    13  	ttl  time.Duration
    14  	tr   transport.Transport
    15  
    16  	sync.Mutex
    17  	conns map[string][]*poolConn
    18  }
    19  
    20  type poolConn struct {
    21  	transport.Client
    22  	id      string
    23  	created time.Time
    24  }
    25  
    26  func newPool(options Options) *pool {
    27  	return &pool{
    28  		size:  options.Size,
    29  		tr:    options.Transport,
    30  		ttl:   options.TTL,
    31  		conns: make(map[string][]*poolConn),
    32  	}
    33  }
    34  
    35  func (p *pool) Close() error {
    36  	p.Lock()
    37  	for k, c := range p.conns {
    38  		for _, conn := range c {
    39  			conn.Client.Close()
    40  		}
    41  		delete(p.conns, k)
    42  	}
    43  	p.Unlock()
    44  	return nil
    45  }
    46  
    47  // NoOp the Close since we manage it
    48  func (p *poolConn) Close() error {
    49  	return nil
    50  }
    51  
    52  func (p *poolConn) Id() string {
    53  	return p.id
    54  }
    55  
    56  func (p *poolConn) Created() time.Time {
    57  	return p.created
    58  }
    59  
    60  func (p *pool) Get(addr string, opts ...transport.DialOption) (Conn, error) {
    61  	p.Lock()
    62  	conns := p.conns[addr]
    63  
    64  	// while we have conns check age and then return one
    65  	// otherwise we'll create a new conn
    66  	for len(conns) > 0 {
    67  		conn := conns[len(conns)-1]
    68  		conns = conns[:len(conns)-1]
    69  		p.conns[addr] = conns
    70  
    71  		// if conn is old kill it and move on
    72  		if d := time.Since(conn.Created()); d > p.ttl {
    73  			conn.Client.Close()
    74  			continue
    75  		}
    76  
    77  		// we got a good conn, lets unlock and return it
    78  		p.Unlock()
    79  
    80  		return conn, nil
    81  	}
    82  
    83  	p.Unlock()
    84  
    85  	// create new conn
    86  	c, err := p.tr.Dial(addr, opts...)
    87  	if err != nil {
    88  		return nil, err
    89  	}
    90  	return &poolConn{
    91  		Client:  c,
    92  		id:      uuid.New().String(),
    93  		created: time.Now(),
    94  	}, nil
    95  }
    96  
    97  func (p *pool) Release(conn Conn, err error) error {
    98  	// don't store the conn if it has errored
    99  	if err != nil {
   100  		return conn.(*poolConn).Client.Close()
   101  	}
   102  
   103  	// otherwise put it back for reuse
   104  	p.Lock()
   105  	conns := p.conns[conn.Remote()]
   106  	if len(conns) >= p.size {
   107  		p.Unlock()
   108  		return conn.(*poolConn).Client.Close()
   109  	}
   110  	p.conns[conn.Remote()] = append(conns, conn.(*poolConn))
   111  	p.Unlock()
   112  
   113  	return nil
   114  }