redis/pubsub.go

392 lines
8.1 KiB
Go
Raw Normal View History

2012-07-25 17:00:50 +04:00
package redis
import (
"fmt"
2015-09-06 13:50:16 +03:00
"net"
"sync"
2014-05-11 11:42:40 +04:00
"time"
2017-02-18 17:42:34 +03:00
"github.com/go-redis/redis/internal"
"github.com/go-redis/redis/internal/pool"
2012-07-25 17:00:50 +04:00
)
2015-05-23 14:15:05 +03:00
// PubSub implements Pub/Sub commands as described in
2015-09-12 09:36:03 +03:00
// http://redis.io/topics/pubsub. It's NOT safe for concurrent use by
// multiple goroutines.
2014-05-11 11:42:40 +04:00
type PubSub struct {
base baseClient
mu sync.Mutex
cn *pool.Conn
closed bool
subMu sync.Mutex
2015-09-06 13:50:16 +03:00
channels []string
patterns []string
2017-04-24 12:43:15 +03:00
cmd *Cmd
2016-09-29 15:07:04 +03:00
}
2017-04-24 12:43:15 +03:00
func (c *PubSub) conn() (*pool.Conn, bool, error) {
c.mu.Lock()
defer c.mu.Unlock()
if c.closed {
return nil, false, pool.ErrClosed
}
if c.cn != nil {
return c.cn, false, nil
}
cn, err := c.base.connPool.NewConn()
2016-09-29 15:07:04 +03:00
if err != nil {
2017-04-24 12:43:15 +03:00
return nil, false, err
2016-09-29 15:07:04 +03:00
}
2017-04-24 12:43:15 +03:00
if !cn.Inited {
if err := c.base.initConn(cn); err != nil {
2017-04-24 14:00:28 +03:00
_ = c.base.connPool.CloseConn(cn)
2017-04-24 12:43:15 +03:00
return nil, false, err
}
2016-09-29 15:07:04 +03:00
}
2017-04-24 12:43:15 +03:00
if err := c.resubscribe(cn); err != nil {
return nil, false, err
}
c.cn = cn
return cn, true, nil
}
2017-04-24 12:43:15 +03:00
func (c *PubSub) resubscribe(cn *pool.Conn) error {
c.subMu.Lock()
2017-04-24 12:43:15 +03:00
defer c.subMu.Unlock()
var firstErr error
2017-04-24 12:43:15 +03:00
if len(c.channels) > 0 {
if err := c._subscribe(cn, "subscribe", c.channels...); err != nil && firstErr == nil {
firstErr = err
}
}
2017-04-24 12:43:15 +03:00
if len(c.patterns) > 0 {
if err := c._subscribe(cn, "psubscribe", c.patterns...); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
2017-04-24 12:43:15 +03:00
func (c *PubSub) putConn(cn *pool.Conn, err error) {
if !internal.IsBadConn(err, true) {
return
}
c.mu.Lock()
if c.cn == cn {
_ = c.closeConn()
}
c.mu.Unlock()
}
func (c *PubSub) closeConn() error {
err := c.base.connPool.CloseConn(c.cn)
c.cn = nil
return err
}
func (c *PubSub) Close() error {
c.mu.Lock()
defer c.mu.Unlock()
if c.closed {
2017-04-24 12:43:15 +03:00
return pool.ErrClosed
}
2017-04-24 12:43:15 +03:00
c.closed = true
if c.cn != nil {
2017-04-24 12:43:15 +03:00
return c.closeConn()
}
2017-04-24 12:43:15 +03:00
return nil
}
2017-04-24 12:43:15 +03:00
func (c *PubSub) subscribe(redisCmd string, channels ...string) error {
cn, isNew, err := c.conn()
if err != nil {
2017-04-24 12:43:15 +03:00
return err
}
2017-04-24 12:43:15 +03:00
if isNew {
return nil
2017-04-18 13:12:38 +03:00
}
2017-04-24 12:43:15 +03:00
err = c._subscribe(cn, redisCmd, channels...)
c.putConn(cn, err)
return err
2012-07-25 17:00:50 +04:00
}
2017-04-24 12:43:15 +03:00
func (c *PubSub) _subscribe(cn *pool.Conn, redisCmd string, channels ...string) error {
2015-09-06 13:50:16 +03:00
args := make([]interface{}, 1+len(channels))
args[0] = redisCmd
2015-09-06 13:50:16 +03:00
for i, channel := range channels {
args[1+i] = channel
}
cmd := NewSliceCmd(args...)
cn.SetWriteTimeout(c.base.opt.WriteTimeout)
2017-04-24 12:43:15 +03:00
return writeCmd(cn, cmd)
2015-09-06 13:50:16 +03:00
}
// Subscribes the client to the specified channels.
func (c *PubSub) Subscribe(channels ...string) error {
c.subMu.Lock()
2017-04-11 16:53:55 +03:00
c.channels = appendIfNotExists(c.channels, channels...)
c.subMu.Unlock()
return c.subscribe("subscribe", channels...)
2015-09-06 13:50:16 +03:00
}
// Subscribes the client to the given patterns.
func (c *PubSub) PSubscribe(patterns ...string) error {
c.subMu.Lock()
2017-04-11 16:53:55 +03:00
c.patterns = appendIfNotExists(c.patterns, patterns...)
c.subMu.Unlock()
return c.subscribe("psubscribe", patterns...)
2015-09-06 13:50:16 +03:00
}
// Unsubscribes the client from the given channels, or from all of
// them if none is given.
func (c *PubSub) Unsubscribe(channels ...string) error {
c.subMu.Lock()
2017-04-11 16:53:55 +03:00
c.channels = remove(c.channels, channels...)
c.subMu.Unlock()
return c.subscribe("unsubscribe", channels...)
2015-09-06 13:50:16 +03:00
}
// Unsubscribes the client from the given patterns, or from all of
// them if none is given.
func (c *PubSub) PUnsubscribe(patterns ...string) error {
c.subMu.Lock()
2017-04-11 16:53:55 +03:00
c.patterns = remove(c.patterns, patterns...)
c.subMu.Unlock()
return c.subscribe("punsubscribe", patterns...)
2015-09-06 13:50:16 +03:00
}
2017-02-23 16:29:38 +03:00
func (c *PubSub) Ping(payload ...string) error {
args := []interface{}{"ping"}
2017-02-23 16:29:38 +03:00
if len(payload) == 1 {
args = append(args, payload[0])
2015-07-11 13:12:47 +03:00
}
cmd := NewCmd(args...)
2017-04-24 12:43:15 +03:00
cn, _, err := c.conn()
if err != nil {
return err
}
cn.SetWriteTimeout(c.base.opt.WriteTimeout)
err = writeCmd(cn, cmd)
c.putConn(cn, err)
return err
2015-07-11 13:12:47 +03:00
}
2015-07-11 13:42:44 +03:00
// Message received after a successful subscription to channel.
type Subscription struct {
// Can be "subscribe", "unsubscribe", "psubscribe" or "punsubscribe".
Kind string
// Channel name we have subscribed to.
Channel string
// Number of channels we are currently subscribed to.
Count int
}
func (m *Subscription) String() string {
return fmt.Sprintf("%s: %s", m.Kind, m.Channel)
}
2015-05-23 18:17:45 +03:00
// Message received as result of a PUBLISH command issued by another client.
2012-07-25 17:00:50 +04:00
type Message struct {
2014-05-11 11:42:40 +04:00
Channel string
2015-09-06 13:50:16 +03:00
Pattern string
2014-05-11 11:42:40 +04:00
Payload string
}
2012-07-25 17:00:50 +04:00
2014-05-11 18:11:55 +04:00
func (m *Message) String() string {
return fmt.Sprintf("Message<%s: %s>", m.Channel, m.Payload)
}
2015-07-11 13:12:47 +03:00
// Pong received as result of a PING command issued by another client.
type Pong struct {
Payload string
}
func (p *Pong) String() string {
if p.Payload != "" {
return fmt.Sprintf("Pong<%s>", p.Payload)
}
return "Pong"
}
func (c *PubSub) newMessage(reply interface{}) (interface{}, error) {
switch reply := reply.(type) {
case string:
2015-07-11 13:12:47 +03:00
return &Pong{
Payload: reply,
2015-07-11 13:12:47 +03:00
}, nil
case []interface{}:
switch kind := reply[0].(string); kind {
case "subscribe", "unsubscribe", "psubscribe", "punsubscribe":
return &Subscription{
Kind: kind,
Channel: reply[1].(string),
Count: int(reply[2].(int64)),
}, nil
case "message":
return &Message{
Channel: reply[1].(string),
Payload: reply[2].(string),
}, nil
case "pmessage":
return &Message{
Pattern: reply[1].(string),
Channel: reply[2].(string),
Payload: reply[3].(string),
}, nil
case "pong":
return &Pong{
Payload: reply[1].(string),
}, nil
default:
return nil, fmt.Errorf("redis: unsupported pubsub message: %q", kind)
}
2015-07-11 13:12:47 +03:00
default:
return nil, fmt.Errorf("redis: unsupported pubsub message: %#v", reply)
2015-07-11 13:12:47 +03:00
}
}
2015-07-11 13:42:44 +03:00
// ReceiveTimeout acts like Receive but returns an error if message
2015-09-06 13:50:16 +03:00
// is not received in time. This is low-level API and most clients
// should use ReceiveMessage.
2015-07-11 13:12:47 +03:00
func (c *PubSub) ReceiveTimeout(timeout time.Duration) (interface{}, error) {
if c.cmd == nil {
c.cmd = NewCmd()
}
2017-04-24 12:43:15 +03:00
cn, _, err := c.conn()
2015-07-11 13:12:47 +03:00
if err != nil {
return nil, err
2014-05-11 11:42:40 +04:00
}
cn.SetReadTimeout(timeout)
err = c.cmd.readReply(cn)
c.putConn(cn, err)
if err != nil {
2015-07-11 13:12:47 +03:00
return nil, err
}
return c.newMessage(c.cmd.Val())
2014-05-11 11:42:40 +04:00
}
2012-07-25 17:00:50 +04:00
2016-04-09 11:45:56 +03:00
// Receive returns a message as a Subscription, Message, Pong or error.
// See PubSub example for details. This is low-level API and most clients
// should use ReceiveMessage.
2015-09-06 13:50:16 +03:00
func (c *PubSub) Receive() (interface{}, error) {
return c.ReceiveTimeout(0)
}
2014-05-11 11:42:40 +04:00
// ReceiveMessage returns a Message or error ignoring Subscription or Pong
// messages. It automatically reconnects to Redis Server and resubscribes
// to channels in case of network errors.
2015-09-06 13:50:16 +03:00
func (c *PubSub) ReceiveMessage() (*Message, error) {
2016-04-09 14:52:01 +03:00
return c.receiveMessage(5 * time.Second)
}
func (c *PubSub) receiveMessage(timeout time.Duration) (*Message, error) {
var errNum uint
2015-09-06 13:50:16 +03:00
for {
2016-04-09 14:52:01 +03:00
msgi, err := c.ReceiveTimeout(timeout)
2015-09-06 13:50:16 +03:00
if err != nil {
if !internal.IsNetworkError(err) {
2015-12-02 16:40:44 +03:00
return nil, err
}
errNum++
if errNum < 3 {
2015-12-02 16:40:44 +03:00
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
2017-02-23 16:29:38 +03:00
err := c.Ping()
if err != nil {
2016-04-09 14:52:01 +03:00
internal.Logf("PubSub.Ping failed: %s", err)
2015-12-02 16:40:44 +03:00
}
2015-09-06 13:50:16 +03:00
}
} else {
// 3 consequent errors - connection is broken or
// Redis Server is down.
// Sleep to not exceed max number of open connections.
2015-12-02 16:40:44 +03:00
time.Sleep(time.Second)
2015-09-06 13:50:16 +03:00
}
2015-12-02 16:40:44 +03:00
continue
2015-09-06 13:50:16 +03:00
}
// Reset error number, because we received a message.
2015-12-02 16:40:44 +03:00
errNum = 0
2015-09-06 13:50:16 +03:00
switch msg := msgi.(type) {
case *Subscription:
// Ignore.
case *Pong:
// Ignore.
case *Message:
return msg, nil
default:
return nil, fmt.Errorf("redis: unknown message: %T", msgi)
}
}
}
2017-04-11 16:18:35 +03:00
// Channel returns a channel for concurrently receiving messages.
// The channel is closed with PubSub.
func (c *PubSub) Channel() <-chan *Message {
ch := make(chan *Message, 100)
go func() {
for {
msg, err := c.ReceiveMessage()
if err != nil {
if err == pool.ErrClosed {
break
}
continue
}
ch <- msg
}
close(ch)
}()
return ch
}
func appendIfNotExists(ss []string, es ...string) []string {
loop:
2016-09-29 15:07:04 +03:00
for _, e := range es {
for _, s := range ss {
2016-09-29 15:07:04 +03:00
if s == e {
continue loop
2016-09-29 15:07:04 +03:00
}
}
ss = append(ss, e)
2016-09-29 15:07:04 +03:00
}
return ss
}
func remove(ss []string, es ...string) []string {
if len(es) == 0 {
return ss[:0]
}
2016-09-29 15:07:04 +03:00
for _, e := range es {
for i, s := range ss {
2016-09-29 15:07:04 +03:00
if s == e {
ss = append(ss[:i], ss[i+1:]...)
break
2016-09-29 15:07:04 +03:00
}
}
}
return ss
}