Fix crash with RuntimeError after continuous realtime connection loss (#366)

* Fix crash with `RuntimeError` after continuous realtime connection loss

* Less agressive backoff

Author: droserasprout

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Fixed
 
 * index: Allow mixing oneshot and regular indexes in a single config.
+* index: Fixed crash with `RuntimeError` after continuous realtime connection loss.
 * index: Fixed `OperationIndexConfig.types` field being partially ignored.
 * tzkt: Fixed `origination` subscription missing when `merge_subscriptions` flag is set.
 

src/dipdup/datasources/datasource.py

@@ -22,6 +22,7 @@
 _logger = logging.getLogger('dipdup.datasource')
 
 
+EmptyCallbackT = Callable[[], Awaitable[None]]
 HeadCallbackT = Callable[['IndexDatasource', HeadBlockData], Awaitable[None]]
 OperationsCallbackT = Callable[['IndexDatasource', Tuple[OperationData, ...]], Awaitable[None]]
 BigMapsCallbackT = Callable[['IndexDatasource', Tuple[BigMapData, ...]], Awaitable[None]]
@@ -69,10 +70,12 @@ class GraphQLDatasource(Datasource):
 class IndexDatasource(Datasource):
     def __init__(self, url: str, http_config: HTTPConfig, merge_subscriptions: bool = False) -> None:
         super().__init__(url, http_config)
-        self._on_head: Set[HeadCallbackT] = set()
-        self._on_operations: Set[OperationsCallbackT] = set()
-        self._on_big_maps: Set[BigMapsCallbackT] = set()
-        self._on_rollback: Set[RollbackCallbackT] = set()
+        self._on_connected_callbacks: Set[EmptyCallbackT] = set()
+        self._on_disconnected_callbacks: Set[EmptyCallbackT] = set()
+        self._on_head_callbacks: Set[HeadCallbackT] = set()
+        self._on_operations_callbacks: Set[OperationsCallbackT] = set()
+        self._on_big_maps_callbacks: Set[BigMapsCallbackT] = set()
+        self._on_rollback_callbacks: Set[RollbackCallbackT] = set()
         self._subscriptions: SubscriptionManager = SubscriptionManager(merge_subscriptions)
         self._subscriptions.add(HeadSubscription())
         self._network: Optional[str] = None
@@ -91,34 +94,48 @@ def network(self) -> str:
     async def subscribe(self) -> None:
         ...
 
-    def on_head(self, fn: HeadCallbackT) -> None:
-        self._on_head.add(fn)
+    def call_on_head(self, fn: HeadCallbackT) -> None:
+        self._on_head_callbacks.add(fn)
 
-    def on_operations(self, fn: OperationsCallbackT) -> None:
-        self._on_operations.add(fn)
+    def call_on_operations(self, fn: OperationsCallbackT) -> None:
+        self._on_operations_callbacks.add(fn)
 
-    def on_big_maps(self, fn: BigMapsCallbackT) -> None:
-        self._on_big_maps.add(fn)
+    def call_on_big_maps(self, fn: BigMapsCallbackT) -> None:
+        self._on_big_maps_callbacks.add(fn)
 
-    def on_rollback(self, fn: RollbackCallbackT) -> None:
-        self._on_rollback.add(fn)
+    def call_on_rollback(self, fn: RollbackCallbackT) -> None:
+        self._on_rollback_callbacks.add(fn)
+
+    def call_on_connected(self, fn: EmptyCallbackT) -> None:
+        self._on_connected_callbacks.add(fn)
+
+    def call_on_disconnected(self, fn: EmptyCallbackT) -> None:
+        self._on_disconnected_callbacks.add(fn)
 
     async def emit_head(self, head: HeadBlockData) -> None:
-        for fn in self._on_head:
+        for fn in self._on_head_callbacks:
             await fn(self, head)
 
     async def emit_operations(self, operations: Tuple[OperationData, ...]) -> None:
-        for fn in self._on_operations:
+        for fn in self._on_operations_callbacks:
             await fn(self, operations)
 
     async def emit_big_maps(self, big_maps: Tuple[BigMapData, ...]) -> None:
-        for fn in self._on_big_maps:
+        for fn in self._on_big_maps_callbacks:
             await fn(self, big_maps)
 
     async def emit_rollback(self, type_: MessageType, from_level: int, to_level: int) -> None:
-        for fn in self._on_rollback:
+        for fn in self._on_rollback_callbacks:
             await fn(self, type_, from_level, to_level)
 
+    async def emit_connected(self) -> None:
+        for fn in self._on_connected_callbacks:
+            await fn()
+
+    async def emit_disconnected(self) -> None:
+        for fn in self._on_disconnected_callbacks:
+            await fn()
+
     def set_network(self, network: str) -> None:
         if self._network:
             raise RuntimeError('Network is already set')

src/dipdup/datasources/tzkt/datasource.py

@@ -359,7 +359,7 @@ class TzktDatasource(IndexDatasource):
     _default_http_config = HTTPConfig(
         cache=True,
         retry_sleep=1,
-        retry_multiplier=2,
+        retry_multiplier=1.1,
         ratelimit_rate=100,
         ratelimit_period=1,
         connection_limit=25,
@@ -894,8 +894,8 @@ def _get_ws_client(self) -> SignalRClient:
             max_size=None,
         )
 
-        self._ws_client.on_open(self._on_connect)
-        self._ws_client.on_close(self._on_disconnect)
+        self._ws_client.on_open(self._on_connected)
+        self._ws_client.on_close(self._on_disconnected)
         self._ws_client.on_error(self._on_error)
 
         self._ws_client.on('operations', partial(self._on_message, MessageType.operation))
@@ -915,6 +915,7 @@ async def _wrapper():
                     await ws.run()
                 except WebsocketConnectionError as e:
                     self._logger.error('Websocket connection error: %s', e)
+                    await self.emit_disconnected()
                     await asyncio.sleep(retry_sleep)
                     retry_sleep *= self._http_config.retry_multiplier
 
@@ -925,13 +926,15 @@ async def _wrapper():
 
         await gather(*tasks)
 
-    async def _on_connect(self) -> None:
+    async def _on_connected(self) -> None:
         self._logger.info('Realtime connection established')
         # NOTE: Subscribing here will block WebSocket loop
+        await self.emit_connected()
 
-    async def _on_disconnect(self) -> None:
+    async def _on_disconnected(self) -> None:
         self._logger.info('Realtime connection lost')
         self._subscriptions.reset()
+        await self.emit_disconnected()
 
     async def _on_error(self, message: CompletionMessage) -> NoReturn:
         """Raise exception from WS server's error message"""

src/dipdup/dipdup.py

@@ -172,15 +172,6 @@ async def _fetch_contracts(self) -> None:
                 self._ctx.config.contracts[contract.name] = contract_config
         self._ctx.config.initialize(skip_imports=True)
 
-    async def _subscribe_to_datasource_events(self) -> None:
-        for datasource in self._ctx.datasources.values():
-            if not isinstance(datasource, IndexDatasource):
-                continue
-            datasource.on_head(self._on_head)
-            datasource.on_operations(self._on_operations)
-            datasource.on_big_maps(self._on_big_maps)
-            datasource.on_rollback(self._on_rollback)
-
     async def _load_index_states(self) -> None:
         if self._indexes:
             raise RuntimeError('Index states are already loaded')
@@ -224,6 +215,22 @@ async def _process(index_state: IndexState) -> None:
         tasks = (create_task(_process(index_state)) for index_state in await IndexState.all())
         await gather(*tasks)
 
+    async def _subscribe_to_datasource_events(self) -> None:
+        for datasource in self._ctx.datasources.values():
+            if not isinstance(datasource, IndexDatasource):
+                continue
+            datasource.call_on_disconnected(self._on_disconnected)
+            datasource.call_on_head(self._on_head)
+            datasource.call_on_operations(self._on_operations)
+            datasource.call_on_big_maps(self._on_big_maps)
+            datasource.call_on_rollback(self._on_rollback)
+
+    async def _on_disconnected(self) -> None:
+        # NOTE: Invalidate realtime queues; sync level will be reset
+        self._logger.info('Datasource disconnected, dropping realtime queues')
+        for index in self._indexes.values():
+            index._queue.clear()
+
     async def _on_head(self, datasource: IndexDatasource, head: HeadBlockData) -> None:
         # NOTE: Do not await query results - blocked database connection may cause Websocket timeout.
         self._tasks.append(

tests/test_dipdup/test_datasources/test_tzkt.py

@@ -213,7 +213,7 @@ async def test_on_operation_message_data(self) -> None:
         message = {'type': 1, 'state': 2, 'data': operations_json}
         async with with_tzkt(1) as tzkt:
             emit_mock = AsyncMock()
-            tzkt.on_operations(emit_mock)
+            tzkt.call_on_operations(emit_mock)
             tzkt.set_sync_level(HeadSubscription(), 1)
 
             level = tzkt.get_channel_level(MessageType.operation)