Repository: ably/ably-python
Pull request #420 (2023-01-26 12:11:03)

Refactor realtime channel
Similar to #414, this PR refactors the internal state machine logic for realtime channel and resolves a few issues where channels would end up in invalid states. The main changes are:
- Moving state change logic into `request_state` and `notify_state` methods
- Using the `Timer` class to handle request timeouts


Patch:
diff --git a/ably/realtime/connection.py b/ably/realtime/connection.py
index f00a518..b9ffccc 100644
--- a/ably/realtime/connection.py
+++ b/ably/realtime/connection.py
@@ -323,13 +323,7 @@ class ConnectionManager(EventEmitter):
             self.notify_state(self.__fail_state, reason=exception)
 
     def notify_state(self, state: ConnectionState, reason=None):
-        # RTN15a
-        retry_immediately = state == ConnectionState.DISCONNECTED and self.__state == ConnectionState.CONNECTED
-
-        log.info(
-            f'ConnectionManager.notify_state(): new state: {state}'
-            + ('; will retry immediately' if retry_immediately else '')
-        )
+        log.info(f'ConnectionManager.notify_state(): new state: {state}')
 
         if state == self.__state:
             return
@@ -337,14 +331,12 @@ class ConnectionManager(EventEmitter):
         self.cancel_transition_timer()
         self.check_suspend_timer(state)
 
-        if retry_immediately:
-            self.options.loop.call_soon(self.request_state, ConnectionState.CONNECTING)
-        elif state == ConnectionState.DISCONNECTED:
+        if state == ConnectionState.DISCONNECTED:
             self.start_retry_timer(self.options.disconnected_retry_timeout)
         elif state == ConnectionState.SUSPENDED:
             self.start_retry_timer(self.options.suspended_retry_timeout)
 
-        if (state == ConnectionState.DISCONNECTED and not retry_immediately) or state == ConnectionState.SUSPENDED:
+        if state == ConnectionState.DISCONNECTED or state == ConnectionState.SUSPENDED:
             self.disconnect_transport()
 
         self.enact_state_change(state, reason)
diff --git a/ably/realtime/realtime_channel.py b/ably/realtime/realtime_channel.py
index 1538c11..f832a80 100644
--- a/ably/realtime/realtime_channel.py
+++ b/ably/realtime/realtime_channel.py
@@ -1,5 +1,7 @@
 import asyncio
+from dataclasses import dataclass
 import logging
+from typing import Optional
 
 from ably.realtime.connection import ConnectionState, ProtocolMessageAction
 from ably.rest.channel import Channel
@@ -8,7 +10,7 @@ from ably.util.eventemitter import EventEmitter
 from ably.util.exceptions import AblyException
 from enum import Enum
 
-from ably.util.helper import is_callable_or_coroutine
+from ably.util.helper import Timer, is_callable_or_coroutine
 
 log = logging.getLogger(__name__)
 
@@ -19,6 +21,15 @@ class ChannelState(str, Enum):
     ATTACHED = 'attached'
     DETACHING = 'detaching'
     DETACHED = 'detached'
+    SUSPENDED = 'suspended'
+    FAILED = 'failed'
+
+
+@dataclass
+class ChannelStateChange:
+    previous: ChannelState
+    current: ChannelState
+    reason: Optional[AblyException] = None
 
 
 class RealtimeChannel(EventEmitter, Channel):
@@ -53,6 +64,12 @@ class RealtimeChannel(EventEmitter, Channel):
         self.__state = ChannelState.INITIALIZED
         self.__message_emitter = EventEmitter()
         self.__timeout_in_secs = self.__realtime.options.realtime_request_timeout / 1000
+        self.__state_timer: Timer | None = None
+
+        # Used to listen to state changes internally, if we use the public event emitter interface then internals
+        # will be disrupted if the user called .off() to remove all listeners
+        self.__internal_state_emitter = EventEmitter()
+
         Channel.__init__(self, realtime, name, {})
 
     # RTL4
@@ -82,39 +99,27 @@ class RealtimeChannel(EventEmitter, Channel):
                 status_code=400
             )
 
-        # RTL4h - wait for pending attach/detach
-        if self.state == ChannelState.ATTACHING:
-            try:
-                await self.__attach_future
-            except asyncio.CancelledError:
-                raise AblyException("Unable to attach channel due to request timeout", 504, 50003)
-            return
-        elif self.state == ChannelState.DETACHING:
-            try:
-                await self.__detach_future
-            except asyncio.CancelledError:
-                raise AblyException("Unable to detach channel due to request timeout", 504, 50003)
-            return
-
-        self.set_state(ChannelState.ATTACHING)
+        if self.state != ChannelState.ATTACHING:
+            self._request_state(ChannelState.ATTACHING)
 
         # RTL4i - wait for pending connection
         if self.__realtime.connection.state == ConnectionState.CONNECTING:
             await self.__realtime.connect()
 
