feat: add the new WithOnSubscriptionHook for the Receive method (#846)

* feat: add the new WithOnSubscriptionHook for the Receive method

Signed-off-by: Rueian <[email protected]>

* feat: apply betteraligment

Signed-off-by: Rueian <[email protected]>

* feat: add the new WithOnSubscriptionHook for the Receive method

Signed-off-by: Rueian <[email protected]>

* feat: add the new WithOnSubscriptionHook for the Receive method

Signed-off-by: Rueian <[email protected]>

---------

Signed-off-by: Rueian <[email protected]>

Author: rueian

README.md

@@ -277,7 +277,12 @@ To receive messages from channels, `client.Receive()` should be used. It support
 
 ```golang
 err = client.Receive(context.Background(), client.B().Subscribe().Channel("ch1", "ch2").Build(), func(msg rueidis.PubSubMessage) {
-    // Handle the message. Note that if you want to call another `client.Do()` here, you need to do it in another goroutine or the `client` will be blocked.
+    // Handle the message. If you need to perform heavy processing or issue
+    // additional commands, do that in a separate goroutine to avoid
+    // blocking the pipeline, e.g.:
+    //   go func() {
+    //       // long work or client.Do(...)
+    //   }()
 })
 ```
 
@@ -294,6 +299,28 @@ While the `client.Receive()` call is blocking, the `Client` is still able to acc
 and they are sharing the same TCP connection. If your message handler may take some time to complete, it is recommended
 to use the `client.Receive()` inside a `client.Dedicated()` for not blocking other concurrent requests.
 
+#### Subscription confirmations
+
+Use `rueidis.WithOnSubscriptionHook` when you need to observe subscribe / unsubscribe confirmations that the server sends during the lifetime of a `client.Receive()`.
+
+The hook can be triggered multiple times because the `client.Receive()` may automatically reconnect and resubscribe.
+
+```go
+ctx := rueidis.WithOnSubscriptionHook(context.Background(), func(s rueidis.PubSubSubscription) {
+    // This hook runs in the pipeline goroutine. If you need to perform
+    // heavy work or invoke additional commands, do it in another
+    // goroutine to avoid blocking the pipeline, for example:
+    //   go func() {
+    //       // long work or client.Do(...)
+    //   }()
+    fmt.Printf("%s %s (count %d)\n", s.Kind, s.Channel, s.Count)
+})
+
+err := client.Receive(ctx, client.B().Subscribe().Channel("news").Build(), func(m rueidis.PubSubMessage) {
+    // ...
+})
+```
+
 ### Alternative PubSub Hooks
 
 The `client.Receive()` requires users to provide a subscription command in advance.
@@ -305,7 +332,12 @@ defer cancel()
 
 wait := c.SetPubSubHooks(rueidis.PubSubHooks{
 	OnMessage: func(m rueidis.PubSubMessage) {
-		// Handle the message. Note that if you want to call another `c.Do()` here, you need to do it in another goroutine or the `c` will be blocked.
+		// Handle the message. If you need to perform heavy processing or issue
+		// additional commands, do that in a separate goroutine to avoid
+		// blocking the pipeline, e.g.:
+		//   go func() {
+		//       // long work or client.Do(...)
+		//   }()
 	}
 })
 c.Do(ctx, c.B().Subscribe().Channel("ch").Build())

pipe.go

@@ -720,26 +720,45 @@ func (p *pipe) handlePush(values []RedisMessage) (reply bool, unsubscribe bool)
 			p.pshks.Load().(*pshks).hooks.OnMessage(m)
 		}
 	case "unsubscribe":
-		p.nsubs.Unsubscribe(values[1].string())
 		if len(values) >= 3 {
-			p.pshks.Load().(*pshks).hooks.OnSubscription(PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen})
+			s := PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen}
+			p.nsubs.Unsubscribe(s)
+			p.pshks.Load().(*pshks).hooks.OnSubscription(s)
 		}
 		return true, true
 	case "punsubscribe":
