feat: allow custom broker

internal/rdb/rdb.go

@@ -32,6 +32,8 @@ type RDB struct {
 	queuesPublished sync.Map
 }
 
+var _ base.Broker = &RDB{}
+
 // NewRDB returns a new instance of RDB.
 func NewRDB(client redis.UniversalClient) *RDB {
 	return &RDB{
@@ -1481,8 +1483,31 @@ func (r *RDB) ClearSchedulerEntries(schedulerID string) error {
 	return nil
 }
 
-// CancelationPubSub returns a pubsub for cancelation messages.
-func (r *RDB) CancelationPubSub() (*redis.PubSub, error) {
+// cancelationSubscription is a wrapper for redis pubsub.
+type cancellationSubscription struct {
+	pubsub *redis.PubSub
+}
+
+func (c *cancellationSubscription) Channel() <-chan string {
+	channelSize := 100 // same as redis defaults
+	ch := make(chan string, channelSize)
+
+	go func() {
+		for msg := range c.pubsub.Channel(redis.WithChannelSize(channelSize)) {
+			ch <- msg.Payload
+		}
+		close(ch)
+	}()
+
+	return ch
+}
+
+func (c *cancellationSubscription) Close() error {
+	return c.pubsub.Close()
+}
+
+// SubscribeCancellation returns a subscription for cancelation messages.
+func (r *RDB) SubscribeCancellation() (base.CancellationSubscription, error) {
 	var op errors.Op = "rdb.CancelationPubSub"
 	ctx := context.Background()
 	pubsub := r.client.Subscribe(ctx, base.CancelChannel)
@@ -1490,7 +1515,7 @@ func (r *RDB) CancelationPubSub() (*redis.PubSub, error) {
 	if err != nil {
 		return nil, errors.E(op, errors.Unknown, fmt.Sprintf("redis pubsub receive error: %v", err))
 	}
-	return pubsub, nil
+	return &cancellationSubscription{pubsub: pubsub}, nil
 }
 
 // PublishCancelation publish cancelation message to all subscribers.

internal/rdb/rdb_test.go

@@ -3236,12 +3236,12 @@ func TestCancelationPubSub(t *testing.T) {
 	r := setup(t)
 	defer r.Close()
 
-	pubsub, err := r.CancelationPubSub()
+	sub, err := r.SubscribeCancellation()
 	if err != nil {
 		t.Fatalf("(*RDB).CancelationPubSub() returned an error: %v", err)
 	}
 
-	cancelCh := pubsub.Channel()
+	cancelCh := sub.Channel()
 
 	var (
 		mu       sync.Mutex
@@ -3249,9 +3249,9 @@ func TestCancelationPubSub(t *testing.T) {
 	)
 
 	go func() {
-		for msg := range cancelCh {
+		for id := range cancelCh {
 			mu.Lock()
-			received = append(received, msg.Payload)
+			received = append(received, id)
 			mu.Unlock()
 		}
 	}()
@@ -3265,7 +3265,7 @@ func TestCancelationPubSub(t *testing.T) {
 	// allow for message to reach subscribers.
 	time.Sleep(time.Second)
 
-	pubsub.Close()
+	sub.Close()
 
 	mu.Lock()
 	if diff := cmp.Diff(publish, received, h.SortStringSliceOpt); diff != "" {