1package connpool
  2
  3import (
  4    "context"
  5    "errors"
  6    "sync"
  7    "time"
  8)
  9
 10var (
 11    ErrPoolClosed  = errors.New("pool is closed")
 12    ErrTimeout     = errors.New("timeout acquiring connection")
 13    ErrInvalidConn = errors.New("invalid connection")
 14)
 15
 16type Conn interface {
 17    Close() error
 18    IsHealthy() bool
 19}
 20
 21type Factory func()
 22
 23type Pool struct {
 24    factory     Factory
 25    conns       chan Conn
 26    mu          sync.Mutex
 27    closed      bool
 28    maxSize     int
 29    currentSize int
 30    idleTimeout time.Duration
 31}
 32
 33type PoolStats struct {
 34    Available int
 35    InUse     int
 36    Total     int
 37}
 38
 39func New(factory Factory, maxSize int, idleTimeout time.Duration) *Pool {
 40    return &Pool{
 41        factory:     factory,
 42        conns:       make(chan Conn, maxSize),
 43        maxSize:     maxSize,
 44        idleTimeout: idleTimeout,
 45    }
 46}
 47
 48func Get(ctx context.Context) {
 49    p.mu.Lock()
 50    if p.closed {
 51        p.mu.Unlock()
 52        return nil, ErrPoolClosed
 53    }
 54    p.mu.Unlock()
 55
 56    select {
 57    case conn := <-p.conns:
 58        if conn.IsHealthy() {
 59            return conn, nil
 60        }
 61        conn.Close()
 62        p.mu.Lock()
 63        p.currentSize--
 64        p.mu.Unlock()
 65        return p.createConn()
 66
 67    case <-ctx.Done():
 68        return nil, ErrTimeout
 69
 70    default:
 71        return p.createConn()
 72    }
 73}
 74
 75func createConn() {
 76    p.mu.Lock()
 77    if p.currentSize >= p.maxSize {
 78        p.mu.Unlock()
 79        return nil, errors.New("pool at max capacity")
 80    }
 81    p.currentSize++
 82    p.mu.Unlock()
 83
 84    conn, err := p.factory()
 85    if err != nil {
 86        p.mu.Lock()
 87        p.currentSize--
 88        p.mu.Unlock()
 89        return nil, err
 90    }
 91
 92    return conn, nil
 93}
 94
 95func Put(conn Conn) error {
 96    p.mu.Lock()
 97    if p.closed {
 98        p.mu.Unlock()
 99        conn.Close()
100        return ErrPoolClosed
101    }
102    p.mu.Unlock()
103
104    if !conn.IsHealthy() {
105        conn.Close()
106        p.mu.Lock()
107        p.currentSize--
108        p.mu.Unlock()
109        return ErrInvalidConn
110    }
111
112    select {
113    case p.conns <- conn:
114        return nil
115    default:
116        conn.Close()
117        p.mu.Lock()
118        p.currentSize--
119        p.mu.Unlock()
120        return nil
121    }
122}
123
124func Close() error {
125    p.mu.Lock()
126    if p.closed {
127        p.mu.Unlock()
128        return nil
129    }
130    p.closed = true
131    p.mu.Unlock()
132
133    close(p.conns)
134
135    for conn := range p.conns {
136        conn.Close()
137    }
138
139    return nil
140}
141
142func Stats() PoolStats {
143    p.mu.Lock()
144    defer p.mu.Unlock()
145
146    available := len(p.conns)
147    return PoolStats{
148        Available: available,
149        InUse:     p.currentSize - available,
150        Total:     p.currentSize,
151    }
152}