-		p.psubs.Unsubscribe(values[1].string())
 		if len(values) >= 3 {
-			p.pshks.Load().(*pshks).hooks.OnSubscription(PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen})
+			s := PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen}
+			p.psubs.Unsubscribe(s)
+			p.pshks.Load().(*pshks).hooks.OnSubscription(s)
 		}
 		return true, true
 	case "sunsubscribe":
-		p.ssubs.Unsubscribe(values[1].string())
 		if len(values) >= 3 {
-			p.pshks.Load().(*pshks).hooks.OnSubscription(PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen})
+			s := PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen}
+			p.ssubs.Unsubscribe(s)
+			p.pshks.Load().(*pshks).hooks.OnSubscription(s)
 		}
 		return true, true
-	case "subscribe", "psubscribe", "ssubscribe":
+	case "subscribe":
 		if len(values) >= 3 {
-			p.pshks.Load().(*pshks).hooks.OnSubscription(PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen})
+			s := PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen}
+			p.nsubs.Confirm(s)
+			p.pshks.Load().(*pshks).hooks.OnSubscription(s)
+		}
+		return true, false
+	case "psubscribe":
+		if len(values) >= 3 {
+			s := PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen}
+			p.psubs.Confirm(s)
+			p.pshks.Load().(*pshks).hooks.OnSubscription(s)
+		}
+		return true, false
+	case "ssubscribe":
+		if len(values) >= 3 {
+			s := PubSubSubscription{Kind: values[0].string(), Channel: values[1].string(), Count: values[2].intlen}
+			p.ssubs.Confirm(s)
+			p.pshks.Load().(*pshks).hooks.OnSubscription(s)
 		}
 		return true, false
 	}
@@ -762,6 +781,25 @@ func (p *pipe) _r2pipe(ctx context.Context) (r2p *pipe) {
 	return r2p
 }
 
