Rework pubsub interface. Add timeout support.

v2/example_test.go

@@ -115,20 +115,19 @@ func ExamplePubSub() {
 	pubsub, err := client.PubSubClient()
 	defer pubsub.Close()
 
-	ch, err := pubsub.Subscribe("mychannel")
+	err = pubsub.Subscribe("mychannel")
-	_ = err
 
-	subscribeMsg := <-ch
+	msg, err := pubsub.Receive(0)
-	fmt.Println(subscribeMsg.Err, subscribeMsg.Name)
+	fmt.Println(msg, err)
 
 	pub := client.Publish("mychannel", "hello")
 	_ = pub.Err()
 
-	msg := <-ch
+	msg, err = pubsub.Receive(0)
-	fmt.Println(msg.Err, msg.Message)
+	fmt.Println(msg, err)
 
-	// Output: <nil> subscribe
+	// Output: &{subscribe mychannel 1} <nil>
-	// <nil> hello
+	// &{mychannel hello} <nil>
 }
 
 func Get(client *redis.Client, key string) *redis.StringReq {

v2/pubsub.go

@@ -2,14 +2,11 @@ package redis
 
 import (
 	"fmt"
-	"sync"
+	"time"
 )
 
 type PubSubClient struct {
 	*baseClient
-
-	ch   chan *Message
-	once sync.Once
 }
 
 func (c *Client) PubSubClient() (*PubSubClient, error) {
@@ -18,8 +15,6 @@ func (c *Client) PubSubClient() (*PubSubClient, error) {
 			opt:      c.opt,
 			connPool: newSingleConnPool(c.connPool, nil, false),
 		},
-
-		ch: make(chan *Message),
 	}, nil
 }
 
@@ -30,90 +25,88 @@ func (c *Client) Publish(channel, message string) *IntReq {
 }
 
 type Message struct {
-	Name, Channel, ChannelPattern, Message string
+	Channel string
-	Number                                 int64
+	Payload string
-
-	Err error
 }
 
