Add support for multiple Streaming Subscriptions subscribers

Signed-off-by: Anton Troshin <[email protected]>

Author: antontroshin

pkg/api/grpc/subscribe.go

@@ -66,7 +66,7 @@ func (a *api) SubscribeTopicEventsAlpha1(stream runtimev1pb.Dapr_SubscribeTopicE
 	}
 
 	key := a.pubsubAdapterStreamer.StreamerKey(req.GetPubsubName(), req.GetTopic())
-	err = a.Universal.CompStore().AddStreamSubscription(&subapi.Subscription{
+	sub := &subapi.Subscription{
 		ObjectMeta: metav1.ObjectMeta{Name: key},
 		Spec: subapi.SubscriptionSpec{
 			Pubsubname:      req.GetPubsubName(),
@@ -75,19 +75,21 @@ func (a *api) SubscribeTopicEventsAlpha1(stream runtimev1pb.Dapr_SubscribeTopicE
 			DeadLetterTopic: req.GetDeadLetterTopic(),
 			Routes:          subapi.Routes{Default: "/"},
 		},
-	})
+	}
+	connectionID := a.pubsubAdapterStreamer.NextIndex()
+	err = a.Universal.CompStore().AddStreamSubscription(sub, connectionID)
 	if err != nil {
 		return err
 	}
 
-	if err = a.processor.Subscriber().StartStreamerSubscription(key); err != nil {
-		a.Universal.CompStore().DeleteStreamSubscription(key)
+	if err = a.processor.Subscriber().StartStreamerSubscription(sub, connectionID); err != nil {
+		a.Universal.CompStore().DeleteStreamSubscription(sub)
 		return err
 	}
 
 	defer func() {
-		a.processor.Subscriber().StopStreamerSubscription(req.GetPubsubName(), key)
-		a.Universal.CompStore().DeleteStreamSubscription(key)
+		a.processor.Subscriber().StopStreamerSubscription(sub, connectionID)
+		a.Universal.CompStore().DeleteStreamSubscription(sub)
 	}()
 
 	if err = stream.Send(&runtimev1pb.SubscribeTopicEventsResponseAlpha1{
@@ -98,5 +100,5 @@ func (a *api) SubscribeTopicEventsAlpha1(stream runtimev1pb.Dapr_SubscribeTopicE
 		return err
 	}
 
-	return a.pubsubAdapterStreamer.Subscribe(stream, req)
+	return a.pubsubAdapterStreamer.Subscribe(stream, req, connectionID)
 }

pkg/runtime/compstore/compstore.go

@@ -82,7 +82,7 @@ func New() *ComponentStore {
 		cryptoProviders:         make(map[string]crypto.SubtleCrypto),
 		subscriptions: &subscriptions{
 			declaratives: make(map[string]*DeclarativeSubscription),
-			streams:      make(map[string]*DeclarativeSubscription),
+			streams:      make(map[string][]*DeclarativeSubscription),
 		},
 		conversations: make(map[string]conversation.Conversation),
 	}

pkg/runtime/compstore/subscriptions.go

@@ -14,14 +14,15 @@ limitations under the License.
 package compstore
 
 import (
-	"fmt"
-
 	subapi "github.com/dapr/dapr/pkg/apis/subscriptions/v2alpha1"
 	runtimev1pb "github.com/dapr/dapr/pkg/proto/runtime/v1"
 	rtpubsub "github.com/dapr/dapr/pkg/runtime/pubsub"
+	"github.com/dapr/kit/logger"
 	"github.com/dapr/kit/ptr"
 )
 
+var log = logger.NewLogger("dapr.runtime.compstore.subscriptions")
+
 type DeclarativeSubscription struct {
 	Comp *subapi.Subscription
 	*NamedSubscription
@@ -30,7 +31,8 @@ type DeclarativeSubscription struct {
 type NamedSubscription struct {
 	// Name is the optional name of the subscription. If not set, the name of the
 	// component is used.
-	Name *string
+	Name         *string
+	ConnectionID rtpubsub.ConnectionID
 	rtpubsub.Subscription
 }
 
@@ -40,7 +42,7 @@ type subscriptions struct {
 	// declarativesList used to track order of declarative subscriptions for
 	// processing priority.
 	declarativesList []string
-	streams          map[string]*DeclarativeSubscription
+	streams          map[string][]*DeclarativeSubscription
 }
 
 // MetadataSubscription is a temporary wrapper for rtpubsub.Subscription to add a Type attribute to be used in GetMetadata
@@ -75,17 +77,19 @@ func (c *ComponentStore) AddDeclarativeSubscription(comp *subapi.Subscription, s
 	c.subscriptions.declarativesList = append(c.subscriptions.declarativesList, comp.Name)
 }
 
-func (c *ComponentStore) AddStreamSubscription(comp *subapi.Subscription) error {
+func (c *ComponentStore) AddStreamSubscription(comp *subapi.Subscription, connectionID rtpubsub.ConnectionID) error {
+	log.Warn("Lock AddStreamSubscription")
 	c.lock.Lock()
-	defer c.lock.Unlock()
-	if _, ok := c.subscriptions.streams[comp.Name]; ok {
-		return fmt.Errorf("streamer already subscribed to pubsub %q topic %q", comp.Spec.Pubsubname, comp.Spec.Topic)
-	}
+	defer func() {
+		c.lock.Unlock()
+		log.Warn("Unlock AddStreamSubscription defer")
+	}()
 
-	c.subscriptions.streams[comp.Name] = &DeclarativeSubscription{
+	sub := &DeclarativeSubscription{
 		Comp: comp,
 		NamedSubscription: &NamedSubscription{
-			Name: ptr.Of(comp.Name),
+			Name:         ptr.Of(comp.Name),
+			ConnectionID: connectionID,
 			Subscription: rtpubsub.Subscription{
 				PubsubName:      comp.Spec.Pubsubname,
 				Topic:           comp.Spec.Topic,
@@ -95,15 +99,28 @@ func (c *ComponentStore) AddStreamSubscription(comp *subapi.Subscription) error
 			},
 		},
 	}
+	c.subscriptions.streams[comp.Name] = append(c.subscriptions.streams[comp.Name], sub)
 
 	return nil
 }
 
-func (c *ComponentStore) DeleteStreamSubscription(names ...string) {
+func (c *ComponentStore) DeleteStreamSubscription(comp *subapi.Subscription) {
+	log.Warn("Lock DeleteStreamSubscription")
 	c.lock.Lock()
-	defer c.lock.Unlock()
-	for _, name := range names {
-		delete(c.subscriptions.streams, name)
+	defer func() {
+		c.lock.Unlock()
+		log.Warn("Unlock DeleteStreamSubscription defer")
+	}()
+	streamingSubscriptions, ok := c.subscriptions.streams[comp.Name]
+	if ok && len(streamingSubscriptions) > 0 {
+		for i, sub := range streamingSubscriptions {
+			if sub.Comp == comp {
+				c.subscriptions.streams[comp.Name] = append(c.subscriptions.streams[comp.Name][:i], c.subscriptions.streams[comp.Name][i+1:]...)
+			}
+		}
+	}
+	if len(c.subscriptions.streams[comp.Name]) == 0 {
+		delete(c.subscriptions.streams, comp.Name)
 	}
 }
 
@@ -155,17 +172,12 @@ func (c *ComponentStore) ListTypedSubscriptions() []TypedSubscription {
 			subs = append(subs, typedSub)
 		}
 	}
-	for i := range c.subscriptions.streams {
-		sub := c.subscriptions.streams[i].Subscription
-		typedSub := TypedSubscription{
-			Subscription: sub,
-			Type:         runtimev1pb.PubsubSubscriptionType_STREAMING,
-		}
-		key := sub.PubsubName + "||" + sub.Topic
-		if j, ok := taken[key]; ok {
-			subs[j] = typedSub
-		} else {
-			taken[key] = len(subs)
+	for _, streamingSubs := range c.subscriptions.streams {
+		for _, sub := range streamingSubs {
+			typedSub := TypedSubscription{
+				Subscription: sub.Subscription,
+				Type:         runtimev1pb.PubsubSubscriptionType_STREAMING,
+			}
 			subs = append(subs, typedSub)
 		}
 	}
@@ -215,9 +227,11 @@ func (c *ComponentStore) ListSubscriptionsStreamByPubSub(name string) []*NamedSu
 	defer c.lock.RUnlock()
 
 	var subs []*NamedSubscription
-	for _, sub := range c.subscriptions.streams {
-		if sub.Subscription.PubsubName == name {
-			subs = append(subs, sub.NamedSubscription)
+	for _, subscriptions := range c.subscriptions.streams {
+		for _, sub := range subscriptions {
+			if sub.Subscription.PubsubName == name {
+				subs = append(subs, sub.NamedSubscription)
+			}
 		}
 	}
 
@@ -235,12 +249,14 @@ func (c *ComponentStore) GetDeclarativeSubscription(name string) (*DeclarativeSu
 	return nil, false
 }
 
-func (c *ComponentStore) GetStreamSubscription(key string) (*NamedSubscription, bool) {
+func (c *ComponentStore) GetStreamSubscription(subscription *subapi.Subscription) (*NamedSubscription, bool) {
 	c.lock.RLock()
 	defer c.lock.RUnlock()
-	for i, sub := range c.subscriptions.streams {
-		if sub.Comp.Name == key {
-			return c.subscriptions.streams[i].NamedSubscription, true
+	for _, subscriptions := range c.subscriptions.streams {
+		for _, sub := range subscriptions {
+			if sub.Comp == subscription {
+				return sub.NamedSubscription, true
+			}
 		}
 	}
 	return nil, false

pkg/runtime/processor/manager.go

@@ -16,6 +16,8 @@ package processor
 import (
 	"context"
 	"fmt"
+	subapi "github.com/dapr/dapr/pkg/apis/subscriptions/v2alpha1"
+	rtpubsub "github.com/dapr/dapr/pkg/runtime/pubsub"
 
 	"github.com/dapr/durabletask-go/backend"
 
@@ -46,8 +48,8 @@ type SubscribeManager interface {
 	StopAppSubscriptions()
 	StopAllSubscriptionsForever()
 	ReloadDeclaredAppSubscription(name, pubsubName string) error
-	StartStreamerSubscription(key string) error
-	StopStreamerSubscription(pubsubName, key string)
+	StartStreamerSubscription(sub *subapi.Subscription, connectionID rtpubsub.ConnectionID) error
+	StopStreamerSubscription(sub *subapi.Subscription, connectionID rtpubsub.ConnectionID)
 	ReloadPubSub(string) error
 	StopPubSub(string)
 }

pkg/runtime/processor/subscriber/subscriber.go

@@ -17,6 +17,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	subapi "github.com/dapr/dapr/pkg/apis/subscriptions/v2alpha1"
 	"sync"
 	"sync/atomic"
 
@@ -60,7 +61,7 @@ type Subscriber struct {
 	adapterStreamer rtpubsub.AdapterStreamer
 
 	appSubs      map[string][]*namedSubscription
-	streamSubs   map[string][]*namedSubscription
+	streamSubs   map[string]map[rtpubsub.ConnectionID]*namedSubscription
 	appSubActive bool
 	hasInitProg  bool
 	lock         sync.RWMutex
@@ -88,7 +89,7 @@ func New(opts Options) *Subscriber {
 		adapter:         opts.Adapter,
 		adapterStreamer: opts.AdapterStreamer,
 		appSubs:         make(map[string][]*namedSubscription),
-		streamSubs:      make(map[string][]*namedSubscription),
+		streamSubs:      make(map[string]map[rtpubsub.ConnectionID]*namedSubscription),
 	}
 }
 
@@ -126,19 +127,19 @@ func (s *Subscriber) ReloadPubSub(name string) error {
 	return errors.Join(errs...)
 }
 
-func (s *Subscriber) StartStreamerSubscription(key string) error {
+func (s *Subscriber) StartStreamerSubscription(subscription *subapi.Subscription, connectionID rtpubsub.ConnectionID) error {
+	log.Warn("Lock StartStreamerSubscription")
 	s.lock.Lock()
-	defer s.lock.Unlock()
+	defer func() {
+		s.lock.Unlock()
+		log.Warn("Unlock StartStreamerSubscription defer")
+	}()
 
 	if s.closed {
 		return apierrors.PubSub("").WithMetadata(nil).DeserializeError(errors.New("subscriber is closed"))
 	}
 
-	sub, ok := s.compStore.GetStreamSubscription(key)
-	if !ok {
-		err := fmt.Errorf("starting stream subscription without connection: %s", key)
-		return apierrors.PubSub("").WithMetadata(nil).DeserializeError(err)
-	}
+	sub, ok := s.compStore.GetStreamSubscription(subscription)
 
 	pubsub, ok := s.compStore.GetPubSub(sub.PubsubName)
 	if !ok {
@@ -150,28 +151,46 @@ func (s *Subscriber) StartStreamerSubscription(key string) error {
 		return fmt.Errorf("failed to create subscription for %s: %s", sub.PubsubName, err)
 	}
 
-	s.streamSubs[sub.PubsubName] = append(s.streamSubs[sub.PubsubName], &namedSubscription{
+	key := s.adapterStreamer.StreamerKey(sub.PubsubName, sub.Topic)
+
+	subscriptions, ok := s.streamSubs[sub.PubsubName]
+	if subscriptions == nil {
+		s.streamSubs[sub.PubsubName] = make(map[rtpubsub.ConnectionID]*namedSubscription)
+	}
+
+	s.streamSubs[sub.PubsubName][connectionID] = &namedSubscription{
 		name:         &key,
 		Subscription: ss,
-	})
-
+	}
 	return nil
 }
 
-func (s *Subscriber) StopStreamerSubscription(pubsubName, key string) {
+func (s *Subscriber) StopStreamerSubscription(subscription *subapi.Subscription, connectionID rtpubsub.ConnectionID) {
+	log.Warn("Lock StopStreamerSubscription")
 	s.lock.Lock()
-	defer s.lock.Unlock()
+	defer func() {
+		s.lock.Unlock()
+		log.Warn("Unlock StopStreamerSubscription defer")
+	}()
 
 	if s.closed {
 		return
 	}
 
-	for i, sub := range s.streamSubs[pubsubName] {
-		if sub.name != nil && *sub.name == key {
-			sub.Stop()
-			s.streamSubs[pubsubName] = append(s.streamSubs[pubsubName][:i], s.streamSubs[pubsubName][i+1:]...)
-			return
-		}
+	subscriptions, ok := s.streamSubs[subscription.Spec.Pubsubname]
+	if !ok {
+		return
+	}
+
+	sub, ok := subscriptions[connectionID]
+	if !ok {
+		return
+	}
+
+	sub.Stop()
+	delete(subscriptions, connectionID)
+	if len(subscriptions) == 0 {
+		delete(s.streamSubs, subscription.Spec.Pubsubname)
 	}
 }
 
@@ -293,8 +312,12 @@ func (s *Subscriber) StopAppSubscriptions() {
 }
 
 func (s *Subscriber) StopAllSubscriptionsForever() {
+	log.Warn("Lock StopAllSubscriptionsForever")
 	s.lock.Lock()
-	defer s.lock.Unlock()
+	defer func() {
+		s.lock.Unlock()
+		log.Warn("Unlock StopAllSubscriptionsForever defer")
+	}()
 
 	s.closed = true
 
@@ -310,7 +333,7 @@ func (s *Subscriber) StopAllSubscriptionsForever() {
 	}
 
 	s.appSubs = make(map[string][]*namedSubscription)
-	s.streamSubs = make(map[string][]*namedSubscription)
+	s.streamSubs = make(map[string]map[rtpubsub.ConnectionID]*namedSubscription)
 }
 
 func (s *Subscriber) InitProgramaticSubscriptions(ctx context.Context) error {
@@ -329,7 +352,8 @@ func (s *Subscriber) reloadPubSubStream(name string, pubsub *rtpubsub.PubsubItem
 		return nil
 	}
 
-	subs := make([]*namedSubscription, 0, len(s.compStore.ListSubscriptionsStreamByPubSub(name)))
+	//subs := make([]*namedSubscription, 0, len(s.compStore.ListSubscriptionsStreamByPubSub(name)))
+	subs := make(map[rtpubsub.ConnectionID]*namedSubscription, len(s.compStore.ListSubscriptionsStreamByPubSub(name)))
 	var errs []error
 	for _, sub := range s.compStore.ListSubscriptionsStreamByPubSub(name) {
 		ss, err := s.startSubscription(pubsub, sub, true)
@@ -338,10 +362,10 @@ func (s *Subscriber) reloadPubSubStream(name string, pubsub *rtpubsub.PubsubItem
 			continue
 		}
 
-		subs = append(subs, &namedSubscription{
+		subs[sub.ConnectionID] = &namedSubscription{
 			name:         sub.Name,
 			Subscription: ss,
-		})
+		}
 	}
 
 	s.streamSubs[name] = subs
@@ -455,5 +479,6 @@ func (s *Subscriber) startSubscription(pubsub *rtpubsub.PubsubItem, comp *compst
 		GRPC:            s.grpc,
 		Adapter:         s.adapter,
 		AdapterStreamer: streamer,
+		ConnectionID:    comp.ConnectionID,
 	})
 }