+type recvCtxKey int
+
+const hookKey recvCtxKey = 0
+
+// WithOnSubscriptionHook attaches a subscription confirmation hook to the provided
+// context and returns a new context for the Receive method.
+//
+// The hook is invoked each time the server sends a subscribe or
+// unsubscribe confirmation, allowing callers to observe the state of a Pub/Sub
+// subscription during the lifetime of a Receive invocation.
+//
+// The hook may be called multiple times because the client can resubscribe after a
+// reconnection. Therefore, the hook implementation must be safe to run more than once.
+// Also, there should not be any blocking operations or another `client.Do()` in the hook
+// since it runs in the same goroutine as the pipeline. Otherwise, the pipeline will be blocked.
+func WithOnSubscriptionHook(ctx context.Context, hook func(PubSubSubscription)) context.Context {
+	return context.WithValue(ctx, hookKey, hook)
+}
+
 func (p *pipe) Receive(ctx context.Context, subscribe Completed, fn func(message PubSubMessage)) error {
 	if p.nsubs == nil || p.psubs == nil || p.ssubs == nil {
 		return p.Error()
@@ -787,7 +825,11 @@ func (p *pipe) Receive(ctx context.Context, subscribe Completed, fn func(message
 		panic(wrongreceive)
 	}
 
-	if ch, cancel := sb.Subscribe(args); ch != nil {
+	var hook func(PubSubSubscription)
+	if v := ctx.Value(hookKey); v != nil {
+		hook = v.(func(PubSubSubscription))
+	}
+	if ch, cancel := sb.Subscribe(args, hook); ch != nil {
 		defer cancel()
 		if err := p.Do(ctx, subscribe).Error(); err != nil {
 			return err

pipe_test.go

@@ -2962,6 +2962,8 @@ func TestPubSub(t *testing.T) {
 			).Reply(strmsg('+', "PONG"))
 		}()
 
+		confirms := make(chan PubSubSubscription, 2)
+		ctx = WithOnSubscriptionHook(ctx, func(s PubSubSubscription) { confirms <- s })
 		if err := p.Receive(ctx, activate, func(msg PubSubMessage) {
 			if msg.Channel == "1" && msg.Message == "2" {
 				if err := p.Do(ctx, deactivate).Error(); err != nil {
@@ -2971,7 +2973,12 @@ func TestPubSub(t *testing.T) {
 		}); err != nil {
 			t.Fatalf("unexpected err %v", err)
 		}
-
+		if s := <-confirms; s.Kind != "subscribe" || s.Channel != "1" {
+			t.Fatalf("unexpected subscription %v", s)
+		}
+		if s := <-confirms; s.Kind != "unsubscribe" || s.Channel != "1" {
+			t.Fatalf("unexpected subscription %v", s)
+		}
 		cancel()
 	})
 
@@ -3003,6 +3010,8 @@ func TestPubSub(t *testing.T) {
 			).Reply(strmsg('+', "PONG"))
 		}()
 
+		confirms := make(chan PubSubSubscription, 2)
+		ctx = WithOnSubscriptionHook(ctx, func(s PubSubSubscription) { confirms <- s })
 		if err := p.Receive(ctx, activate, func(msg PubSubMessage) {
 			if msg.Channel == "1" && msg.Message == "2" {
 				if err := p.Do(ctx, deactivate).Error(); err != nil {
@@ -3012,7 +3021,12 @@ func TestPubSub(t *testing.T) {
 		}); err != nil {
 			t.Fatalf("unexpected err %v", err)
 		}
-
+		if s := <-confirms; s.Kind != "ssubscribe" || s.Channel != "1" {
+			t.Fatalf("unexpected subscription %v", s)
+		}
+		if s := <-confirms; s.Kind != "sunsubscribe" || s.Channel != "1" {
+			t.Fatalf("unexpected subscription %v", s)
+		}
 		cancel()
 	})
 
@@ -3045,6 +3059,8 @@ func TestPubSub(t *testing.T) {
 			).Reply(strmsg('+', "PONG"))
 		}()
 
+		confirms := make(chan PubSubSubscription, 2)
+		ctx = WithOnSubscriptionHook(ctx, func(s PubSubSubscription) { confirms <- s })
 		if err := p.Receive(ctx, activate, func(msg PubSubMessage) {
 			if msg.Pattern == "1" && msg.Channel == "2" && msg.Message == "3" {
 				if err := p.Do(ctx, deactivate).Error(); err != nil {
@@ -3054,7 +3070,12 @@ func TestPubSub(t *testing.T) {
 		}); err != nil {
 			t.Fatalf("unexpected err %v", err)
 		}
-
+		if s := <-confirms; s.Kind != "psubscribe" || s.Channel != "1" {
+			t.Fatalf("unexpected subscription %v", s)
+		}
+		if s := <-confirms; s.Kind != "punsubscribe" || s.Channel != "1" {
+			t.Fatalf("unexpected subscription %v", s)
+		}
 		cancel()
 	})
 

pubsub.go

@@ -15,9 +15,9 @@ type PubSubMessage struct {
 	Message string
 }
 
-// PubSubSubscription represent a pubsub "subscribe", "unsubscribe", "psubscribe" or "punsubscribe" event.
+// PubSubSubscription represent a pubsub "subscribe", "unsubscribe", "ssubscribe", "sunsubscribe", "psubscribe" or "punsubscribe" event.
 type PubSubSubscription struct {
-	// Kind is "subscribe", "unsubscribe", "psubscribe" or "punsubscribe"
+	// Kind is "subscribe", "unsubscribe", "ssubscribe", "sunsubscribe", "psubscribe" or "punsubscribe"
 	Kind string
 	// Channel is the event subject.
 	Channel string
@@ -54,6 +54,7 @@ type chs struct {
 
 type sub struct {
 	ch chan PubSubMessage
+	fn func(PubSubSubscription)
 	cs []string
 }
 
@@ -67,12 +68,12 @@ func (s *subs) Publish(channel string, msg PubSubMessage) {
 	}
 }
 
-func (s *subs) Subscribe(channels []string) (ch chan PubSubMessage, cancel func()) {
+func (s *subs) Subscribe(channels []string, fn func(PubSubSubscription)) (ch chan PubSubMessage, cancel func()) {
 	id := atomic.AddUint64(&s.cnt, 1)
 	s.mu.Lock()
 	if s.chs != nil {
 		ch = make(chan PubSubMessage, 16)
-		sb := &sub{cs: channels, ch: ch}
+		sb := &sub{cs: channels, ch: ch, fn: fn}
 		s.sub[id] = sb
 		for _, channel := range channels {
 			c := s.chs[channel].sub
@@ -110,13 +111,28 @@ func (s *subs) remove(id uint64) {
 	}
 }
 
-func (s *subs) Unsubscribe(channel string) {
+func (s *subs) Confirm(sub PubSubSubscription) {
+	if atomic.LoadUint64(&s.cnt) != 0 {
+		s.mu.RLock()
+		for _, sb := range s.chs[sub.Channel].sub {
+			if sb.fn != nil {
+				sb.fn(sub)
+			}
+		}
+		s.mu.RUnlock()
+	}
+}
+
+func (s *subs) Unsubscribe(sub PubSubSubscription) {
 	if atomic.LoadUint64(&s.cnt) != 0 {
 		s.mu.Lock()
-		for id := range s.chs[channel].sub {
+		for id, sb := range s.chs[sub.Channel].sub {
+			if sb.fn != nil {
+				sb.fn(sub)
+			}
 			s.remove(id)
 		}
-		delete(s.chs, channel)
+		delete(s.chs, sub.Channel)
 		s.mu.Unlock()
 	}
 }

pubsub_test.go

@@ -14,9 +14,24 @@ func TestSubs_Publish(t *testing.T) {
 
 	t.Run("with multiple subs", func(t *testing.T) {
 		s := newSubs()
-		ch1, cancel1 := s.Subscribe([]string{"a"})
-		ch2, cancel2 := s.Subscribe([]string{"a"})
-		ch3, cancel3 := s.Subscribe([]string{"b"})
+		counts := map[string]int{
+			"a": 0,
+			"b": 0,
+		}
+		subFn := func(s PubSubSubscription) {
+			counts[s.Channel]++
+		}
+
+		ch1, cancel1 := s.Subscribe([]string{"a"}, subFn)
+		ch2, cancel2 := s.Subscribe([]string{"a"}, subFn)
+		ch3, cancel3 := s.Subscribe([]string{"b"}, subFn)
+		s.Confirm(PubSubSubscription{Channel: "a"})
+		s.Confirm(PubSubSubscription{Channel: "b"})
+
+		if counts["a"] != 2 || counts["b"] != 1 {
+			t.Fatalf("unexpected counts %v", counts)
+		}
+
 		m1 := PubSubMessage{Pattern: "1", Channel: "2", Message: "3"}
 		m2 := PubSubMessage{Pattern: "11", Channel: "22", Message: "33"}
 		go func() {
@@ -45,7 +60,7 @@ func TestSubs_Publish(t *testing.T) {
 
 	t.Run("drain ch", func(t *testing.T) {
 		s := newSubs()
-		ch, cancel := s.Subscribe([]string{"a"})
+		ch, cancel := s.Subscribe([]string{"a"}, nil)
 		s.Publish("a", PubSubMessage{})
 		if len(ch) != 1 {
 			t.Fatalf("unexpected ch len %v", len(ch))
@@ -60,15 +75,22 @@ func TestSubs_Publish(t *testing.T) {
 func TestSubs_Unsubscribe(t *testing.T) {
 	defer ShouldNotLeaked(SetupLeakDetection())
 	s := newSubs()
-	ch, _ := s.Subscribe([]string{"1", "2"})
+	counts := map[string]int{"1": 0, "2": 0}
+	subFn := func(s PubSubSubscription) {
+		counts[s.Channel]++
+	}
+	ch, _ := s.Subscribe([]string{"1", "2"}, subFn)
 	go func() {
 		s.Publish("1", PubSubMessage{})
 	}()
 	_, ok := <-ch
 	if !ok {
 		t.Fatalf("unexpected ch closed")
 	}
-	s.Unsubscribe("1")
+	s.Unsubscribe(PubSubSubscription{Channel: "1"})
+	if counts["1"] != 1 {
+		t.Fatalf("unexpected counts %v", counts)
+	}
 	_, ok = <-ch
 	if ok {
 		t.Fatalf("unexpected ch unclosed")