ledisdb/cmd/vendor/github.com/siddontang/goredis/client.go

package goredis

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

type PoolConn struct {
	*Conn
	c *Client
}

func (c *PoolConn) Close() {
	if c.Conn.isClosed() {
		return
	}

	c.c.put(c.Conn)
}

// force close inner connection and not put it into pool
func (c *PoolConn) Finalize() {
	c.Conn.Close()
}

type Client struct {
	sync.Mutex

	addr            string
	maxIdleConns    int
	readBufferSize  int
	writeBufferSize int
	password        string

	conns *list.List

	quit chan struct{}
	wg   sync.WaitGroup
}

func getProto(addr string) string {
	if strings.Contains(addr, "/") {
		return "unix"
	} else {
		return "tcp"
	}
}

func NewClient(addr string, password string) *Client {
	c := new(Client)

	c.addr = addr
	c.maxIdleConns = 4
	c.readBufferSize = 1024
	c.writeBufferSize = 1024
	c.password = password

	c.conns = list.New()
	c.quit = make(chan struct{})

	c.wg.Add(1)
	go c.onCheck()

	return c
}

func (c *Client) SetPassword(pass string) {
	c.password = pass
}

func (c *Client) SetReadBufferSize(s int) {
	c.readBufferSize = s
}

func (c *Client) SetWriteBufferSize(s int) {
	c.writeBufferSize = s
}

func (c *Client) SetMaxIdleConns(n int) {
	c.maxIdleConns = n
}

func (c *Client) Do(cmd string, args ...interface{}) (interface{}, error) {
	var co *Conn
	var err error
	var r interface{}

	for i := 0; i < 2; i++ {
		co, err = c.get()
		if err != nil {
			return nil, err
		}

		r, err = co.Do(cmd, args...)
		if err != nil {
			co.Close()

			if e, ok := err.(*net.OpError); ok && strings.Contains(e.Error(), "use of closed network connection") {
				//send to a closed connection, try again
				continue
			}

			return nil, err
		} else {
			c.put(co)
		}

		return r, nil
	}

	return nil, err
}

func (c *Client) Close() {
	c.Lock()
	defer c.Unlock()

	close(c.quit)
	c.wg.Wait()

	for c.conns.Len() > 0 {
		e := c.conns.Front()
		co := e.Value.(*Conn)
		c.conns.Remove(e)

		co.Close()
	}
}

func (c *Client) Get() (*PoolConn, error) {
	co, err := c.get()
	if err != nil {
		return nil, err
	}

	return &PoolConn{co, c}, err
}

func (c *Client) get() (co *Conn, err error) {
	c.Lock()
	if c.conns.Len() == 0 {
		c.Unlock()

		co, err = c.newConn(c.addr, c.password)
	} else {
		e := c.conns.Front()
		co = e.Value.(*Conn)
		c.conns.Remove(e)

		c.Unlock()
	}

	return
}

func (c *Client) put(conn *Conn) {
	c.Lock()
	defer c.Unlock()

	for c.conns.Len() >= c.maxIdleConns {
		// remove back
		e := c.conns.Back()
		co := e.Value.(*Conn)
		c.conns.Remove(e)

		co.Close()
	}

	c.conns.PushFront(conn)
}

func (c *Client) getIdle() *Conn {
	c.Lock()
	defer c.Unlock()

	if c.conns.Len() == 0 {
		return nil
	} else {
		e := c.conns.Back()
		co := e.Value.(*Conn)
		c.conns.Remove(e)
		return co
	}
}

func (c *Client) checkIdle() {
	co := c.getIdle()
	if co == nil {
		return
	}

	_, err := co.Do("PING")
	if err != nil {
		co.Close()
	} else {
		c.put(co)
	}
}

func (c *Client) onCheck() {
	t := time.NewTicker(3 * time.Second)

	defer func() {
		t.Stop()
		c.wg.Done()
	}()

	for {
		select {
		case <-t.C:
			c.checkIdle()
		case <-c.quit:
			return
		}
	}
}