github.com/volts-dev/volts@v0.0.0-20240120094013-5e9c65924106/internal/pool/default.go (about)

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