go-micro.dev/v5@v5.12.0/util/pool/default.go (about)

     1  package pool
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/google/uuid"
     9  
    10  	"go-micro.dev/v5/transport"
    11  )
    12  
    13  type pool struct {
    14  	tr transport.Transport
    15  
    16  	closeTimeout time.Duration
    17  	conns        map[string][]*poolConn
    18  	mu           sync.Mutex
    19  	size         int
    20  	ttl          time.Duration
    21  }
    22  
    23  type poolConn struct {
    24  	transport.Client
    25  
    26  	closeTimeout time.Duration
    27  	created      time.Time
    28  	id           string
    29  }
    30  
    31  func newPool(options Options) *pool {
    32  	return &pool{
    33  		size:         options.Size,
    34  		tr:           options.Transport,
    35  		ttl:          options.TTL,
    36  		closeTimeout: options.CloseTimeout,
    37  		conns:        make(map[string][]*poolConn),
    38  	}
    39  }
    40  
    41  func (p *pool) Close() error {
    42  	p.mu.Lock()
    43  	defer p.mu.Unlock()
    44  
    45  	var err error
    46  
    47  	for k, c := range p.conns {
    48  		for _, conn := range c {
    49  			if nerr := conn.close(); nerr != nil {
    50  				err = nerr
    51  			}
    52  		}
    53  
    54  		delete(p.conns, k)
    55  	}
    56  
    57  	return err
    58  }
    59  
    60  // NoOp the Close since we manage it.
    61  func (p *poolConn) Close() error {
    62  	return nil
    63  }
    64  
    65  func (p *poolConn) Id() string {
    66  	return p.id
    67  }
    68  
    69  func (p *poolConn) Created() time.Time {
    70  	return p.created
    71  }
    72  
    73  func (p *pool) Get(addr string, opts ...transport.DialOption) (Conn, error) {
    74  	p.mu.Lock()
    75  	conns := p.conns[addr]
    76  
    77  	// While we have conns check age and then return one
    78  	// otherwise we'll create a new conn
    79  	for len(conns) > 0 {
    80  		conn := conns[len(conns)-1]
    81  		conns = conns[:len(conns)-1]
    82  		p.conns[addr] = conns
    83  
    84  		// If conn is old kill it and move on
    85  		if d := time.Since(conn.Created()); d > p.ttl {
    86  			if err := conn.close(); err != nil {
    87  				p.mu.Unlock()
    88  				c, errConn := p.newConn(addr, opts)
    89  				if errConn != nil {
    90  					return nil, errConn
    91  				}
    92  				return c, err
    93  			}
    94  
    95  			continue
    96  		}
    97  
    98  		// We got a good conn, lets unlock and return it
    99  		p.mu.Unlock()
   100  
   101  		return conn, nil
   102  	}
   103  
   104  	p.mu.Unlock()
   105  
   106  	return p.newConn(addr, opts)
   107  }
   108  
   109  func (p *pool) newConn(addr string, opts []transport.DialOption) (Conn, error) {
   110  	// create new conn
   111  	c, err := p.tr.Dial(addr, opts...)
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  
   116  	return &poolConn{
   117  		Client:       c,
   118  		id:           uuid.New().String(),
   119  		closeTimeout: p.closeTimeout,
   120  		created:      time.Now(),
   121  	}, nil
   122  }
   123  
   124  func (p *pool) Release(conn Conn, err error) error {
   125  	// don't store the conn if it has errored
   126  	if err != nil {
   127  		return conn.(*poolConn).close()
   128  	}
   129  
   130  	// otherwise put it back for reuse
   131  	p.mu.Lock()
   132  	defer p.mu.Unlock()
   133  
   134  	conns := p.conns[conn.Remote()]
   135  	if len(conns) >= p.size {
   136  		return conn.(*poolConn).close()
   137  	}
   138  
   139  	p.conns[conn.Remote()] = append(conns, conn.(*poolConn))
   140  
   141  	return nil
   142  }
   143  
   144  func (p *poolConn) close() error {
   145  	ch := make(chan error)
   146  	go func() {
   147  		defer close(ch)
   148  		ch <- p.Client.Close()
   149  	}()
   150  	t := time.NewTimer(p.closeTimeout)
   151  	var err error
   152  	select {
   153  	case <-t.C:
   154  		err = errors.New("unable to close in time")
   155  	case err = <-ch:
   156  		t.Stop()
   157  	}
   158  	return err
   159  }