Add MinIdleConns

Author: vmihailenco

internal/pool/conn.go

@@ -19,6 +19,7 @@ type Conn struct {
 	concurrentReadWrite bool
 
 	Inited bool
+	pooled bool
 	usedAt atomic.Value
 }
 

internal/pool/pool.go

@@ -53,6 +53,7 @@ type Options struct {
 	OnClose func(*Conn) error
 
 	PoolSize           int
+	MinIdleConns       int
 	PoolTimeout        time.Duration
 	IdleTimeout        time.Duration
 	IdleCheckFrequency time.Duration
@@ -63,16 +64,16 @@ type ConnPool struct {
 
 	dialErrorsNum uint32 // atomic
 
-	lastDialError   error
 	lastDialErrorMu sync.RWMutex
+	lastDialError   error
 
 	queue chan struct{}
 
-	connsMu sync.Mutex
-	conns   []*Conn
-
-	idleConnsMu sync.RWMutex
-	idleConns   []*Conn
+	connsMu      sync.Mutex
+	conns        []*Conn
+	idleConns    []*Conn
+	poolSize     int
+	idleConnsLen int
 
 	stats Stats
 
@@ -90,26 +91,64 @@ func NewConnPool(opt *Options) *ConnPool {
 		idleConns: make([]*Conn, 0, opt.PoolSize),
 	}
 
+	for i := 0; i < opt.MinIdleConns; i++ {
+		p.checkMinIdleConns()
+	}
+
 	if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 {
 		go p.reaper(opt.IdleCheckFrequency)
 	}
 
 	return p
 }
 
+func (p *ConnPool) checkMinIdleConns() {
+	if p.opt.MinIdleConns == 0 {
+		return
+	}
+	if p.poolSize < p.opt.PoolSize && p.idleConnsLen < p.opt.MinIdleConns {
+		p.poolSize++
+		p.idleConnsLen++
+		go p.addIdleConn()
+	}
+}
+
+func (p *ConnPool) addIdleConn() {
+	cn, err := p.newConn(true)
+	if err != nil {
+		return
+	}
+
+	p.connsMu.Lock()
+	p.conns = append(p.conns, cn)
+	p.idleConns = append(p.idleConns, cn)
+	p.connsMu.Unlock()
+}
+
 func (p *ConnPool) NewConn() (*Conn, error) {
-	cn, err := p.newConn()
+	return p._NewConn(false)
+}
+
+func (p *ConnPool) _NewConn(pooled bool) (*Conn, error) {
+	cn, err := p.newConn(pooled)
 	if err != nil {
 		return nil, err
 	}
 
 	p.connsMu.Lock()
 	p.conns = append(p.conns, cn)
+	if pooled {
+		if p.poolSize < p.opt.PoolSize {
+			p.poolSize++
+		} else {
+			cn.pooled = false
+		}
+	}
 	p.connsMu.Unlock()
 	return cn, nil
 }
 
-func (p *ConnPool) newConn() (*Conn, error) {
+func (p *ConnPool) newConn(pooled bool) (*Conn, error) {
 	if p.closed() {
 		return nil, ErrClosed
 	}
@@ -127,7 +166,9 @@ func (p *ConnPool) newConn() (*Conn, error) {
 		return nil, err
 	}
 
-	return NewConn(netConn), nil
+	cn := NewConn(netConn)
+	cn.pooled = pooled
+	return cn, nil
 }
 
 func (p *ConnPool) tryDial() {
@@ -174,9 +215,9 @@ func (p *ConnPool) Get() (*Conn, error) {
 	}
 
 	for {
-		p.idleConnsMu.Lock()
+		p.connsMu.Lock()
 		cn := p.popIdle()
-		p.idleConnsMu.Unlock()
+		p.connsMu.Unlock()
 
 		if cn == nil {
 			break
@@ -193,7 +234,7 @@ func (p *ConnPool) Get() (*Conn, error) {
 
 	atomic.AddUint32(&p.stats.Misses, 1)
 
-	newcn, err := p.NewConn()
+	newcn, err := p._NewConn(true)
 	if err != nil {
 		p.freeTurn()
 		return nil, err
@@ -241,7 +282,8 @@ func (p *ConnPool) popIdle() *Conn {
 	idx := len(p.idleConns) - 1
 	cn := p.idleConns[idx]
 	p.idleConns = p.idleConns[:idx]
-
+	p.idleConnsLen--
+	p.checkMinIdleConns()
 	return cn
 }
 
@@ -253,9 +295,15 @@ func (p *ConnPool) Put(cn *Conn) {
 		return
 	}
 
-	p.idleConnsMu.Lock()
+	if !cn.pooled {
+		p.Remove(cn)
+		return
+	}
+
+	p.connsMu.Lock()
 	p.idleConns = append(p.idleConns, cn)
-	p.idleConnsMu.Unlock()
+	p.idleConnsLen++
+	p.connsMu.Unlock()
 	p.freeTurn()
 }
 
@@ -275,6 +323,10 @@ func (p *ConnPool) removeConn(cn *Conn) {
 	for i, c := range p.conns {
 		if c == cn {
 			p.conns = append(p.conns[:i], p.conns[i+1:]...)
+			if cn.pooled {
+				p.poolSize--
+				p.checkMinIdleConns()
+			}
 			break
 		}
 	}
@@ -291,17 +343,17 @@ func (p *ConnPool) closeConn(cn *Conn) error {
 // Len returns total number of connections.
 func (p *ConnPool) Len() int {
 	p.connsMu.Lock()
-	l := len(p.conns)
+	n := p.poolSize
 	p.connsMu.Unlock()
-	return l
+	return n
 }
 
 // FreeLen returns number of idle connections.
 func (p *ConnPool) IdleLen() int {
-	p.idleConnsMu.RLock()
-	l := len(p.idleConns)
-	p.idleConnsMu.RUnlock()
-	return l
+	p.connsMu.Lock()
+	n := p.idleConnsLen
+	p.connsMu.Unlock()
+	return n
 }
 
 func (p *ConnPool) Stats() *Stats {
@@ -349,11 +401,10 @@ func (p *ConnPool) Close() error {
 		}
 	}
 	p.conns = nil
-	p.connsMu.Unlock()
-
-	p.idleConnsMu.Lock()
+	p.poolSize = 0
 	p.idleConns = nil
-	p.idleConnsMu.Unlock()
+	p.idleConnsLen = 0
+	p.connsMu.Unlock()
 
 	return firstErr
 }
@@ -369,6 +420,7 @@ func (p *ConnPool) reapStaleConn() *Conn {
 	}
 
 	p.idleConns = append(p.idleConns[:0], p.idleConns[1:]...)
+	p.idleConnsLen--
 
 	return cn
 }
@@ -378,9 +430,9 @@ func (p *ConnPool) ReapStaleConns() (int, error) {
 	for {
 		p.getTurn()
 
-		p.idleConnsMu.Lock()
+		p.connsMu.Lock()
 		cn := p.reapStaleConn()
-		p.idleConnsMu.Unlock()
+		p.connsMu.Unlock()
 
 		if cn != nil {
 			p.removeConn(cn)