redis/v2/pool.go

324 lines
5.3 KiB
Go
Raw Normal View History

2013-07-02 15:17:31 +04:00
package redis
import (
"container/list"
2013-11-07 18:20:15 +04:00
"errors"
2013-07-02 15:17:31 +04:00
"net"
"sync"
"time"
2013-11-07 18:20:15 +04:00
"github.com/golang/glog"
2013-07-02 15:17:31 +04:00
"github.com/vmihailenco/bufio"
)
2013-11-07 18:20:15 +04:00
var (
errPoolClosed = errors.New("attempt to use closed connection pool")
)
2013-07-02 15:17:31 +04:00
type pool interface {
Get() (*conn, bool, error)
Put(*conn) error
Remove(*conn) error
Len() int
2013-11-07 18:20:15 +04:00
Size() int
2013-07-02 15:17:31 +04:00
Close() error
}
//------------------------------------------------------------------------------
type conn struct {
2013-09-29 12:11:18 +04:00
cn net.Conn
rd reader
2013-11-07 18:20:15 +04:00
inUse bool
2013-09-29 12:11:18 +04:00
usedAt time.Time
2013-07-02 15:17:31 +04:00
readTimeout, writeTimeout time.Duration
2013-11-07 18:20:15 +04:00
elem *list.Element
2013-07-02 15:17:31 +04:00
}
2013-11-07 18:20:15 +04:00
func newConnFunc(dial func() (net.Conn, error)) func() (*conn, error) {
return func() (*conn, error) {
netcn, err := dial()
if err != nil {
return nil, err
}
cn := &conn{
cn: netcn,
}
cn.rd = bufio.NewReader(cn)
return cn, nil
2013-07-02 15:17:31 +04:00
}
}
func (cn *conn) Read(b []byte) (int, error) {
2013-09-11 19:06:36 +04:00
if cn.readTimeout != 0 {
2013-09-29 12:11:18 +04:00
cn.cn.SetReadDeadline(time.Now().Add(cn.readTimeout))
2013-09-11 20:22:10 +04:00
} else {
2013-09-29 12:11:18 +04:00
cn.cn.SetReadDeadline(time.Time{})
2013-07-02 15:17:31 +04:00
}
2013-09-29 12:11:18 +04:00
return cn.cn.Read(b)
2013-07-02 15:17:31 +04:00
}
func (cn *conn) Write(b []byte) (int, error) {
2013-09-11 19:06:36 +04:00
if cn.writeTimeout != 0 {
2013-09-29 12:11:18 +04:00
cn.cn.SetWriteDeadline(time.Now().Add(cn.writeTimeout))
2013-09-11 20:22:10 +04:00
} else {
2013-09-29 12:11:18 +04:00
cn.cn.SetReadDeadline(time.Time{})
2013-07-02 15:17:31 +04:00
}
2013-09-29 12:11:18 +04:00
return cn.cn.Write(b)
2013-07-02 15:17:31 +04:00
}
2013-11-04 11:53:48 +04:00
func (cn *conn) Close() error {
return cn.cn.Close()
}
2013-07-02 15:17:31 +04:00
//------------------------------------------------------------------------------
type connPool struct {
2013-11-07 18:20:15 +04:00
New func() (*conn, error)
2013-07-02 15:17:31 +04:00
cond *sync.Cond
conns *list.List
2013-11-07 18:20:15 +04:00
len int
maxSize int
idleTimeout time.Duration
closed bool
2013-07-02 15:17:31 +04:00
}
func newConnPool(
2013-11-07 18:20:15 +04:00
dial func() (*conn, error),
2013-07-02 15:17:31 +04:00
maxSize int,
2013-09-11 20:22:10 +04:00
idleTimeout time.Duration,
2013-07-02 15:17:31 +04:00
) *connPool {
return &connPool{
2013-11-07 18:20:15 +04:00
New: dial,
2013-07-02 15:17:31 +04:00
cond: sync.NewCond(&sync.Mutex{}),
conns: list.New(),
2013-09-11 20:22:10 +04:00
maxSize: maxSize,
idleTimeout: idleTimeout,
2013-07-02 15:17:31 +04:00
}
}
func (p *connPool) Get() (*conn, bool, error) {
p.cond.L.Lock()
2013-11-07 18:20:15 +04:00
if p.closed {
p.cond.L.Unlock()
return nil, false, errPoolClosed
2013-07-02 15:17:31 +04:00
}
if p.idleTimeout > 0 {
for e := p.conns.Front(); e != nil; e = e.Next() {
cn := e.Value.(*conn)
2013-11-07 18:20:15 +04:00
if cn.inUse {
break
}
2013-09-29 12:11:18 +04:00
if time.Since(cn.usedAt) > p.idleTimeout {
2013-11-07 18:20:15 +04:00
if err := p.Remove(cn); err != nil {
glog.Errorf("Remove failed: %s", err)
}
2013-07-02 15:17:31 +04:00
}
}
}
2013-11-07 18:20:15 +04:00
for p.conns.Len() >= p.maxSize && p.len == 0 {
p.cond.Wait()
}
if p.len > 0 {
elem := p.conns.Front()
cn := elem.Value.(*conn)
if cn.inUse {
panic("pool: precondition failed")
}
cn.inUse = true
p.conns.MoveToBack(elem)
p.len--
p.cond.L.Unlock()
return cn, false, nil
}
if p.conns.Len() < p.maxSize {
cn, err := p.New()
2013-07-02 15:17:31 +04:00
if err != nil {
2013-11-07 18:20:15 +04:00
p.cond.L.Unlock()
2013-07-02 15:17:31 +04:00
return nil, false, err
}
2013-11-07 18:20:15 +04:00
cn.inUse = true
cn.elem = p.conns.PushBack(cn)
p.cond.L.Unlock()
return cn, true, nil
2013-07-02 15:17:31 +04:00
}
2013-11-07 18:20:15 +04:00
panic("not reached")
2013-07-02 15:17:31 +04:00
}
func (p *connPool) Put(cn *conn) error {
2013-09-29 12:11:18 +04:00
if cn.rd.Buffered() != 0 {
panic("redis: attempt to put connection with buffered data")
2013-09-17 13:03:17 +04:00
}
2013-07-02 15:17:31 +04:00
p.cond.L.Lock()
2013-11-07 18:20:15 +04:00
if p.closed {
p.cond.L.Unlock()
return errPoolClosed
}
cn.inUse = false
2013-09-29 12:11:18 +04:00
cn.usedAt = time.Now()
2013-11-07 18:20:15 +04:00
p.conns.MoveToFront(cn.elem)
p.len++
2013-07-02 15:17:31 +04:00
p.cond.Signal()
p.cond.L.Unlock()
return nil
}
2013-11-07 18:20:15 +04:00
func (p *connPool) Remove(cn *conn) (err error) {
p.cond.L.Lock()
if p.closed {
// Noop, connection is already closed.
p.cond.L.Unlock()
return nil
}
2013-07-02 15:17:31 +04:00
if cn != nil {
2013-11-04 11:53:48 +04:00
err = cn.Close()
2013-07-02 15:17:31 +04:00
}
2013-11-07 18:20:15 +04:00
p.conns.Remove(cn.elem)
cn.elem = nil
2013-07-02 15:17:31 +04:00
p.cond.Signal()
p.cond.L.Unlock()
return err
}
2013-11-07 18:20:15 +04:00
// Returns number of idle connections.
2013-07-02 15:17:31 +04:00
func (p *connPool) Len() int {
2013-09-11 19:06:36 +04:00
defer p.cond.L.Unlock()
p.cond.L.Lock()
2013-11-07 18:20:15 +04:00
return p.len
2013-07-02 15:17:31 +04:00
}
2013-11-07 18:20:15 +04:00
// Returns size of the pool.
2013-07-02 15:17:31 +04:00
func (p *connPool) Size() int {
2013-09-11 19:06:36 +04:00
defer p.cond.L.Unlock()
p.cond.L.Lock()
2013-11-07 18:20:15 +04:00
return p.conns.Len()
2013-07-02 15:17:31 +04:00
}
func (p *connPool) Close() error {
defer p.cond.L.Unlock()
p.cond.L.Lock()
2013-11-07 18:20:15 +04:00
if p.closed {
return nil
}
p.closed = true
var retErr error
2013-07-02 15:17:31 +04:00
for e := p.conns.Front(); e != nil; e = e.Next() {
2013-11-07 18:20:15 +04:00
cn := e.Value.(*conn)
if err := cn.Close(); err != nil {
glog.Errorf("cn.Close failed: %s", err)
retErr = err
2013-07-02 15:17:31 +04:00
}
2013-11-07 18:20:15 +04:00
cn.elem = nil
2013-07-02 15:17:31 +04:00
}
2013-11-07 18:20:15 +04:00
p.conns = nil
return retErr
2013-07-02 15:17:31 +04:00
}
//------------------------------------------------------------------------------
type singleConnPool struct {
pool pool
l sync.RWMutex
cn *conn
reusable bool
}
func newSingleConnPool(pool pool, cn *conn, reusable bool) *singleConnPool {
return &singleConnPool{
pool: pool,
cn: cn,
reusable: reusable,
}
}
func (p *singleConnPool) Get() (*conn, bool, error) {
p.l.RLock()
if p.cn != nil {
p.l.RUnlock()
return p.cn, false, nil
}
p.l.RUnlock()
p.l.Lock()
cn, isNew, err := p.pool.Get()
if err != nil {
2013-11-07 18:20:15 +04:00
p.l.Unlock()
2013-07-02 15:17:31 +04:00
return nil, false, err
}
p.cn = cn
2013-11-07 18:20:15 +04:00
p.l.Unlock()
2013-07-02 15:17:31 +04:00
return cn, isNew, nil
}
func (p *singleConnPool) Put(cn *conn) error {
p.l.Lock()
if p.cn != cn {
panic("p.cn != cn")
}
2013-11-07 18:20:15 +04:00
p.l.Unlock()
2013-07-02 15:17:31 +04:00
return nil
}
func (p *singleConnPool) Remove(cn *conn) error {
p.l.Lock()
if p.cn != cn {
panic("p.cn != cn")
}
p.cn = nil
2013-11-07 18:20:15 +04:00
p.l.Unlock()
2013-07-02 15:17:31 +04:00
return nil
}
func (p *singleConnPool) Len() int {
defer p.l.Unlock()
p.l.Lock()
if p.cn == nil {
return 0
}
return 1
}
2013-11-07 18:20:15 +04:00
func (p *singleConnPool) Size() int {
defer p.l.Unlock()
p.l.Lock()
if p.cn == nil {
return 0
}
return 1
}
2013-07-02 15:17:31 +04:00
func (p *singleConnPool) Close() error {
defer p.l.Unlock()
p.l.Lock()
var err error
if p.cn != nil {
if p.reusable {
err = p.pool.Put(p.cn)
} else {
err = p.pool.Remove(p.cn)
}
}
p.cn = nil
return err
}