Refactor subscriber index move to compstore

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

Author: antontroshin

pkg/api/grpc/subscribe.go

@@ -76,7 +76,7 @@ func (a *api) SubscribeTopicEventsAlpha1(stream runtimev1pb.Dapr_SubscribeTopicE
 			Routes:          subapi.Routes{Default: "/"},
 		},
 	}
-	connectionID := a.pubsubAdapterStreamer.NextIndex()
+	connectionID := a.Universal.CompStore().NextSubscriberIndex()
 	err = a.Universal.CompStore().AddStreamSubscription(sub, connectionID)
 	if err != nil {
 		return err

pkg/runtime/compstore/compstore.go

@@ -15,6 +15,7 @@ package compstore
 
 import (
 	"sync"
+	"sync/atomic"
 
 	"github.com/dapr/durabletask-go/backend"
 
@@ -63,6 +64,9 @@ type ComponentStore struct {
 
 	compPendingLock sync.Mutex
 	compPending     *compsv1alpha1.Component
+
+	// subscribersIndex is used to generate unique connection IDs for streaming subscriptions
+	subscribersIndex atomic.Uint64
 }
 
 func New() *ComponentStore {

pkg/runtime/compstore/subscriptions.go

@@ -277,3 +277,7 @@ func (c *ComponentStore) ListProgramaticSubscriptions() []rtpubsub.Subscription
 	defer c.lock.RUnlock()
 	return c.subscriptions.programmatics
 }
+
+func (c *ComponentStore) NextSubscriberIndex() rtpubsub.ConnectionID {
+	return rtpubsub.ConnectionID(c.subscribersIndex.Add(1))
+}

pkg/runtime/compstore/subscriptions_index_test.go

@@ -0,0 +1,65 @@
+/*
+Copyright 2025 The Dapr Authors
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package compstore
+
+import (
+	"sync"
+	"testing"
+
+	rtpubsub "github.com/dapr/dapr/pkg/runtime/pubsub"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNextSubscriberIndex(t *testing.T) {
+	t.Run("sequential calls return incrementing values", func(t *testing.T) {
+		store := New()
+
+		id1 := store.NextSubscriberIndex()
+		id2 := store.NextSubscriberIndex()
+		id3 := store.NextSubscriberIndex()
+
+		assert.Equal(t, rtpubsub.ConnectionID(1), id1)
+		assert.Equal(t, rtpubsub.ConnectionID(2), id2)
+		assert.Equal(t, rtpubsub.ConnectionID(3), id3)
+	})
+
+	t.Run("concurrent calls return unique values", func(t *testing.T) {
+		store := New()
+		const numGoroutines = 100
+		const numCallsPerGoroutine = 10
+		var wg sync.WaitGroup
+		var mu sync.Mutex
+		ids := make(map[rtpubsub.ConnectionID]bool)
+
+		wg.Add(numGoroutines)
+		for i := 0; i < numGoroutines; i++ {
+			go func() {
+				defer wg.Done()
+				for j := 0; j < numCallsPerGoroutine; j++ {
+					id := store.NextSubscriberIndex()
+					mu.Lock()
+					ids[id] = true
+					mu.Unlock()
+				}
+			}()
+		}
+		wg.Wait()
+
+		assert.Len(t, ids, numGoroutines*numCallsPerGoroutine, "Expected all IDs to be unique")
+
+		for i := 1; i <= numGoroutines*numCallsPerGoroutine; i++ {
+			assert.True(t, ids[rtpubsub.ConnectionID(i)], "Expected ID %d to be present", i)
+		}
+	})
+}

pkg/runtime/pubsub/adapter.go

@@ -47,7 +47,6 @@ type AdapterStreamer interface {
 	Subscribe(rtv1pb.Dapr_SubscribeTopicEventsAlpha1Server, *rtv1pb.SubscribeTopicEventsRequestInitialAlpha1, ConnectionID) error
 	Publish(context.Context, *SubscribedMessage) error
 	StreamerKey(pubsub, topic string) string
-	NextIndex() ConnectionID
 	Close(key string, connectionID ConnectionID)
 }
 

pkg/runtime/pubsub/streamer/streamer.go

@@ -20,7 +20,6 @@ import (
 	"io"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"google.golang.org/grpc/codes"
@@ -44,9 +43,8 @@ type Connections map[rtpubsub.ConnectionID]*conn
 type Subscribers map[string]Connections
 
 type streamer struct {
-	tracingSpec      *config.TracingSpec
-	subscribers      Subscribers
-	subscribersIndex atomic.Uint64
+	tracingSpec *config.TracingSpec
+	subscribers Subscribers
 
 	lock sync.RWMutex
 }
@@ -228,10 +226,6 @@ func (s *streamer) StreamerKey(pubsub, topic string) string {
 	return "___" + pubsub + "||" + topic
 }
 
-func (s *streamer) NextIndex() rtpubsub.ConnectionID {
-	return rtpubsub.ConnectionID(s.subscribersIndex.Add(1))
-}
-
 func (s *streamer) Close(key string, connectionID rtpubsub.ConnectionID) {
 	log.Warn("RLock Close")
 	s.lock.RLock()