fix: subscriptions bookkeeping in pubsub clients (#510)

anitarua · web-flow · commit 0d5e5144a3ad · 2025-06-18T16:22:25.000-07:00
Pubsub clients should do some bookkeeping of number of active subscriptions to be able to gracefully reject new subscribe requests rather than letting them silently queue up on the client.

Updated the client to use a pool of 4 unary grpc managers by default to match the other SDKs.
The stream managers use the existing with_max_subscriptions configuration to statically set the number of grpc managers in the stream pool.
diff --git a/.github/workflows/on-pull-request.yml b/.github/workflows/on-pull-request.yml
@@ -57,7 +57,7 @@ jobs:
         run: poetry run ruff format --check --diff src tests
 
       - name: Run tests
-        run: poetry run pytest -p no:sugar -q -m "not local"
+        run: poetry run pytest -p no:sugar -q -m "not local and not subscription_initialization"
 
   test-examples:
     runs-on: ubuntu-24.04
diff --git a/Makefile b/Makefile
@@ -50,13 +50,16 @@ gen-sync: do-gen-sync format lint
 .PHONY: test
 ## Run unit and integration tests with pytest
 test:
-	@poetry run pytest -m "not local"
+	@poetry run pytest -m "not local and not subscription_initialization"
 
 .PHONY: test-local
 ## Run the integration tests that require Momento Local
 test-local:
 	@poetry run pytest -m local
 
+test-subscription-initialization:
+	@poetry run pytest -m subscription_initialization
+
 .PHONY: precommit
 ## Run format, lint, and test as a step before committing.
 precommit: gen-sync format lint test
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -44,6 +44,7 @@ pytest = "^7.1.3"
 pytest-asyncio = "^0.19.0"
 pytest-describe = "^2.0.1"
 pytest-sugar = "^0.9.5"
+pytest-timeout = "^2.4.0"
 
 [tool.poetry.group.lint.dependencies]
 mypy = "^1.0"
@@ -61,6 +62,7 @@ log_cli_format = "%(asctime)s [%(levelname)s] %(message)s"
 log_cli_date_format = "%Y-%m-%d %H:%M:%S.%f"
 markers = [
   "local: tests that require Momento Local",
+  "subscription_initialization: tests that require higher subscriptions limit when testing against the live service",
 ]
 
 [tool.mypy]
diff --git a/src/momento/errors/exceptions.py b/src/momento/errors/exceptions.py
@@ -371,3 +371,21 @@ def __init__(
             transport_details,
             message_wrapper="Service returned an unknown response; please contact us at support@momentohq.com",
         )
+
+
+class ClientResourceExhaustedException(SdkException):
+    """Client resource (such as memory or number of concurrent streams) exhausted."""
+
+    def __init__(
+        self,
+        message: str,
+        service: Service,
+        transport_details: Optional[MomentoErrorTransportDetails] = None,
+    ):
+        super().__init__(
+            message,
+            MomentoErrorCode.CLIENT_RESOURCE_EXHAUSTED,
+            service,
+            transport_details,
+            message_wrapper="Client resource (such as memory or number of concurrent streams) exhausted",
+        )
diff --git a/src/momento/internal/aio/_scs_grpc_manager.py b/src/momento/internal/aio/_scs_grpc_manager.py
@@ -13,7 +13,7 @@
 from momento.auth import CredentialProvider
 from momento.config import Configuration, TopicConfiguration
 from momento.config.auth_configuration import AuthConfiguration
-from momento.errors.exceptions import ConnectionException
+from momento.errors.exceptions import ClientResourceExhaustedException, ConnectionException
 from momento.internal._utilities import PYTHON_RUNTIME_VERSION, ClientType
 from momento.internal._utilities._channel_credentials import (
     channel_credentials_from_root_certs_or_default,
@@ -220,13 +220,23 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede
                     configuration.get_transport_strategy().get_grpc_configuration()
                 ),
             )
+        self._active_streams_count = 0
 
     async def close(self) -> None:
         await self._channel.close()
 
     def async_stub(self) -> pubsub_client.PubsubStub:
+        if self._active_streams_count >= 100:
+            raise ClientResourceExhaustedException(
+                message="Already at max number of concurrent streams",
+                service=Service.TOPICS,
+            )
+        self._active_streams_count += 1
         return pubsub_client.PubsubStub(self._channel)  # type: ignore[no-untyped-call]
 
+    def decrement_stream_count(self) -> None:
+        self._active_streams_count -= 1
+
 
 class _TokenGrpcManager:
     """Internal gRPC token manager."""
diff --git a/src/momento/internal/aio/_scs_pubsub_client.py b/src/momento/internal/aio/_scs_pubsub_client.py
@@ -2,6 +2,7 @@
 
 import math
 from datetime import timedelta
