Simplify connection management with sticky connection pool. Fixes #260.

pubsub.go

@@ -17,16 +17,18 @@ func (c *Client) Publish(channel, message string) *IntCmd {
 // http://redis.io/topics/pubsub. It's NOT safe for concurrent use by
 // multiple goroutines.
 type PubSub struct {
-	*baseClient
+	base *baseClient
 
 	channels []string
 	patterns []string
+
+	nsub int // number of active subscriptions
 }
 
 // Deprecated. Use Subscribe/PSubscribe instead.
 func (c *Client) PubSub() *PubSub {
 	return &PubSub{
-		baseClient: &baseClient{
+		base: &baseClient{
 			opt:      c.opt,
 			connPool: newStickyConnPool(c.connPool, false),
 		},
@@ -46,7 +48,7 @@ func (c *Client) PSubscribe(channels ...string) (*PubSub, error) {
 }
 
 func (c *PubSub) subscribe(cmd string, channels ...string) error {
-	cn, _, err := c.conn()
+	cn, _, err := c.base.conn()
 	if err != nil {
 		return err
 	}
@@ -65,6 +67,7 @@ func (c *PubSub) Subscribe(channels ...string) error {
 	err := c.subscribe("SUBSCRIBE", channels...)
 	if err == nil {
 		c.channels = append(c.channels, channels...)
+		c.nsub += len(channels)
 	}
 	return err
 }
@@ -74,6 +77,7 @@ func (c *PubSub) PSubscribe(patterns ...string) error {
 	err := c.subscribe("PSUBSCRIBE", patterns...)
 	if err == nil {
 		c.patterns = append(c.patterns, patterns...)
+		c.nsub += len(patterns)
 	}
 	return err
 }
@@ -113,8 +117,12 @@ func (c *PubSub) PUnsubscribe(patterns ...string) error {
 	return err
 }
 
+func (c *PubSub) Close() error {
+	return c.base.Close()
+}
+
 func (c *PubSub) Ping(payload string) error {
-	cn, _, err := c.conn()
+	cn, _, err := c.base.conn()
 	if err != nil {
 		return err
 	}
@@ -178,7 +186,7 @@ func (p *Pong) String() string {
 	return "Pong"
 }
 
-func newMessage(reply []interface{}) (interface{}, error) {
+func (c *PubSub) newMessage(reply []interface{}) (interface{}, error) {
 	switch kind := reply[0].(string); kind {
 	case "subscribe", "unsubscribe", "psubscribe", "punsubscribe":
 		return &Subscription{
@@ -210,7 +218,11 @@ func newMessage(reply []interface{}) (interface{}, error) {
 // is not received in time. This is low-level API and most clients
 // should use ReceiveMessage.
 func (c *PubSub) ReceiveTimeout(timeout time.Duration) (interface{}, error) {
-	cn, _, err := c.conn()
+	if c.nsub == 0 {
+		c.resubscribe()
+	}
+
+	cn, _, err := c.base.conn()
 	if err != nil {
 		return nil, err
 	}
@@ -222,7 +234,8 @@ func (c *PubSub) ReceiveTimeout(timeout time.Duration) (interface{}, error) {
 	if err != nil {
 		return nil, err
 	}
-	return newMessage(cmd.Val())
+
+	return c.newMessage(cmd.Val())
 }
 
 // Receive returns a message as a Subscription, Message, PMessage,
@@ -232,22 +245,6 @@ func (c *PubSub) Receive() (interface{}, error) {
 	return c.ReceiveTimeout(0)
 }
 
-func (c *PubSub) reconnect(reason error) {
-	// Close current connection.
-	c.connPool.(*stickyConnPool).Reset(reason)
-
-	if len(c.channels) > 0 {
-		if err := c.Subscribe(c.channels...); err != nil {
-			Logger.Printf("Subscribe failed: %s", err)
-		}
-	}
-	if len(c.patterns) > 0 {
-		if err := c.PSubscribe(c.patterns...); err != nil {
-			Logger.Printf("PSubscribe failed: %s", err)
-		}
-	}
-}
-
 // ReceiveMessage returns a message or error. It automatically
 // reconnects to Redis in case of network errors.
 func (c *PubSub) ReceiveMessage() (*Message, error) {
@@ -259,10 +256,8 @@ func (c *PubSub) ReceiveMessage() (*Message, error) {
 				return nil, err
 			}
 
-			goodConn := errNum == 0
 			errNum++
-
-			if goodConn {
+			if errNum < 3 {
 				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
 					err := c.Ping("")
 					if err == nil {
@@ -270,16 +265,16 @@ func (c *PubSub) ReceiveMessage() (*Message, error) {
 					}
 					Logger.Printf("PubSub.Ping failed: %s", err)
 				}
-			}
-
-			if errNum > 2 {
+			} else {
+				// 3 consequent errors - connection is bad
+				// and/or Redis Server is down.
+				// Sleep to not exceed max number of open connections.
 				time.Sleep(time.Second)
 			}
-			c.reconnect(err)
 			continue
 		}
 
-		// Reset error number.
+		// Reset error number, because we received a message.
 		errNum = 0
 
 		switch msg := msgi.(type) {
@@ -300,3 +295,22 @@ func (c *PubSub) ReceiveMessage() (*Message, error) {
 		}
 	}
 }
+
+func (c *PubSub) putConn(cn *conn, err error) {
+	if !c.base.putConn(cn, err) {
+		c.nsub = 0
+	}
+}
+
+func (c *PubSub) resubscribe() {
+	if len(c.channels) > 0 {
+		if err := c.Subscribe(c.channels...); err != nil {
+			Logger.Printf("Subscribe failed: %s", err)
+		}
+	}
+	if len(c.patterns) > 0 {
+		if err := c.PSubscribe(c.patterns...); err != nil {
+			Logger.Printf("PSubscribe failed: %s", err)
+		}
+	}
+}