You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

394 lines
9.9 KiB

  1. // Copyright 2012 Gary Burd
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License"): you may
  4. // not use this file except in compliance with the License. You may obtain
  5. // a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11. // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12. // License for the specific language governing permissions and limitations
  13. // under the License.
  14. package redis
  15. import (
  16. "bytes"
  17. "container/list"
  18. "crypto/rand"
  19. "crypto/sha1"
  20. "errors"
  21. "io"
  22. "strconv"
  23. "sync"
  24. "time"
  25. "github.com/garyburd/redigo/internal"
  26. )
  27. var nowFunc = time.Now // for testing
  28. // ErrPoolExhausted is returned from a pool connection method (Do, Send,
  29. // Receive, Flush, Err) when the maximum number of database connections in the
  30. // pool has been reached.
  31. var ErrPoolExhausted = errors.New("redigo: connection pool exhausted")
  32. var (
  33. errPoolClosed = errors.New("redigo: connection pool closed")
  34. errConnClosed = errors.New("redigo: connection closed")
  35. )
  36. // Pool maintains a pool of connections. The application calls the Get method
  37. // to get a connection from the pool and the connection's Close method to
  38. // return the connection's resources to the pool.
  39. //
  40. // The following example shows how to use a pool in a web application. The
  41. // application creates a pool at application startup and makes it available to
  42. // request handlers using a global variable.
  43. //
  44. // func newPool(server, password string) *redis.Pool {
  45. // return &redis.Pool{
  46. // MaxIdle: 3,
  47. // IdleTimeout: 240 * time.Second,
  48. // Dial: func () (redis.Conn, error) {
  49. // c, err := redis.Dial("tcp", server)
  50. // if err != nil {
  51. // return nil, err
  52. // }
  53. // if _, err := c.Do("AUTH", password); err != nil {
  54. // c.Close()
  55. // return nil, err
  56. // }
  57. // return c, err
  58. // },
  59. // TestOnBorrow: func(c redis.Conn, t time.Time) error {
  60. // _, err := c.Do("PING")
  61. // return err
  62. // },
  63. // }
  64. // }
  65. //
  66. // var (
  67. // pool *redis.Pool
  68. // redisServer = flag.String("redisServer", ":6379", "")
  69. // redisPassword = flag.String("redisPassword", "", "")
  70. // )
  71. //
  72. // func main() {
  73. // flag.Parse()
  74. // pool = newPool(*redisServer, *redisPassword)
  75. // ...
  76. // }
  77. //
  78. // A request handler gets a connection from the pool and closes the connection
  79. // when the handler is done:
  80. //
  81. // func serveHome(w http.ResponseWriter, r *http.Request) {
  82. // conn := pool.Get()
  83. // defer conn.Close()
  84. // ....
  85. // }
  86. //
  87. type Pool struct {
  88. // Dial is an application supplied function for creating and configuring a
  89. // connection.
  90. //
  91. // The connection returned from Dial must not be in a special state
  92. // (subscribed to pubsub channel, transaction started, ...).
  93. Dial func() (Conn, error)
  94. // TestOnBorrow is an optional application supplied function for checking
  95. // the health of an idle connection before the connection is used again by
  96. // the application. Argument t is the time that the connection was returned
  97. // to the pool. If the function returns an error, then the connection is
  98. // closed.
  99. TestOnBorrow func(c Conn, t time.Time) error
  100. // Maximum number of idle connections in the pool.
  101. MaxIdle int
  102. // Maximum number of connections allocated by the pool at a given time.
  103. // When zero, there is no limit on the number of connections in the pool.
  104. MaxActive int
  105. // Close connections after remaining idle for this duration. If the value
  106. // is zero, then idle connections are not closed. Applications should set
  107. // the timeout to a value less than the server's timeout.
  108. IdleTimeout time.Duration
  109. // If Wait is true and the pool is at the MaxActive limit, then Get() waits
  110. // for a connection to be returned to the pool before returning.
  111. Wait bool
  112. // mu protects fields defined below.
  113. mu sync.Mutex
  114. cond *sync.Cond
  115. closed bool
  116. active int
  117. // Stack of idleConn with most recently used at the front.
  118. idle list.List
  119. }
  120. type idleConn struct {
  121. c Conn
  122. t time.Time
  123. }
  124. // NewPool creates a new pool.
  125. //
  126. // Deprecated: Initialize the Pool directory as shown in the example.
  127. func NewPool(newFn func() (Conn, error), maxIdle int) *Pool {
  128. return &Pool{Dial: newFn, MaxIdle: maxIdle}
  129. }
  130. // Get gets a connection. The application must close the returned connection.
  131. // This method always returns a valid connection so that applications can defer
  132. // error handling to the first use of the connection. If there is an error
  133. // getting an underlying connection, then the connection Err, Do, Send, Flush
  134. // and Receive methods return that error.
  135. func (p *Pool) Get() Conn {
  136. c, err := p.get()
  137. if err != nil {
  138. return errorConnection{err}
  139. }
  140. return &pooledConnection{p: p, c: c}
  141. }
  142. // ActiveCount returns the number of active connections in the pool.
  143. func (p *Pool) ActiveCount() int {
  144. p.mu.Lock()
  145. active := p.active
  146. p.mu.Unlock()
  147. return active
  148. }
  149. // Close releases the resources used by the pool.
  150. func (p *Pool) Close() error {
  151. p.mu.Lock()
  152. idle := p.idle
  153. p.idle.Init()
  154. p.closed = true
  155. p.active -= idle.Len()
  156. if p.cond != nil {
  157. p.cond.Broadcast()
  158. }
  159. p.mu.Unlock()
  160. for e := idle.Front(); e != nil; e = e.Next() {
  161. e.Value.(idleConn).c.Close()
  162. }
  163. return nil
  164. }
  165. // release decrements the active count and signals waiters. The caller must
  166. // hold p.mu during the call.
  167. func (p *Pool) release() {
  168. p.active -= 1
  169. if p.cond != nil {
  170. p.cond.Signal()
  171. }
  172. }
  173. // get prunes stale connections and returns a connection from the idle list or
  174. // creates a new connection.
  175. func (p *Pool) get() (Conn, error) {
  176. p.mu.Lock()
  177. // Prune stale connections.
  178. if timeout := p.IdleTimeout; timeout > 0 {
  179. for i, n := 0, p.idle.Len(); i < n; i++ {
  180. e := p.idle.Back()
  181. if e == nil {
  182. break
  183. }
  184. ic := e.Value.(idleConn)
  185. if ic.t.Add(timeout).After(nowFunc()) {
  186. break
  187. }
  188. p.idle.Remove(e)
  189. p.release()
  190. p.mu.Unlock()
  191. ic.c.Close()
  192. p.mu.Lock()
  193. }
  194. }
  195. for {
  196. // Get idle connection.
  197. for i, n := 0, p.idle.Len(); i < n; i++ {
  198. e := p.idle.Front()
  199. if e == nil {
  200. break
  201. }
  202. ic := e.Value.(idleConn)
  203. p.idle.Remove(e)
  204. test := p.TestOnBorrow
  205. p.mu.Unlock()
  206. if test == nil || test(ic.c, ic.t) == nil {
  207. return ic.c, nil
  208. }
  209. ic.c.Close()
  210. p.mu.Lock()
  211. p.release()
  212. }
  213. // Check for pool closed before dialing a new connection.
  214. if p.closed {
  215. p.mu.Unlock()
  216. return nil, errors.New("redigo: get on closed pool")
  217. }
  218. // Dial new connection if under limit.
  219. if p.MaxActive == 0 || p.active < p.MaxActive {
  220. dial := p.Dial
  221. p.active += 1
  222. p.mu.Unlock()
  223. c, err := dial()
  224. if err != nil {
  225. p.mu.Lock()
  226. p.release()
  227. p.mu.Unlock()
  228. c = nil
  229. }
  230. return c, err
  231. }
  232. if !p.Wait {
  233. p.mu.Unlock()
  234. return nil, ErrPoolExhausted
  235. }
  236. if p.cond == nil {
  237. p.cond = sync.NewCond(&p.mu)
  238. }
  239. p.cond.Wait()
  240. }
  241. }
  242. func (p *Pool) put(c Conn, forceClose bool) error {
  243. err := c.Err()
  244. p.mu.Lock()
  245. if !p.closed && err == nil && !forceClose {
  246. p.idle.PushFront(idleConn{t: nowFunc(), c: c})
  247. if p.idle.Len() > p.MaxIdle {
  248. c = p.idle.Remove(p.idle.Back()).(idleConn).c
  249. } else {
  250. c = nil
  251. }
  252. }
  253. if c == nil {
  254. if p.cond != nil {
  255. p.cond.Signal()
  256. }
  257. p.mu.Unlock()
  258. return nil
  259. }
  260. p.release()
  261. p.mu.Unlock()
  262. return c.Close()
  263. }
  264. type pooledConnection struct {
  265. p *Pool
  266. c Conn
  267. state int
  268. }
  269. var (
  270. sentinel []byte
  271. sentinelOnce sync.Once
  272. )
  273. func initSentinel() {
  274. p := make([]byte, 64)
  275. if _, err := rand.Read(p); err == nil {
  276. sentinel = p
  277. } else {
  278. h := sha1.New()
  279. io.WriteString(h, "Oops, rand failed. Use time instead.")
  280. io.WriteString(h, strconv.FormatInt(time.Now().UnixNano(), 10))
  281. sentinel = h.Sum(nil)
  282. }
  283. }
  284. func (pc *pooledConnection) Close() error {
  285. c := pc.c
  286. if _, ok := c.(errorConnection); ok {
  287. return nil
  288. }
  289. pc.c = errorConnection{errConnClosed}
  290. if pc.state&internal.MultiState != 0 {
  291. c.Send("DISCARD")
  292. pc.state &^= (internal.MultiState | internal.WatchState)
  293. } else if pc.state&internal.WatchState != 0 {
  294. c.Send("UNWATCH")
  295. pc.state &^= internal.WatchState
  296. }
  297. if pc.state&internal.SubscribeState != 0 {
  298. c.Send("UNSUBSCRIBE")
  299. c.Send("PUNSUBSCRIBE")
  300. // To detect the end of the message stream, ask the server to echo
  301. // a sentinel value and read until we see that value.
  302. sentinelOnce.Do(initSentinel)
  303. c.Send("ECHO", sentinel)
  304. c.Flush()
  305. for {
  306. p, err := c.Receive()
  307. if err != nil {
  308. break
  309. }
  310. if p, ok := p.([]byte); ok && bytes.Equal(p, sentinel) {
  311. pc.state &^= internal.SubscribeState
  312. break
  313. }
  314. }
  315. }
  316. c.Do("")
  317. pc.p.put(c, pc.state != 0)
  318. return nil
  319. }
  320. func (pc *pooledConnection) Err() error {
  321. return pc.c.Err()
  322. }
  323. func (pc *pooledConnection) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
  324. ci := internal.LookupCommandInfo(commandName)
  325. pc.state = (pc.state | ci.Set) &^ ci.Clear
  326. return pc.c.Do(commandName, args...)
  327. }
  328. func (pc *pooledConnection) Send(commandName string, args ...interface{}) error {
  329. ci := internal.LookupCommandInfo(commandName)
  330. pc.state = (pc.state | ci.Set) &^ ci.Clear
  331. return pc.c.Send(commandName, args...)
  332. }
  333. func (pc *pooledConnection) Flush() error {
  334. return pc.c.Flush()
  335. }
  336. func (pc *pooledConnection) Receive() (reply interface{}, err error) {
  337. return pc.c.Receive()
  338. }
  339. type errorConnection struct{ err error }
  340. func (ec errorConnection) Do(string, ...interface{}) (interface{}, error) { return nil, ec.err }
  341. func (ec errorConnection) Send(string, ...interface{}) error { return ec.err }
  342. func (ec errorConnection) Err() error { return ec.err }
  343. func (ec errorConnection) Close() error { return ec.err }
  344. func (ec errorConnection) Flush() error { return ec.err }
  345. func (ec errorConnection) Receive() (interface{}, error) { return nil, ec.err }