+from typing import Callable
 
 from momento_wire_types import cachepubsub_pb2 as pubsub_pb
 from momento_wire_types import cachepubsub_pb2_grpc as pubsub_grpc
@@ -10,6 +11,7 @@
 from momento.auth import CredentialProvider
 from momento.config import TopicConfiguration
 from momento.errors import convert_error
+from momento.errors.exceptions import ClientResourceExhaustedException
 from momento.internal._utilities import _validate_cache_name, _validate_topic_name
 from momento.internal.aio._scs_grpc_manager import (
     _PubsubGrpcManager,
@@ -27,8 +29,6 @@
 class _ScsPubsubClient:
     """Internal pubsub client."""
 
-    stream_topic_manager_count = 0
-
     def __init__(self, configuration: TopicConfiguration, credential_provider: CredentialProvider):
         endpoint = credential_provider.cache_endpoint
         self._logger = logs.logger
@@ -38,20 +38,23 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede
         default_deadline: timedelta = configuration.get_transport_strategy().get_grpc_configuration().get_deadline()
         self._default_deadline_seconds = default_deadline.total_seconds()
 
-        num_subscriptions = configuration.get_max_subscriptions()
         # Default to a single channel and scale up if necessary. Each channel can support
         # 100 subscriptions. Issuing more subscribe requests than you have channels to handle
-        # will cause the last request to hang indefinitely, so it's important to get this right.
+        # will cause a ClientResourceExhaustedException.
         num_channels = 1
+        num_subscriptions = configuration.get_max_subscriptions()
         if num_subscriptions > 0:
             num_channels = math.ceil(num_subscriptions / 100.0)
             self._logger.debug(f"creating {num_channels} subscription channels")
-
-        self._grpc_manager = _PubsubGrpcManager(configuration, credential_provider)
         self._stream_managers = [
             _PubsubGrpcStreamManager(configuration, credential_provider) for i in range(0, num_channels)
         ]
 
+        # Default to 4 unary pubsub channels. TODO: Make this configurable.
+        self._unary_managers = [_PubsubGrpcManager(configuration, credential_provider) for i in range(0, 4)]
+        self._stream_manager_count = 0
+        self._unary_manager_count = 0
+
     @property
     def endpoint(self) -> str:
         return self._endpoint
@@ -72,7 +75,7 @@ async def publish(self, cache_name: str, topic_name: str, value: str | bytes) ->
                 value=topic_value,
             )
 
