3 Star 0 Fork 0

Gitee 极速下载/go-redis-redissource

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
此仓库是为了提升国内下载速度的镜像仓库,每日同步一次。 原始仓库: https://github.com/go-redis/redis
克隆/下载
pubsub.go 8.17 KB
一键复制 编辑 原始数据 按行查看 历史
Vladimir Mihailenco 提交于 2017-04-24 14:00 . Fix conn management in PubSub
package redis
import (
"fmt"
"net"
"sync"
"time"
"github.com/go-redis/redis/internal"
"github.com/go-redis/redis/internal/pool"
)
// PubSub implements Pub/Sub commands as described in
// http://redis.io/topics/pubsub. It's NOT safe for concurrent use by
// multiple goroutines.
type PubSub struct {
base baseClient
mu sync.Mutex
cn *pool.Conn
closed bool
subMu sync.Mutex
channels []string
patterns []string
cmd *Cmd
}
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()
if err != nil {
return nil, false, err
}
if !cn.Inited {
if err := c.base.initConn(cn); err != nil {
_ = c.base.connPool.CloseConn(cn)
return nil, false, err
}
}
if err := c.resubscribe(cn); err != nil {
_ = c.base.connPool.CloseConn(cn)
return nil, false, err
}
c.cn = cn
return cn, true, nil
}
func (c *PubSub) resubscribe(cn *pool.Conn) error {
c.subMu.Lock()
defer c.subMu.Unlock()
var firstErr error
if len(c.channels) > 0 {
if err := c._subscribe(cn, "subscribe", c.channels...); err != nil && firstErr == nil {
firstErr = err
}
}
if len(c.patterns) > 0 {
if err := c._subscribe(cn, "psubscribe", c.patterns...); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
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 {
return pool.ErrClosed
}
c.closed = true
if c.cn != nil {
return c.closeConn()
}
return nil
}
func (c *PubSub) subscribe(redisCmd string, channels ...string) error {
cn, isNew, err := c.conn()
if err != nil {
return err
}
if isNew {
return nil
}
err = c._subscribe(cn, redisCmd, channels...)
c.putConn(cn, err)
return err
}
func (c *PubSub) _subscribe(cn *pool.Conn, redisCmd string, channels ...string) error {
args := make([]interface{}, 1+len(channels))
args[0] = redisCmd
for i, channel := range channels {
args[1+i] = channel
}
cmd := NewSliceCmd(args...)
cn.SetWriteTimeout(c.base.opt.WriteTimeout)
return writeCmd(cn, cmd)
}
// Subscribes the client to the specified channels.
func (c *PubSub) Subscribe(channels ...string) error {
c.subMu.Lock()
c.channels = appendIfNotExists(c.channels, channels...)
c.subMu.Unlock()
return c.subscribe("subscribe", channels...)
}
// Subscribes the client to the given patterns.
func (c *PubSub) PSubscribe(patterns ...string) error {
c.subMu.Lock()
c.patterns = appendIfNotExists(c.patterns, patterns...)
c.subMu.Unlock()
return c.subscribe("psubscribe", patterns...)
}
// 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()
c.channels = remove(c.channels, channels...)
c.subMu.Unlock()
return c.subscribe("unsubscribe", channels...)
}
// 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()
c.patterns = remove(c.patterns, patterns...)
c.subMu.Unlock()
return c.subscribe("punsubscribe", patterns...)
}
func (c *PubSub) Ping(payload ...string) error {
args := []interface{}{"ping"}
if len(payload) == 1 {
args = append(args, payload[0])
}
cmd := NewCmd(args...)
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
}
// 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)
}
// Message received as result of a PUBLISH command issued by another client.
type Message struct {
Channel string
Pattern string
Payload string
}
func (m *Message) String() string {
return fmt.Sprintf("Message<%s: %s>", m.Channel, m.Payload)
}
// 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:
return &Pong{
Payload: reply,
}, 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)
}
default:
return nil, fmt.Errorf("redis: unsupported pubsub message: %#v", reply)
}
}
// ReceiveTimeout acts like Receive but returns an error if message
// 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) {
if c.cmd == nil {
c.cmd = NewCmd()
}
cn, _, err := c.conn()
if err != nil {
return nil, err
}
cn.SetReadTimeout(timeout)
err = c.cmd.readReply(cn)
c.putConn(cn, err)
if err != nil {
return nil, err
}
return c.newMessage(c.cmd.Val())
}
// 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.
func (c *PubSub) Receive() (interface{}, error) {
return c.ReceiveTimeout(0)
}
// 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.
func (c *PubSub) ReceiveMessage() (*Message, error) {
return c.receiveMessage(5 * time.Second)
}
func (c *PubSub) receiveMessage(timeout time.Duration) (*Message, error) {
var errNum uint
for {
msgi, err := c.ReceiveTimeout(timeout)
if err != nil {
if !internal.IsNetworkError(err) {
return nil, err
}
errNum++
if errNum < 3 {
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
err := c.Ping()
if err != nil {
internal.Logf("PubSub.Ping failed: %s", err)
}
}
} else {
// 3 consequent errors - connection is broken or
// Redis Server is down.
// Sleep to not exceed max number of open connections.
time.Sleep(time.Second)
}
continue
}
// Reset error number, because we received a message.
errNum = 0
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)
}
}
}
// 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:
for _, e := range es {
for _, s := range ss {
if s == e {
continue loop
}
}
ss = append(ss, e)
}
return ss
}
func remove(ss []string, es ...string) []string {
if len(es) == 0 {
return ss[:0]
}
for _, e := range es {
for i, s := range ss {
if s == e {
ss = append(ss[:i], ss[i+1:]...)
break
}
}
}
return ss
}
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/mirrors/redissource.git
git@gitee.com:mirrors/redissource.git
mirrors
redissource
go-redis-redissource
v6.3.1

搜索帮助