-func (c *PubSubClient) consumeMessages(cn *conn) {
+type PMessage struct {
-	req := NewIfaceSliceReq()
+	Channel string
-
+	Pattern string
-	for {
+	Payload string
-		msg := &Message{}
-
-		replyIface, err := req.ParseReply(cn.Rd)
-		if err != nil {
-			msg.Err = err
-			c.ch <- msg
-			return
-		}
-		reply, ok := replyIface.([]interface{})
-		if !ok {
-			msg.Err = fmt.Errorf("redis: unexpected reply type %T", replyIface)
-			c.ch <- msg
-			return
-		}
-
-		msgName := reply[0].(string)
-		switch msgName {
-		case "subscribe", "unsubscribe", "psubscribe", "punsubscribe":
-			msg.Name = msgName
-			msg.Channel = reply[1].(string)
-			msg.Number = reply[2].(int64)
-		case "message":
-			msg.Name = msgName
-			msg.Channel = reply[1].(string)
-			msg.Message = reply[2].(string)
-		case "pmessage":
-			msg.Name = msgName
-			msg.ChannelPattern = reply[1].(string)
-			msg.Channel = reply[2].(string)
-			msg.Message = reply[3].(string)
-		default:
-			msg.Err = fmt.Errorf("redis: unsupported message name: %q", msgName)
-		}
-		c.ch <- msg
-	}
 }
 
-func (c *PubSubClient) subscribe(cmd string, channels ...string) (chan *Message, error) {
+type Subscription struct {
-	args := append([]string{cmd}, channels...)
+	Kind    string
-	req := NewIfaceSliceReq(args...)
+	Channel string
+	Count   int
+}
 
+func (c *PubSubClient) Receive(timeout time.Duration) (interface{}, error) {
 	cn, err := c.conn()
 	if err != nil {
 		return nil, err
 	}
+	cn.readTimeout = timeout
 
-	if err := c.writeReq(cn, req); err != nil {
+	replyIface, err := NewIfaceSliceReq().ParseReply(cn.Rd)
+	if err != nil {
 		return nil, err
 	}
+	reply, ok := replyIface.([]interface{})
+	if !ok {
+		return nil, fmt.Errorf("redis: unexpected reply type %T", replyIface)
+	}
 
-	c.once.Do(func() {
+	switch msgName := reply[0].(string); msgName {
-		go c.consumeMessages(cn)
+	case "subscribe", "unsubscribe", "psubscribe", "punsubscribe":
-	})
+		return &Subscription{
-
+			Kind:    msgName,
-	return c.ch, nil
+			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 &PMessage{
+			Pattern: reply[1].(string),
+			Channel: reply[2].(string),
+			Payload: reply[3].(string),
+		}, nil
+	default:
+		return nil, fmt.Errorf("redis: unsupported message name: %q", msgName)
+	}
 }
 
-func (c *PubSubClient) Subscribe(channels ...string) (chan *Message, error) {
+func (c *PubSubClient) subscribe(cmd string, channels ...string) error {
-	return c.subscribe("SUBSCRIBE", channels...)
-}
-
-func (c *PubSubClient) PSubscribe(patterns ...string) (chan *Message, error) {
-	return c.subscribe("PSUBSCRIBE", patterns...)
-}
-
-func (c *PubSubClient) unsubscribe(cmd string, channels ...string) error {
-	args := append([]string{cmd}, channels...)
-	req := NewIfaceSliceReq(args...)
-
 	cn, err := c.conn()
 	if err != nil {
 		return err
 	}
 
+	args := append([]string{cmd}, channels...)
+	req := NewIfaceSliceReq(args...)
+	return c.writeReq(cn, req)
+}
+
+func (c *PubSubClient) Subscribe(channels ...string) error {
+	return c.subscribe("SUBSCRIBE", channels...)
+}
+
+func (c *PubSubClient) PSubscribe(patterns ...string) error {
+	return c.subscribe("PSUBSCRIBE", patterns...)
+}
+
+func (c *PubSubClient) unsubscribe(cmd string, channels ...string) error {
+	cn, err := c.conn()
+	if err != nil {
+		return err
+	}
+
+	args := append([]string{cmd}, channels...)
+	req := NewIfaceSliceReq(args...)
 	return c.writeReq(cn, req)
 }
 

v2/redis_test.go

@@ -2268,46 +2268,45 @@ func (t *RedisTest) TestPatternPubSub(c *C) {
 		c.Assert(pubsub.Close(), IsNil)
 	}()
 
-	ch, err := pubsub.PSubscribe("mychannel*")
+	c.Assert(pubsub.PSubscribe("mychannel*"), IsNil)
-	c.Assert(err, IsNil)
-	c.Assert(ch, Not(IsNil))
 
 	pub := t.client.Publish("mychannel1", "hello")
 	c.Assert(pub.Err(), IsNil)
 	c.Assert(pub.Val(), Equals, int64(1))
 
-	err = pubsub.PUnsubscribe("mychannel*")
+	c.Assert(pubsub.PUnsubscribe("mychannel*"), IsNil)
+
+	{
+		msgi, err := pubsub.Receive(time.Second)
 		c.Assert(err, IsNil)
-
+		subscr := msgi.(*redis.Subscription)
-	select {
+		c.Assert(subscr.Kind, Equals, "psubscribe")
-	case msg := <-ch:
+		c.Assert(subscr.Channel, Equals, "mychannel*")
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(subscr.Count, Equals, 1)
-		c.Assert(msg.Name, Equals, "psubscribe")
-		c.Assert(msg.Channel, Equals, "mychannel*")
-		c.Assert(msg.Number, Equals, int64(1))
-	case <-time.After(time.Second):
-		c.Error("Channel is empty.")
 	}
 
-	select {
+	{
-	case msg := <-ch:
+		msgi, err := pubsub.Receive(time.Second)
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(err, IsNil)
-		c.Assert(msg.Name, Equals, "pmessage")
+		subscr := msgi.(*redis.PMessage)
-		c.Assert(msg.ChannelPattern, Equals, "mychannel*")
+		c.Assert(subscr.Channel, Equals, "mychannel1")
-		c.Assert(msg.Channel, Equals, "mychannel1")
+		c.Assert(subscr.Pattern, Equals, "mychannel*")
-		c.Assert(msg.Message, Equals, "hello")
+		c.Assert(subscr.Payload, Equals, "hello")
-	case <-time.After(time.Second):
-		c.Error("Channel is empty.")
 	}
 
-	select {
+	{
-	case msg := <-ch:
+		msgi, err := pubsub.Receive(time.Second)
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(err, IsNil)
-		c.Assert(msg.Name, Equals, "punsubscribe")
+		subscr := msgi.(*redis.Subscription)
-		c.Assert(msg.Channel, Equals, "mychannel*")
+		c.Assert(subscr.Kind, Equals, "punsubscribe")
-		c.Assert(msg.Number, Equals, int64(0))
+		c.Assert(subscr.Channel, Equals, "mychannel*")
-	case <-time.After(time.Second):
+		c.Assert(subscr.Count, Equals, 0)
-		c.Error("Channel is empty.")
+	}
+
+	{
+		msgi, err := pubsub.Receive(time.Second)
+		c.Assert(err.(net.Error).Timeout(), Equals, true)
+		c.Assert(msgi, IsNil)
 	}
 }
 
@@ -2318,13 +2317,7 @@ func (t *RedisTest) TestPubSub(c *C) {
 		c.Assert(pubsub.Close(), IsNil)
 	}()
 
-	ch, err := pubsub.Subscribe("mychannel")
+	c.Assert(pubsub.Subscribe("mychannel", "mychannel2"), IsNil)
-	c.Assert(err, IsNil)
-	c.Assert(ch, Not(IsNil))
-
-	ch2, err := pubsub.Subscribe("mychannel2")
-	c.Assert(err, IsNil)
-	c.Assert(ch2, Equals, ch)
 
 	pub := t.client.Publish("mychannel", "hello")
 	c.Assert(pub.Err(), IsNil)
@@ -2334,70 +2327,64 @@ func (t *RedisTest) TestPubSub(c *C) {
 	c.Assert(pub.Err(), IsNil)
 	c.Assert(pub.Val(), Equals, int64(1))
 
-	err = pubsub.Unsubscribe("mychannel")
+	c.Assert(pubsub.Unsubscribe("mychannel", "mychannel2"), IsNil)
+
+	{
+		msgi, err := pubsub.Receive(time.Second)
 		c.Assert(err, IsNil)
+		subscr := msgi.(*redis.Subscription)
+		c.Assert(subscr.Kind, Equals, "subscribe")
+		c.Assert(subscr.Channel, Equals, "mychannel")
+		c.Assert(subscr.Count, Equals, 1)
+	}
 
-	err = pubsub.Unsubscribe("mychannel2")
+	{
+		msgi, err := pubsub.Receive(time.Second)
 		c.Assert(err, IsNil)
-
+		subscr := msgi.(*redis.Subscription)
-	select {
+		c.Assert(subscr.Kind, Equals, "subscribe")
-	case msg := <-ch:
+		c.Assert(subscr.Channel, Equals, "mychannel2")
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(subscr.Count, Equals, 2)
-		c.Assert(msg.Name, Equals, "subscribe")
-		c.Assert(msg.Channel, Equals, "mychannel")
-		c.Assert(msg.Number, Equals, int64(1))
-	case <-time.After(time.Second):
-		c.Error("Channel is empty.")
 	}
 
-	select {
+	{
-	case msg := <-ch:
+		msgi, err := pubsub.Receive(time.Second)
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(err, IsNil)
-		c.Assert(msg.Name, Equals, "subscribe")
+		subscr := msgi.(*redis.Message)
+		c.Assert(subscr.Channel, Equals, "mychannel")
+		c.Assert(subscr.Payload, Equals, "hello")
+	}
+
+	{
+		msgi, err := pubsub.Receive(time.Second)
+		c.Assert(err, IsNil)
+		msg := msgi.(*redis.Message)
 		c.Assert(msg.Channel, Equals, "mychannel2")
-		c.Assert(msg.Number, Equals, int64(2))
+		c.Assert(msg.Payload, Equals, "hello2")
-	case <-time.After(time.Second):
-		c.Error("Channel is empty.")
 	}
 
-	select {
+	{
-	case msg := <-ch:
+		msgi, err := pubsub.Receive(time.Second)
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(err, IsNil)
-		c.Assert(msg.Name, Equals, "message")
+		subscr := msgi.(*redis.Subscription)
-		c.Assert(msg.Channel, Equals, "mychannel")
+		c.Assert(subscr.Kind, Equals, "unsubscribe")
-		c.Assert(msg.Message, Equals, "hello")
+		c.Assert(subscr.Channel, Equals, "mychannel")
-	case <-time.After(time.Second):
+		c.Assert(subscr.Count, Equals, 1)
-		c.Error("Channel is empty.")
 	}
 
-	select {
+	{
-	case msg := <-ch:
+		msgi, err := pubsub.Receive(time.Second)
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(err, IsNil)
-		c.Assert(msg.Name, Equals, "message")
+		subscr := msgi.(*redis.Subscription)
-		c.Assert(msg.Channel, Equals, "mychannel2")
+		c.Assert(subscr.Kind, Equals, "unsubscribe")
-		c.Assert(msg.Message, Equals, "hello2")
+		c.Assert(subscr.Channel, Equals, "mychannel2")
-	case <-time.After(time.Second):
+		c.Assert(subscr.Count, Equals, 0)
-		c.Error("Channel is empty.")
 	}
 
-	select {
+	{
-	case msg := <-ch:
+		msgi, err := pubsub.Receive(time.Second)
-		c.Assert(msg.Err, Equals, nil)
+		c.Assert(err.(net.Error).Timeout(), Equals, true)
-		c.Assert(msg.Name, Equals, "unsubscribe")
+		c.Assert(msgi, IsNil)
-		c.Assert(msg.Channel, Equals, "mychannel")
-		c.Assert(msg.Number, Equals, int64(1))
-	case <-time.After(time.Second):
-		c.Error("Channel is empty.")
-	}
-
-	select {
-	case msg := <-ch:
-		c.Assert(msg.Err, Equals, nil)
-		c.Assert(msg.Name, Equals, "unsubscribe")
-		c.Assert(msg.Channel, Equals, "mychannel2")
-		c.Assert(msg.Number, Equals, int64(0))
-	case <-time.After(time.Second):
-		c.Error("Channel is empty.")
 	}
 }