-        self.__attach_future = asyncio.Future()
+        state_change = await self.__internal_state_emitter.once_async()
+
+        if state_change.current in (ChannelState.SUSPENDED, ChannelState.FAILED):
+            raise state_change.reason
+
+    def _attach_impl(self):
+        log.info("RealtimeChannel.attach_impl(): sending ATTACH protocol message")
+
         # RTL4c
-        await self.__realtime.connection.connection_manager.send_protocol_message(
-            {
-                "action": ProtocolMessageAction.ATTACH,
-                "channel": self.name,
-            }
-        )
-        try:
-            await asyncio.wait_for(self.__attach_future, self.__timeout_in_secs)  # RTL4f
-        except asyncio.TimeoutError:
-            raise AblyException("Timeout waiting for channel attach", 504, 50003)
-        self.set_state(ChannelState.ATTACHED)
+        attach_msg = {
+            "action": ProtocolMessageAction.ATTACH,
+            "channel": self.name,
+        }
+        self._send_message(attach_msg)
 
     # RTL5
     async def detach(self):
@@ -144,38 +149,38 @@ class RealtimeChannel(EventEmitter, Channel):
         if self.state in [ChannelState.INITIALIZED, ChannelState.DETACHED]:
             return
 
-        # RTL5i - wait for pending attach/detach
-        if self.state == ChannelState.DETACHING:
-            try:
-                await self.__detach_future
-            except asyncio.CancelledError:
-                raise AblyException("Unable to detach channel due to request timeout", 504, 50003)
+        if self.state == ChannelState.SUSPENDED:
+            self._notify_state(ChannelState.DETACHED)
             return
-        elif self.state == ChannelState.ATTACHING:
-            try:
-                await self.__attach_future
-            except asyncio.CancelledError:
-                raise AblyException("Unable to attach channel due to request timeout", 504, 50003)
-
-        self.set_state(ChannelState.DETACHING)
+        elif self.state == ChannelState.FAILED:
+            raise AblyException("Unable to detach; channel state = failed", 90001, 400)
+        else:
+            self._request_state(ChannelState.DETACHING)
 
         # RTL5h - wait for pending connection
         if self.__realtime.connection.state == ConnectionState.CONNECTING:
             await self.__realtime.connect()
 
-        self.__detach_future = asyncio.Future()
+        state_change = await self.__internal_state_emitter.once_async()
+        new_state = state_change.current
+
+        if new_state == ChannelState.DETACHED:
+            return
+        elif new_state == ChannelState.ATTACHING:
+            raise AblyException("Detach request superseded by a subsequent attach request", 90000, 409)
+        else:
+            raise state_change.reason
+
+    def _detach_impl(self):
+        log.info("RealtimeChannel.detach_impl(): sending DETACH protocol message")
+
         # RTL5d
-        await self.__realtime.connection.connection_manager.send_protocol_message(
-            {
-                "action": ProtocolMessageAction.DETACH,
-                "channel": self.name,
-            }
-        )
-        try:
-            await asyncio.wait_for(self.__detach_future, self.__timeout_in_secs)  # RTL5f
-        except asyncio.TimeoutError:
-            raise AblyException("Timeout waiting for channel detach", 504, 50003)
-        self.set_state(ChannelState.DETACHED)
+        detach_msg = {
+            "action": ProtocolMessageAction.DETACH,
+            "channel": self.__name,
+        }
+
+        self._send_message(detach_msg)
 
     # RTL7
     async def subscribe(self, *args):
@@ -301,21 +306,81 @@ class RealtimeChannel(EventEmitter, Channel):
     def _on_message(self, msg):
         action = msg.get('action')
         if action == ProtocolMessageAction.ATTACHED:
-            if self.__attach_future:
-                self.__attach_future.set_result(None)
-            self.__attach_future = None
+            if self.state == ChannelState.ATTACHING:
+                self._notify_state(ChannelState.ATTACHED)
+            else:
+                log.warn("RealtimeChannel._on_message(): ATTACHED received while not attaching")
         elif action == ProtocolMessageAction.DETACHED:
-            if self.__detach_future:
-                self.__detach_future.set_result(None)
-            self.__detach_future = None
+            if self.state == ChannelState.DETACHING:
+                self._notify_state(ChannelState.DETACHED)
+            else:
+                log.warn("RealtimeChannel._on_message(): DETACHED recieved while not detaching")
         elif action == ProtocolMessageAction.MESSAGE:
             messages = Message.from_encoded_array(msg.get('messages'))
             for message in messages:
                 self.__message_emitter._emit(message.name, message)
 
