redis/pool.go

427 lines
7.4 KiB
Go
Raw Normal View History

2013-07-02 15:17:31 +04:00
package redis
import (
2013-11-07 18:20:15 +04:00
"errors"
2015-01-17 12:56:34 +03:00
"fmt"
2014-07-31 17:01:54 +04:00
"log"
2013-07-02 15:17:31 +04:00
"sync"
2015-01-31 16:20:37 +03:00
"sync/atomic"
2013-07-02 15:17:31 +04:00
"time"
"gopkg.in/bsm/ratelimit.v1"
2013-07-02 15:17:31 +04:00
)
2013-11-07 18:20:15 +04:00
var (
2015-01-31 16:20:37 +03:00
errClosed = errors.New("redis: client is closed")
errPoolTimeout = errors.New("redis: connection pool timeout")
2013-11-07 18:20:15 +04:00
)
2013-07-02 15:17:31 +04:00
type pool interface {
First() *conn
Get() (*conn, error)
2013-07-02 15:17:31 +04:00
Put(*conn) error
Remove(*conn) error
Len() int
FreeLen() int
2013-07-02 15:17:31 +04:00
Close() error
}
type connList struct {
cns []*conn
mx sync.Mutex
len int32 // atomic
size int32
2013-07-02 15:17:31 +04:00
}
func newConnList(size int) *connList {
return &connList{
cns: make([]*conn, 0, size),
size: int32(size),
2013-07-02 15:17:31 +04:00
}
}
func (l *connList) Len() int {
return int(atomic.LoadInt32(&l.len))
}
// Reserve reserves place in the list and returns true on success. The
// caller must add or remove connection if place was reserved.
func (l *connList) Reserve() bool {
len := atomic.AddInt32(&l.len, 1)
reserved := len <= l.size
if !reserved {
atomic.AddInt32(&l.len, -1)
2015-01-24 15:12:48 +03:00
}
return reserved
}
2015-01-24 15:12:48 +03:00
// Add adds connection to the list. The caller must reserve place first.
func (l *connList) Add(cn *conn) {
l.mx.Lock()
l.cns = append(l.cns, cn)
l.mx.Unlock()
2015-01-24 15:12:48 +03:00
}
func (l *connList) Remove(cn *conn) error {
defer l.mx.Unlock()
l.mx.Lock()
if cn == nil {
atomic.AddInt32(&l.len, -1)
return nil
2013-07-02 15:17:31 +04:00
}
for i, c := range l.cns {
if c == cn {
l.cns = append(l.cns[:i], l.cns[i+1:]...)
atomic.AddInt32(&l.len, -1)
return cn.Close()
}
2013-07-02 15:17:31 +04:00
}
2014-05-11 18:11:55 +04:00
panic("conn not found in the list")
2013-07-02 15:17:31 +04:00
}
func (l *connList) Replace(cn, newcn *conn) error {
defer l.mx.Unlock()
l.mx.Lock()
for i, c := range l.cns {
if c == cn {
l.cns[i] = newcn
return cn.Close()
}
}
panic("conn not found in the list")
2013-11-04 11:53:48 +04:00
}
func (l *connList) Close() (retErr error) {
l.mx.Lock()
for _, c := range l.cns {
if err := c.Close(); err != nil {
retErr = err
}
}
l.cns = nil
atomic.StoreInt32(&l.len, 0)
l.mx.Unlock()
return retErr
2015-01-31 16:20:37 +03:00
}
type connPoolOptions struct {
Dialer func() (*conn, error)
PoolSize int
PoolTimeout time.Duration
IdleTimeout time.Duration
IdleCheckFrequency time.Duration
}
2013-07-02 15:17:31 +04:00
type connPool struct {
rl *ratelimit.RateLimiter
opt *connPoolOptions
conns *connList
freeConns chan *conn
2014-05-11 18:11:55 +04:00
_closed int32
2015-01-17 12:56:34 +03:00
lastDialErr error
2013-07-02 15:17:31 +04:00
}
func newConnPool(opt *connPoolOptions) *connPool {
p := &connPool{
rl: ratelimit.New(2*opt.PoolSize, time.Second),
opt: opt,
conns: newConnList(opt.PoolSize),
freeConns: make(chan *conn, opt.PoolSize),
2015-01-31 16:20:37 +03:00
}
if p.opt.IdleTimeout > 0 && p.opt.IdleCheckFrequency > 0 {
go p.reaper()
}
return p
2015-01-31 16:20:37 +03:00
}
func (p *connPool) closed() bool {
return atomic.LoadInt32(&p._closed) == 1
}
func (p *connPool) isIdle(cn *conn) bool {
return p.opt.IdleTimeout > 0 && time.Since(cn.usedAt) > p.opt.IdleTimeout
}
2015-01-31 16:20:37 +03:00
// First returns first non-idle connection from the pool or nil if
// there are no connections.
func (p *connPool) First() *conn {
2015-01-31 16:20:37 +03:00
for {
select {
case cn := <-p.freeConns:
if p.isIdle(cn) {
p.conns.Remove(cn)
continue
2015-01-31 16:20:37 +03:00
}
return cn
2015-01-31 16:20:37 +03:00
default:
return nil
}
}
panic("not reached")
}
2013-07-02 15:17:31 +04:00
// wait waits for free non-idle connection. It returns nil on timeout.
func (p *connPool) wait(timeout time.Duration) *conn {
deadline := time.After(timeout)
2015-01-31 16:20:37 +03:00
for {
select {
case cn := <-p.freeConns:
if p.isIdle(cn) {
p.Remove(cn)
continue
2015-01-31 16:20:37 +03:00
}
return cn
2015-01-31 16:20:37 +03:00
case <-deadline:
return nil
2015-01-31 16:20:37 +03:00
}
2013-07-02 15:17:31 +04:00
}
2015-01-31 16:20:37 +03:00
panic("not reached")
2013-07-02 15:17:31 +04:00
}
2015-01-31 16:20:37 +03:00
// Establish a new connection
func (p *connPool) new() (*conn, error) {
if p.rl.Limit() {
2015-01-17 12:56:34 +03:00
err := fmt.Errorf(
"redis: you open connections too fast (last error: %v)",
2015-01-17 12:56:34 +03:00
p.lastDialErr,
)
return nil, err
}
cn, err := p.opt.Dialer()
2015-01-17 12:56:34 +03:00
if err != nil {
p.lastDialErr = err
return nil, err
}
return cn, nil
}
// Get returns existed connection from the pool or creates a new one
// if needed.
func (p *connPool) Get() (*conn, error) {
if p.closed() {
return nil, errClosed
2013-07-02 15:17:31 +04:00
}
2015-01-31 16:20:37 +03:00
// Fetch first non-idle connection, if available
if cn := p.First(); cn != nil {
return cn, nil
2013-11-07 18:20:15 +04:00
}
// Try to create a new one.
if p.conns.Reserve() {
2014-06-28 15:47:37 +04:00
cn, err := p.new()
2013-07-02 15:17:31 +04:00
if err != nil {
p.conns.Remove(nil)
return nil, err
2013-07-02 15:17:31 +04:00
}
p.conns.Add(cn)
return cn, nil
2013-07-02 15:17:31 +04:00
}
2015-01-31 16:20:37 +03:00
// Otherwise, wait for the available connection
if cn := p.wait(p.opt.PoolTimeout); cn != nil {
return cn, nil
}
return nil, errPoolTimeout
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 {
b, _ := cn.rd.ReadN(cn.rd.Buffered())
2014-07-31 17:01:54 +04:00
log.Printf("redis: connection has unread data: %q", b)
return p.Remove(cn)
2013-09-17 13:03:17 +04:00
}
2014-05-11 18:11:55 +04:00
if p.opt.IdleTimeout > 0 {
cn.usedAt = time.Now()
}
p.freeConns <- cn
2013-07-02 15:17:31 +04:00
return nil
}
func (p *connPool) Remove(cn *conn) error {
if p.closed() {
// Close already closed all connections.
return nil
}
// Replace existing connection with new one and unblock waiter.
newcn, err := p.new()
if err != nil {
return p.conns.Remove(cn)
2013-11-07 18:20:15 +04:00
}
p.freeConns <- newcn
return p.conns.Replace(cn, newcn)
}
// Len returns total number of connections.
2013-07-02 15:17:31 +04:00
func (p *connPool) Len() int {
return p.conns.Len()
2013-07-02 15:17:31 +04:00
}
// FreeLen returns number of free connections.
func (p *connPool) FreeLen() int {
return len(p.freeConns)
2013-07-02 15:17:31 +04:00
}
func (p *connPool) Close() error {
if !atomic.CompareAndSwapInt32(&p._closed, 0, 1) {
return errClosed
2013-11-07 18:20:15 +04:00
}
return p.conns.Close()
}
2015-01-31 16:20:37 +03:00
func (p *connPool) reaper() {
ticker := time.NewTicker(p.opt.IdleCheckFrequency)
defer ticker.Stop()
for _ = range ticker.C {
if p.closed() {
2014-05-11 18:11:55 +04:00
break
}
// pool.First removes idle connections from the pool and
// returns first non-idle connection. So just put returned
// connection back.
if cn := p.First(); cn != nil {
p.Put(cn)
2013-07-02 15:17:31 +04:00
}
}
}
//------------------------------------------------------------------------------
type singleConnPool struct {
pool pool
2014-06-28 15:47:37 +04:00
cnMtx sync.Mutex
cn *conn
2013-07-02 15:17:31 +04:00
reusable bool
closed bool
2013-07-02 15:17:31 +04:00
}
2014-06-28 15:47:37 +04:00
func newSingleConnPool(pool pool, reusable bool) *singleConnPool {
2013-07-02 15:17:31 +04:00
return &singleConnPool{
pool: pool,
reusable: reusable,
}
}
2014-06-28 15:47:37 +04:00
func (p *singleConnPool) SetConn(cn *conn) {
p.cnMtx.Lock()
p.cn = cn
p.cnMtx.Unlock()
}
func (p *singleConnPool) First() *conn {
defer p.cnMtx.Unlock()
p.cnMtx.Lock()
return p.cn
}
func (p *singleConnPool) Get() (*conn, error) {
2014-06-28 15:47:37 +04:00
defer p.cnMtx.Unlock()
p.cnMtx.Lock()
if p.closed {
return nil, errClosed
}
2013-07-02 15:17:31 +04:00
if p.cn != nil {
return p.cn, nil
2013-07-02 15:17:31 +04:00
}
cn, err := p.pool.Get()
2013-07-02 15:17:31 +04:00
if err != nil {
return nil, err
2013-07-02 15:17:31 +04:00
}
p.cn = cn
2014-06-28 15:47:37 +04:00
return p.cn, nil
2013-07-02 15:17:31 +04:00
}
func (p *singleConnPool) Put(cn *conn) error {
2014-06-28 15:47:37 +04:00
defer p.cnMtx.Unlock()
p.cnMtx.Lock()
2013-07-02 15:17:31 +04:00
if p.cn != cn {
panic("p.cn != cn")
}
if p.closed {
return errClosed
}
2013-07-02 15:17:31 +04:00
return nil
}
2014-06-28 15:47:37 +04:00
func (p *singleConnPool) put() error {
err := p.pool.Put(p.cn)
p.cn = nil
return err
}
2013-07-02 15:17:31 +04:00
func (p *singleConnPool) Remove(cn *conn) error {
2014-06-28 15:47:37 +04:00
defer p.cnMtx.Unlock()
p.cnMtx.Lock()
2014-05-11 18:11:55 +04:00
if p.cn == nil {
panic("p.cn == nil")
}
2013-07-02 15:17:31 +04:00
if p.cn != cn {
panic("p.cn != cn")
}
if p.closed {
return errClosed
}
2014-05-11 18:11:55 +04:00
return p.remove()
}
func (p *singleConnPool) remove() error {
err := p.pool.Remove(p.cn)
2013-07-02 15:17:31 +04:00
p.cn = nil
2014-05-11 18:11:55 +04:00
return err
2013-07-02 15:17:31 +04:00
}
func (p *singleConnPool) Len() int {
2014-06-28 15:47:37 +04:00
defer p.cnMtx.Unlock()
p.cnMtx.Lock()
2013-07-02 15:17:31 +04:00
if p.cn == nil {
return 0
}
return 1
}
func (p *singleConnPool) FreeLen() int {
2014-06-28 15:47:37 +04:00
defer p.cnMtx.Unlock()
p.cnMtx.Lock()
2013-11-07 18:20:15 +04:00
if p.cn == nil {
return 0
}
return 1
}
2013-07-02 15:17:31 +04:00
func (p *singleConnPool) Close() error {
2014-06-28 15:47:37 +04:00
defer p.cnMtx.Unlock()
p.cnMtx.Lock()
if p.closed {
return nil
}
p.closed = true
2013-07-02 15:17:31 +04:00
var err error
if p.cn != nil {
if p.reusable {
2014-06-28 15:47:37 +04:00
err = p.put()
2013-07-02 15:17:31 +04:00
} else {
2014-06-28 15:47:37 +04:00
err = p.remove()
2013-07-02 15:17:31 +04:00
}
}
return err
}