-            await self._get_stub().Publish(  # type: ignore[misc]
+            await self._get_unary_stub().Publish(  # type: ignore[misc]
                 request,
                 timeout=self._default_deadline_seconds,
             )
@@ -98,7 +101,8 @@ async def subscribe(
                 resume_at_topic_sequence_number=resume_at_topic_sequence_number,
                 sequence_page=resume_at_topic_sequence_page,
             )
-            stream = self._get_stream_stub().Subscribe(  # type: ignore[misc]
+            stub, decrement_stream_count = self._get_stream_stub()
+            stream = stub.Subscribe(  # type: ignore[misc]
                 request,
             )
 
@@ -112,23 +116,48 @@ async def subscribe(
                 err = Exception(f"expected a heartbeat message but got '{msg_type}'")
                 self._log_request_error("subscribe", err)
                 return TopicSubscribe.Error(convert_error(err, Service.TOPICS))
-            return TopicSubscribe.SubscriptionAsync(cache_name, topic_name, client_stream=stream)  # type: ignore[misc]
+            return TopicSubscribe.SubscriptionAsync(
+                cache_name,
+                topic_name,
+                client_stream=stream,  # type: ignore[misc]
+                decrement_stream_count_method=decrement_stream_count,
+            )
         except Exception as e:
             self._log_request_error("subscribe", e)
             return TopicSubscribe.Error(convert_error(e, Service.TOPICS))
 
     def _log_request_error(self, request_type: str, e: Exception) -> None:
         self._logger.warning(f"{request_type} failed with exception: {e}")
 
-    def _get_stub(self) -> pubsub_grpc.PubsubStub:
-        return self._grpc_manager.async_stub()
-
-    def _get_stream_stub(self) -> pubsub_grpc.PubsubStub:
-        stub = self._stream_managers[self.stream_topic_manager_count % len(self._stream_managers)].async_stub()
-        self.stream_topic_manager_count += 1
-        return stub
+    def _get_unary_stub(self) -> pubsub_grpc.PubsubStub:
+        # Simply round-robin through the unary managers.
+        # Unary requests will eventually complete (unlike long-lived subscriptions),
+        # so we do not need the same bookkeeping logic here.
+        manager = self._unary_managers[self._unary_manager_count % len(self._unary_managers)]
+        self._unary_manager_count += 1
+        return manager.async_stub()
+
+    def _get_stream_stub(self) -> tuple[pubsub_grpc.PubsubStub, Callable[[], None]]:
+        # Try to get a client with capacity for another subscription by round-robining through the stubs.
+        # Allow up to max_stream_capacity attempts to account for large bursts of requests.
+        max_stream_capacity = len(self._stream_managers) * 100
+        for _ in range(0, max_stream_capacity):
+            try:
+                manager = self._stream_managers[self._stream_manager_count % len(self._stream_managers)]
+                self._stream_manager_count += 1
+                return manager.async_stub(), manager.decrement_stream_count
+            except ClientResourceExhaustedException:
+                # If the stub is at capacity, continue to the next one.
+                continue
+
+        # Otherwise return an error if no stubs have capacity.
+        raise ClientResourceExhaustedException(
+            message="Maximum number of active subscriptions reached",
+            service=Service.TOPICS,
+        )
 
     async def close(self) -> None:
-        await self._grpc_manager.close()
+        for unary_client in self._unary_managers:
+            await unary_client.close()
         for stream_client in self._stream_managers:
             await stream_client.close()
diff --git a/src/momento/internal/synchronous/_scs_grpc_manager.py b/src/momento/internal/synchronous/_scs_grpc_manager.py
@@ -17,7 +17,7 @@
 from momento.config.middleware import MiddlewareRequestHandlerContext
 from momento.config.middleware.models import CONNECTION_ID_KEY
 from momento.config.middleware.synchronous import Middleware
-from momento.errors.exceptions import ConnectionException
+from momento.errors.exceptions import ClientResourceExhaustedException, ConnectionException
 from momento.internal._utilities import PYTHON_RUNTIME_VERSION, ClientType
 from momento.internal._utilities._channel_credentials import (
     channel_credentials_from_root_certs_or_default,
@@ -158,7 +158,7 @@ def on_state_change(state: grpc.ChannelConnectivity) -> None:
             elif state == connecting:
                 self._logger.debug("State transitioned to CONNECTING; waiting to get READY")
             else:
-                self._logger.warn(f"Unexpected connection state while trying to eagerly connect: {state}")
+                self._logger.warning(f"Unexpected connection state while trying to eagerly connect: {state}")
                 # we could not connect within the timeout and we no longer need this subscription
                 self._channel.unsubscribe(on_state_change)
                 connection_event.set()
@@ -233,13 +233,23 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede
             self._secure_channel, *_stream_interceptors(credential_provider.auth_token, ClientType.TOPIC)
         )
         self._stub = pubsub_client.PubsubStub(intercept_channel)  # type: ignore[no-untyped-call]
+        self._active_streams_count = 0
 
     def close(self) -> None:
         self._secure_channel.close()
 
     def stub(self) -> pubsub_client.PubsubStub:
+        if self._active_streams_count >= 100:
+            raise ClientResourceExhaustedException(
+                message="Already at max number of concurrent streams",
+                service=Service.TOPICS,
+            )
+        self._active_streams_count += 1
         return self._stub
 
+    def decrement_stream_count(self) -> None:
+        self._active_streams_count -= 1
+
 
 class _TokenGrpcManager:
     """Internal gRPC token manager."""
diff --git a/src/momento/internal/synchronous/_scs_pubsub_client.py b/src/momento/internal/synchronous/_scs_pubsub_client.py
@@ -2,6 +2,7 @@
 
 import math
 from datetime import timedelta
+from typing import Callable
 
 from momento_wire_types import cachepubsub_pb2 as pubsub_pb
 from momento_wire_types import cachepubsub_pb2_grpc as pubsub_grpc
@@ -10,6 +11,7 @@
 from momento.auth import CredentialProvider
 from momento.config import TopicConfiguration
 from momento.errors import convert_error
+from momento.errors.exceptions import ClientResourceExhaustedException
 from momento.internal._utilities import _validate_cache_name, _validate_topic_name
 from momento.internal.services import Service
 from momento.internal.synchronous._scs_grpc_manager import (
@@ -27,8 +29,6 @@
 class _ScsPubsubClient:
     """Internal pubsub client."""
 
-    stream_topic_manager_count = 0
-
     def __init__(self, configuration: TopicConfiguration, credential_provider: CredentialProvider):
         endpoint = credential_provider.cache_endpoint
         self._logger = logs.logger
@@ -38,20 +38,23 @@ def __init__(self, configuration: TopicConfiguration, credential_provider: Crede
         default_deadline: timedelta = configuration.get_transport_strategy().get_grpc_configuration().get_deadline()
         self._default_deadline_seconds = default_deadline.total_seconds()
 
-        num_subscriptions = configuration.get_max_subscriptions()
         # Default to a single channel and scale up if necessary. Each channel can support
         # 100 subscriptions. Issuing more subscribe requests than you have channels to handle
-        # will cause the last request to hang indefinitely, so it's important to get this right.
+        # will cause a ClientResourceExhaustedException.
         num_channels = 1
+        num_subscriptions = configuration.get_max_subscriptions()
         if num_subscriptions > 0:
             num_channels = math.ceil(num_subscriptions / 100.0)
             self._logger.debug(f"creating {num_channels} subscription channels")
-
-        self._grpc_manager = _PubsubGrpcManager(configuration, credential_provider)
         self._stream_managers = [
             _PubsubGrpcStreamManager(configuration, credential_provider) for i in range(0, num_channels)
         ]
 
+        # Default to 4 unary pubsub channels. TODO: Make this configurable.
+        self._unary_managers = [_PubsubGrpcManager(configuration, credential_provider) for i in range(0, 4)]
+        self._stream_manager_count = 0
+        self._unary_manager_count = 0
+
     @property
     def endpoint(self) -> str:
         return self._endpoint
@@ -72,7 +75,7 @@ def publish(self, cache_name: str, topic_name: str, value: str | bytes) -> Topic
                 value=topic_value,
             )
 
-            self._get_stub().Publish(
+            self._get_unary_stub().Publish(
                 request,
                 timeout=self._default_deadline_seconds,
             )
@@ -98,7 +101,8 @@ def subscribe(
                 resume_at_topic_sequence_number=resume_at_topic_sequence_number,
                 sequence_page=resume_at_topic_sequence_page,
             )
-            stream = self._get_stream_stub().Subscribe(  # type: ignore[misc]
+            stub, decrement_stream_count = self._get_stream_stub()
+            stream = stub.Subscribe(  # type: ignore[misc]
                 request,
             )
 
@@ -112,23 +116,48 @@ def subscribe(
                 err = Exception(f"expected a heartbeat message but got '{msg_type}'")
                 self._log_request_error("subscribe", err)
                 return TopicSubscribe.Error(convert_error(err, Service.TOPICS))
-            return TopicSubscribe.Subscription(cache_name, topic_name, client_stream=stream)  # type: ignore[misc]
+            return TopicSubscribe.Subscription(
+                cache_name,
+                topic_name,
+                client_stream=stream,  # type: ignore[misc]
+                decrement_stream_count_method=decrement_stream_count,
+            )
         except Exception as e:
             self._log_request_error("subscribe", e)
             return TopicSubscribe.Error(convert_error(e, Service.TOPICS))
 
     def _log_request_error(self, request_type: str, e: Exception) -> None:
         self._logger.warning(f"{request_type} failed with exception: {e}")
 
-    def _get_stub(self) -> pubsub_grpc.PubsubStub:
-        return self._grpc_manager.stub()
-
-    def _get_stream_stub(self) -> pubsub_grpc.PubsubStub:
-        stub = self._stream_managers[self.stream_topic_manager_count % len(self._stream_managers)].stub()
-        self.stream_topic_manager_count += 1
+    def _get_unary_stub(self) -> pubsub_grpc.PubsubStub:
+        # Simply round-robin through the unary managers.
+        # Unary requests will eventually complete (unlike long-lived subscriptions),
+        # so we do not need the same bookkeeping logic here.
+        stub = self._unary_managers[self._unary_manager_count % len(self._unary_managers)].stub()
+        self._unary_manager_count += 1
         return stub
 
+    def _get_stream_stub(self) -> tuple[pubsub_grpc.PubsubStub, Callable[[], None]]:
+        # Try to get a client with capacity for another subscription by round-robining through the stubs.
+        # Allow up to max_stream_capacity attempts to account for large bursts of requests.
+        max_stream_capacity = len(self._stream_managers) * 100
+        for _ in range(0, max_stream_capacity):
+            try:
+                manager = self._stream_managers[self._stream_manager_count % len(self._stream_managers)]
+                self._stream_manager_count += 1
+                return manager.stub(), manager.decrement_stream_count
+            except ClientResourceExhaustedException:
+                # If the stub is at capacity, continue to the next one.
+                continue
+
+        # Otherwise return an error if no stubs have capacity.
+        raise ClientResourceExhaustedException(
+            message="Maximum number of active subscriptions reached",
+            service=Service.TOPICS,
+        )
+
     def close(self) -> None:
-        self._grpc_manager.close()
+        for unary_manager in self._unary_managers:
+            unary_manager.close()
         for stream_client in self._stream_managers:
             stream_client.close()
diff --git a/src/momento/responses/pubsub/subscribe.py b/src/momento/responses/pubsub/subscribe.py
diff --git a/tests/momento/local/test_topic_subscribe.py b/tests/momento/local/test_topic_subscribe.py
diff --git a/tests/momento/local/test_topic_subscribe_async.py b/tests/momento/local/test_topic_subscribe_async.py
