redigo中PubSub的一点小坑

最近在用 golang 做一些 redis 相关的操作，选用了 redigo 这个第三方库。然后在使用 Pub/Sub 的时候，却发现了一个小坑……

Redis Client

首先，我们来初始化一个带连接池的 Redis Client：

import (
	"github.com/gomodule/redigo/redis"
)
 
type RedisClient struct {
	pool *redis.Pool
}
 
func NewRedisClient(addr string, db int, passwd string) *RedisClient {
	pool := &redis.Pool{
		MaxIdle:     10,
		IdleTimeout: 300 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", addr, redis.DialPassword(passwd), redis.DialDatabase(db))
			if err != nil {
				return nil, err
			}
			return c, nil
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := c.Do("PING")
			return err
		},
	}
	log.Printf("new redis pool at %s", addr)
	client := &RedisClient{
		pool: pool,
	}
	return client
}

Publish

然后我们可以简单的实现一个 publish 方法：

func (r *RedisClient) Publish(channel, message string) (int, error) {
	c := r.pool.Get()
	defer c.Close()
	n, err := redis.Int(c.Do("PUBLISH", channel, message))
	if err != nil {
		return 0, fmt.Errorf("redis publish %s %s, err: %v", channel, message, err)
	}
	return n, nil
}

Subscribe

接下来就是一个稍微复杂点的带有心跳的 subscribe 方法：

func (r *RedisClient) Subscribe(ctx context.Context, consume ConsumeFunc, channel ...string) error {
	psc := redis.PubSubConn{Conn: r.pool.Get()}
	defer psc.Close()
	log.Printf("redis pubsub subscribe channel: %v", channel)
	if err := psc.Subscribe(redis.Args{}.AddFlat(channel)...); err != nil {
		return err
	}
	done := make(chan error, 1)
	// start a new goroutine to receive message
	go func() {
		for {
			switch msg := psc.Receive().(type) {
			case error:
				done <- fmt.Errorf("redis pubsub receive err: %v", msg)
				return
			case redis.Message:
				if err := consume(msg); err != nil {
					done <- err
					return
				}
			case redis.Subscription:
				if msg.Count == 0 {
					// all channels are unsubscribed
					done <- nil
					return
				}
			}
		}
	}()
 
	// health check
	tick := time.NewTicker(time.Minute)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			if err := psc.Unsubscribe(); err != nil {
				return fmt.Errorf("redis pubsub unsubscribe err: %v", err)
			}
			return nil
		case err := <-done:
			return err
		case <-tick.C:
			if err := psc.Ping(""); err != nil {
				return err
			}
		}
	}
 
	return nil
}

最后，我们写一个简单地 main 函数来调用 publish & subscribe：

func main(){
	r := NewRedisClient("172.16.8.128:6379", 0, "")
	consume := func(msg redis.Message) error {
		log.Printf("recv msg: %s", msg.Data)
		return nil
	}
	for i := 0; i <10; i++{
		log.Printf("-------------- %d -----------------", i)
		ctx, cancel := context.WithCancel(context.Background())
		go func(){
			if err := r.Subscribe(ctx, consume, "channel"); err != nil {
				log.Println("subscribe err: %v", err)
			}
		}()
		time.Sleep(time.Second)
		_, err:= r.Publish("channel", "hello, " + strconv.Itoa(i))
		if err != nil {
			log.Fatal(err)
		}
		time.Sleep(time.Second)
		cancel()
	}
	forever := make(chan struct{})
	<-forever
}

坑

咋一看之下，好像并没有什么异常？然而，如果我们这时候去看 redis 的 tcp 连接，就可以发现一些猫腻：

$sudo netstat -antp | grep redis
tcp        0      0 0.0.0.0:6379            0.0.0.0:*               LISTEN      940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55010        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55015        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55009        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55005        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55012        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55011        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55013        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55007        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55006        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:55014        ESTABLISHED 940/redis-server 0. 
tcp        0      0 172.16.8.128:6379       172.16.8.1:54972        ESTABLISHED 940/redis-server 0. 

竟然是每一次 subscribe 就新建了一个连接，而 connection pool 似乎没有什么作用。

更进一步地调试，我们发现在 defer psc.Close() 的时候就卡住了，也就是上面的 10 个 goroutine 其实并没有正常退出。

Concurrent

排查许久之后，终于定位到了问题！引用 redigo 的 说明 ：

也就是说，虽然一个连接可以在不同的 goroutine 并发调用 Receive() 和 Subscribe() （subscribe调用了send和flush） ，但是却不能再有其他并发操作 （比如 Close()） 。

其他相似的问题还可以参考 issue

Fix

知道了上面的原因之后，我们稍微修改一下 defer psc.Close() 的位置即可解决问题：

	// start a new goroutine to receive message
	go func() {
		// IMPORTANT!
		defer psc.Close()
		for {
			switch msg := psc.Receive().(type) {
			case error: