redis/v2/pool.go

package redis

import (
	"container/list"
	"net"
	"sync"
	"time"

	"github.com/vmihailenco/bufio"
)

type pool interface {
	Get() (*conn, bool, error)
	Put(*conn) error
	Remove(*conn) error
	Len() int
	Close() error
}

//------------------------------------------------------------------------------

type conn struct {
	Cn net.Conn
	Rd reader

	UsedAt time.Time

	readTimeout, writeTimeout time.Duration
}

func newConn(netcn net.Conn, readTimeout, writeTimeout time.Duration) *conn {
	cn := &conn{
		Cn: netcn,

		readTimeout:  readTimeout,
		writeTimeout: writeTimeout,
	}
	cn.Rd = bufio.NewReaderSize(netcn, 1024)
	return cn
}

func (cn *conn) Read(b []byte) (int, error) {
	if cn.readTimeout > 0 {
		cn.Cn.SetReadDeadline(time.Now().Add(cn.readTimeout))
	}
	return cn.Cn.Read(b)
}

func (cn *conn) Write(b []byte) (int, error) {
	if cn.writeTimeout > 0 {
		cn.Cn.SetWriteDeadline(time.Now().Add(cn.writeTimeout))
	}
	return cn.Cn.Write(b)
}

//------------------------------------------------------------------------------

type connPool struct {
	dial  func() (net.Conn, error)
	close func(net.Conn) error

	cond  *sync.Cond
	conns *list.List

	readTimeout, writeTimeout time.Duration

	size, maxSize int
	idleTimeout   time.Duration
}

func newConnPool(
	dial func() (net.Conn, error),
	close func(net.Conn) error,
	maxSize int,
	readTimeout, writeTimeout, idleTimeout time.Duration,
) *connPool {
	return &connPool{
		dial:  dial,
		close: close,

		cond:  sync.NewCond(&sync.Mutex{}),
		conns: list.New(),

		maxSize: maxSize,

		readTimeout:  readTimeout,
		writeTimeout: writeTimeout,
		idleTimeout:  idleTimeout,
	}
}

func (p *connPool) Get() (*conn, bool, error) {
	defer p.cond.L.Unlock()
	p.cond.L.Lock()

	for p.conns.Len() == 0 && p.size >= p.maxSize {
		p.cond.Wait()
	}

	if p.idleTimeout > 0 {
		for e := p.conns.Front(); e != nil; e = e.Next() {
			cn := e.Value.(*conn)
			if time.Since(cn.UsedAt) > p.idleTimeout {
				p.conns.Remove(e)
			}
		}
	}

	if p.conns.Len() == 0 {
		rw, err := p.dial()
		if err != nil {
			return nil, false, err
		}

		p.size++
		return newConn(rw, p.readTimeout, p.writeTimeout), true, nil
	}

	elem := p.conns.Front()
	p.conns.Remove(elem)
	return elem.Value.(*conn), false, nil
}

func (p *connPool) Put(cn *conn) error {
	p.cond.L.Lock()
	cn.UsedAt = time.Now()
	p.conns.PushFront(cn)
	p.cond.Signal()
	p.cond.L.Unlock()
	return nil
}

func (p *connPool) Remove(cn *conn) error {
	var err error
	if cn != nil {
		err = p.closeConn(cn)
	}
	p.cond.L.Lock()
	p.size--
	p.cond.Signal()
	p.cond.L.Unlock()
	return err
}

func (p *connPool) Len() int {
	return p.conns.Len()
}

func (p *connPool) Size() int {
	return p.size
}

func (p *connPool) Close() error {
	defer p.cond.L.Unlock()
	p.cond.L.Lock()

	for e := p.conns.Front(); e != nil; e = e.Next() {
		if err := p.closeConn(e.Value.(*conn)); err != nil {
			return err
		}
	}
	p.conns.Init()
	p.size = 0

	return nil
}

func (p *connPool) closeConn(cn *conn) error {
	if p.close != nil {
		return p.close(cn.Cn)
	} else {
		return cn.Cn.Close()
	}
}

//------------------------------------------------------------------------------

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()

	defer p.l.Unlock()
	p.l.Lock()

	cn, isNew, err := p.pool.Get()
	if err != nil {
		return nil, false, err
	}
	p.cn = cn

	return cn, isNew, nil
}

func (p *singleConnPool) Put(cn *conn) error {
	defer p.l.Unlock()
	p.l.Lock()
	if p.cn != cn {
		panic("p.cn != cn")
	}
	return nil
}

func (p *singleConnPool) Remove(cn *conn) error {
	defer p.l.Unlock()
	p.l.Lock()
	if p.cn != cn {
		panic("p.cn != cn")
	}
	p.cn = nil
	return nil
}

func (p *singleConnPool) Len() int {
	defer p.l.Unlock()
	p.l.Lock()
	if p.cn == nil {
		return 0
	}
	return 1
}

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
}