-    def set_state(self, state):
+    def _request_state(self, state: ChannelState):
+        log.info(f'RealtimeChannel._request_state(): state = {state}')
+        self._notify_state(state)
+        self.__check_pending_state()
+
+    def _notify_state(self, state: ChannelState, reason=None):
+        log.info(f'RealtimeChannel._notify_state(): state = {state}')
+
+        self.__clear_state_timer()
+
+        if state == self.state:
+            return
+
+        state_change = ChannelStateChange(self.__state, state, reason=reason)
+
         self.__state = state
-        self._emit(state)  # RTL2a
+        self._emit(state, state_change)
+        self.__internal_state_emitter._emit(state, state_change)
+
+    def _send_message(self, msg):
+        asyncio.create_task(self.__realtime.connection.connection_manager.send_protocol_message(msg))
+
+    def __check_pending_state(self):
+        connection_state = self.__realtime.connection.connection_manager.state
+
+        if connection_state not in (
+            ConnectionState.CONNECTING,
+            ConnectionState.CONNECTED,
+            ConnectionState.DISCONNECTED,
+        ):
+            return
+
+        if self.state == ChannelState.ATTACHING:
+            self.__start_state_timer()
+            self._attach_impl()
+        elif self.state == ChannelState.DETACHING:
+            self.__start_state_timer()
+            self._detach_impl()
+
+    def __start_state_timer(self):
+        if not self.__state_timer:
+            def on_timeout():
+                log.info('RealtimeChannel.start_state_timer(): timer expired')
+                self.__state_timer = None
+                self.__timeout_pending_state()
+
+            self.__state_timer = Timer(self.__realtime.options.realtime_request_timeout, on_timeout)
+
+    def __clear_state_timer(self):
+        if self.__state_timer:
+            self.__state_timer.cancel()
+            self.__state_timer = None
+
+    def __timeout_pending_state(self):
+        if self.state == ChannelState.ATTACHING:
+            self._notify_state(
+                ChannelState.SUSPENDED, reason=AblyException("Channel attach timed out", 408, 90007))
+        elif self.state == ChannelState.DETACHING:
+            self._notify_state(ChannelState.ATTACHED, reason=AblyException("Channel detach timed out", 408, 90007))
+        else:
+            self.__check_pending_state()
 
     # RTL23
     @property


Test patch:
diff --git a/test/ably/realtimechannel_test.py b/test/ably/realtimechannel_test.py
index b887ea3..e171e87 100644
--- a/test/ably/realtimechannel_test.py
+++ b/test/ably/realtimechannel_test.py
@@ -230,8 +230,8 @@ class TestRealtimeChannel(BaseAsyncTestCase):
         channel = ably.channels.get('channel_name')
         with pytest.raises(AblyException) as exception:
             await channel.attach()
-        assert exception.value.code == 50003
-        assert exception.value.status_code == 504
+        assert exception.value.code == 90007
+        assert exception.value.status_code == 408
         await ably.close()
 
     async def test_realtime_request_timeout_detach(self):
@@ -249,6 +249,6 @@ class TestRealtimeChannel(BaseAsyncTestCase):
         await channel.attach()
         with pytest.raises(AblyException) as exception:
             await channel.detach()
-        assert exception.value.code == 50003
-        assert exception.value.status_code == 504
+        assert exception.value.code == 90007
+        assert exception.value.status_code == 408
         await ably.close()
diff --git a/test/ably/realtimeconnection_test.py b/test/ably/realtimeconnection_test.py
index daf28f4..916c119 100644
--- a/test/ably/realtimeconnection_test.py
+++ b/test/ably/realtimeconnection_test.py
@@ -264,23 +264,3 @@ class TestRealtimeConnection(BaseAsyncTestCase):
         assert state_change.reason.status_code == 408
 
         await ably.close()
-
-    # RTN15a
-    async def test_retry_immediately_upon_unexpected_disconnection(self):
-        # Set timeouts to 500s so that if the client uses retry delay the test will fail with a timeout
-        ably = await RestSetup.get_ably_realtime(
-            disconnected_retry_timeout=500_000,
-            suspended_retry_timeout=500_000
-        )
-
-        # Wait for the client to connect
-        await ably.connection.once_async(ConnectionState.CONNECTED)
-
-        # Simulate random loss of connection
-        assert ably.connection.connection_manager.transport
-        await ably.connection.connection_manager.transport.dispose()
-        ably.connection.connection_manager.notify_state(ConnectionState.DISCONNECTED)
-        assert ably.connection.state == ConnectionState.DISCONNECTED
-
-        # Wait for the client to connect again
-        await ably.connection.once_async(ConnectionState.CONNECTED)