package overflow_redis_pool import ( "bytes" "container/list" "context" "crypto/rand" "crypto/sha1" "errors" "io" "strconv" "sync" "time" "git.loafle.net/commons_go/redis_pool/internal" "github.com/garyburd/redigo/redis" ) var nowFunc = time.Now // for testing // ErrPoolExhausted is returned from a pool connection method (Do, Send, // Receive, Flush, Err) when the maximum number of database connections in the // pool has been reached. var ErrPoolExhausted = errors.New("redigo: connection pool exhausted") var ( errPoolClosed = errors.New("redigo: connection pool closed") errConnClosed = errors.New("redigo: connection closed") ) // Pool maintains a pool of connections. The application calls the Get method // to get a connection from the pool and the connection's Close method to // return the connection's resources to the pool. // // The following example shows how to use a pool in a web application. The // application creates a pool at application startup and makes it available to // request handlers using a package level variable. The pool configuration used // here is an example, not a recommendation. // // func newPool(addr string) *redis.Pool { // return &redis.Pool{ // MaxIdle: 3, // IdleTimeout: 240 * time.Second, // Dial: func () (redis.Conn, error) { return redis.Dial("tcp", addr) }, // } // } // // var ( // pool *redis.Pool // redisServer = flag.String("redisServer", ":6379", "") // ) // // func main() { // flag.Parse() // pool = newPool(*redisServer) // ... // } // // A request handler gets a connection from the pool and closes the connection // when the handler is done: // // func serveHome(w http.ResponseWriter, r *http.Request) { // conn := pool.Get() // defer conn.Close() // ... // } // // Use the Dial function to authenticate connections with the AUTH command or // select a database with the SELECT command: // // pool := &redis.Pool{ // // Other pool configuration not shown in this example. // Dial: func () (redis.Conn, error) { // c, err := redis.Dial("tcp", server) // if err != nil { // return nil, err // } // if _, err := c.Do("AUTH", password); err != nil { // c.Close() // return nil, err // } // if _, err := c.Do("SELECT", db); err != nil { // c.Close() // return nil, err // } // return c, nil // } // } // // Use the TestOnBorrow function to check the health of an idle connection // before the connection is returned to the application. This example PINGs // connections that have been idle more than a minute: // // pool := &redis.Pool{ // // Other pool configuration not shown in this example. // TestOnBorrow: func(c redis.Conn, t time.Time) error { // if time.Since(t) < time.Minute { // return nil // } // _, err := c.Do("PING") // return err // }, // } // type Pool interface { Get() redis.Conn ActiveCount() int Close() error } type pool struct { ctx context.Context handler PoolHandler // mu protects fields defined below. mu sync.Mutex cond *sync.Cond closed bool active int // Stack of idleConn with most recently used at the front. idle list.List } type idleConn struct { c redis.Conn t time.Time } // NewPool creates a new pool. // // Deprecated: Initialize the Pool directory as shown in the example. func NewPool(ctx context.Context, handler PoolHandler) Pool { p := &pool{ ctx: ctx, handler: handler, } return p } // Get gets a connection. The application must close the returned connection. // This method always returns a valid connection so that applications can defer // error handling to the first use of the connection. If there is an error // getting an underlying connection, then the connection Err, Do, Send, Flush // and Receive methods return that error. func (p *pool) Get() redis.Conn { c, err := p.get() if err != nil { return errorConnection{err} } return &pooledConnection{p: p, c: c} } // ActiveCount returns the number of active connections in the pool. func (p *pool) ActiveCount() int { p.mu.Lock() active := p.active p.mu.Unlock() return active } // Close releases the resources used by the pool. func (p *pool) Close() error { p.mu.Lock() idle := p.idle p.idle.Init() p.closed = true p.active -= idle.Len() if p.cond != nil { p.cond.Broadcast() } p.mu.Unlock() for e := idle.Front(); e != nil; e = e.Next() { e.Value.(idleConn).c.Close() } return nil } // release decrements the active count and signals waiters. The caller must // hold p.mu during the call. func (p *pool) release() { p.active -= 1 if p.cond != nil { p.cond.Signal() } } // get prunes stale connections and returns a connection from the idle list or // creates a new connection. func (p *pool) get() (redis.Conn, error) { p.mu.Lock() // Prune stale connections. if timeout := p.handler.GetIdleTimeout(); timeout > 0 { for i, n := 0, p.idle.Len(); i < n; i++ { e := p.idle.Back() if e == nil { break } ic := e.Value.(idleConn) if ic.t.Add(timeout).After(nowFunc()) { break } p.idle.Remove(e) p.release() p.mu.Unlock() ic.c.Close() p.mu.Lock() } } for { // Get idle connection. for i, n := 0, p.idle.Len(); i < n; i++ { e := p.idle.Front() if e == nil { break } ic := e.Value.(idleConn) p.idle.Remove(e) test := p.handler.TestOnBorrow p.mu.Unlock() if test == nil || test(ic.c, ic.t) == nil { return ic.c, nil } ic.c.Close() p.mu.Lock() p.release() } // Check for pool closed before dialing a new connection. if p.closed { p.mu.Unlock() return nil, errors.New("redigo: get on closed pool") } // Dial new connection if under limit. if p.handler.GetMaxActive() == 0 || p.active < p.handler.GetMaxActive() { dial := p.handler.Dial p.active += 1 p.mu.Unlock() c, err := dial() if err != nil { p.mu.Lock() p.release() p.mu.Unlock() c = nil } return c, err } if !p.handler.IsWait() { p.mu.Unlock() return nil, ErrPoolExhausted } if p.cond == nil { p.cond = sync.NewCond(&p.mu) } p.cond.Wait() } } func (p *pool) put(c redis.Conn, forceClose bool) error { err := c.Err() p.mu.Lock() if !p.closed && err == nil && !forceClose { p.idle.PushFront(idleConn{t: nowFunc(), c: c}) if p.idle.Len() > p.handler.GetMaxIdle() { c = p.idle.Remove(p.idle.Back()).(idleConn).c } else { c = nil } } if c == nil { if p.cond != nil { p.cond.Signal() } p.mu.Unlock() return nil } p.release() p.mu.Unlock() return c.Close() } type pooledConnection struct { p *pool c redis.Conn state int } var ( sentinel []byte sentinelOnce sync.Once ) func initSentinel() { p := make([]byte, 64) if _, err := rand.Read(p); err == nil { sentinel = p } else { h := sha1.New() io.WriteString(h, "Oops, rand failed. Use time instead.") io.WriteString(h, strconv.FormatInt(time.Now().UnixNano(), 10)) sentinel = h.Sum(nil) } } func (pc *pooledConnection) Close() error { c := pc.c if _, ok := c.(errorConnection); ok { return nil } pc.c = errorConnection{errConnClosed} if pc.state&internal.MultiState != 0 { c.Send("DISCARD") pc.state &^= (internal.MultiState | internal.WatchState) } else if pc.state&internal.WatchState != 0 { c.Send("UNWATCH") pc.state &^= internal.WatchState } if pc.state&internal.SubscribeState != 0 { c.Send("UNSUBSCRIBE") c.Send("PUNSUBSCRIBE") // To detect the end of the message stream, ask the server to echo // a sentinel value and read until we see that value. sentinelOnce.Do(initSentinel) c.Send("ECHO", sentinel) c.Flush() for { p, err := c.Receive() if err != nil { break } if p, ok := p.([]byte); ok && bytes.Equal(p, sentinel) { pc.state &^= internal.SubscribeState break } } } c.Do("") pc.p.put(c, pc.state != 0) return nil } func (pc *pooledConnection) Err() error { return pc.c.Err() } func (pc *pooledConnection) Do(commandName string, args ...interface{}) (reply interface{}, err error) { ci := internal.LookupCommandInfo(commandName) pc.state = (pc.state | ci.Set) &^ ci.Clear return pc.c.Do(commandName, args...) } func (pc *pooledConnection) Send(commandName string, args ...interface{}) error { ci := internal.LookupCommandInfo(commandName) pc.state = (pc.state | ci.Set) &^ ci.Clear return pc.c.Send(commandName, args...) } func (pc *pooledConnection) Flush() error { return pc.c.Flush() } func (pc *pooledConnection) Receive() (reply interface{}, err error) { return pc.c.Receive() } type errorConnection struct{ err error } func (ec errorConnection) Do(string, ...interface{}) (interface{}, error) { return nil, ec.err } func (ec errorConnection) Send(string, ...interface{}) error { return ec.err } func (ec errorConnection) Err() error { return ec.err } func (ec errorConnection) Close() error { return ec.err } func (ec errorConnection) Flush() error { return ec.err } func (ec errorConnection) Receive() (interface{}, error) { return nil, ec.err }