Fixed loss of real-time subscriptions occurred after TzKT API outage (#163)

* Fixed loss of real-time subscriptions occurred after TzKT API outage

* Changelog

Author: droserasprout

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 3.1.1 - [unreleased]
+
+### Fixed
+
+* Fixed loss of real-time subscriptions occurred after TzKT API outage.
+* Fixed updating schema hash in `schema approve` command.
+* Fixed possible crash occurred while Hasura is not ready.
+
 ## 3.1.0 - 2021-10-12
 
 ### Added

src/dipdup/cli.py

@@ -27,7 +27,7 @@
 from dipdup.migrations import DipDupMigrationManager, deprecated_handlers
 from dipdup.models import Schema
 from dipdup.utils import iter_files
-from dipdup.utils.database import set_decimal_context, tortoise_wrapper, wipe_schema
+from dipdup.utils.database import get_schema_hash, set_decimal_context, tortoise_wrapper, wipe_schema
 
 _logger = logging.getLogger('dipdup.cli')
 
@@ -238,6 +238,9 @@ async def schema(ctx):
     ...
 
 
+from dipdup.enums import ReindexingReason
+
+
 @schema.command(name='approve', help='Continue indexing with the same schema after crashing with `ReindexingRequiredError`')
 @click.pass_context
 @cli_wrapper
@@ -247,7 +250,15 @@ async def schema_approve(ctx):
     models = f'{config.package}.models'
 
     async with tortoise_wrapper(url, models):
-        await Schema.filter().update(reindex=None)
+        schema = await Schema.filter().get()
+        if not schema.reindex:
+            return
+
+        if schema.reindex == ReindexingReason.SCHEMA_HASH_MISMATCH:
+            conn = get_connection(None)
+            schema.hash = get_schema_hash(conn)
+        schema.reindex = None
+        await schema.save()
 
 
 @schema.command(name='wipe', help='Drop all database tables, functions and views')

src/dipdup/datasources/tzkt/datasource.py

@@ -31,6 +31,7 @@
 from dipdup.enums import MessageType
 from dipdup.models import BigMapAction, BigMapData, BlockData, HeadBlockData, OperationData, QuoteData
 from dipdup.utils import split_by_chunks
+from dipdup.utils.watchdog import Watchdog
 
 TZKT_ORIGINATIONS_REQUEST_LIMIT = 100
 
@@ -266,9 +267,11 @@ def __init__(
         self,
         url: str,
         http_config: Optional[HTTPConfig] = None,
+        watchdog: Optional[Watchdog] = None,
     ) -> None:
         super().__init__(url, self._default_http_config.merge(http_config))
         self._logger = logging.getLogger('dipdup.tzkt')
+        self._watchdog = watchdog
 
         self._transaction_subscriptions: Set[str] = set()
         self._origination_subscriptions: bool = False
@@ -546,11 +549,13 @@ def _get_ws_client(self) -> BaseHubConnection:
         return self._ws_client
 
     async def run(self) -> None:
-        """Main loop. Sync indexes via REST, start WS connection"""
-        self._logger.info('Starting datasource')
+        self._logger.info('Establishing realtime connection')
+        tasks = [asyncio.create_task(self._get_ws_client().start())]
 
-        self._logger.info('Starting websocket client')
-        await self._get_ws_client().start()
+        if self._watchdog:
+            tasks.append(asyncio.create_task(self._watchdog.run()))
+
+        await asyncio.gather(*tasks)
 
     async def _on_connect(self) -> None:
         """Subscribe to all required channels on established WS connection"""
@@ -680,6 +685,9 @@ async def _on_big_maps_message(self, message: List[Dict[str, Any]]) -> None:
     async def _on_head_message(self, message: List[Dict[str, Any]]) -> None:
         """Parse and emit raw head block from WS"""
         async for data in self._extract_message_data(MessageType.head, message):
+            if self._watchdog:
+                self._watchdog.reset()
+
             block = self.convert_head_block(data)
             await self.emit_head(block)
 

src/dipdup/dipdup.py

@@ -38,6 +38,7 @@
 from dipdup.scheduler import add_job, create_scheduler
 from dipdup.utils import slowdown
 from dipdup.utils.database import generate_schema, get_schema_hash, prepare_models, set_schema, tortoise_wrapper, validate_models
+from dipdup.utils.watchdog import Watchdog
 
 
 class IndexDispatcher:
@@ -293,6 +294,7 @@ async def _create_datasources(self) -> None:
                 datasource = TzktDatasource(
                     url=datasource_config.url,
                     http_config=datasource_config.http,
+                    watchdog=Watchdog(120.0, self._ctx.restart),
                 )
             elif isinstance(datasource_config, BcdDatasourceConfig):
                 datasource = BcdDatasource(

src/dipdup/hasura.py

@@ -8,7 +8,7 @@
 from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple
 
 import humps  # type: ignore
-from aiohttp import ClientConnectorError, ClientOSError
+from aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError
 from pydantic.dataclasses import dataclass
 from tortoise import fields
 from tortoise.transactions import get_connection
@@ -145,7 +145,7 @@ async def _healthcheck(self) -> None:
         self._logger.info('Waiting for Hasura instance to be ready')
         timeout = self._http_config.connection_timeout or 60
         for _ in range(timeout):
-            with suppress(ClientConnectorError, ClientOSError):
+            with suppress(ClientConnectorError, ClientOSError, ServerDisconnectedError):
                 response = await self._http._session.get(f'{self._hasura_config.url}/healthz')
                 if response.status == 200:
                     break

src/dipdup/utils/watchdog.py

@@ -0,0 +1,26 @@
+import asyncio
+import logging
+from typing import Awaitable, Callable
+
+
+class Watchdog:
+    """Simple keepalive task."""
+
+    def __init__(self, timeout: float, fn: Callable[..., Awaitable[None]]) -> None:
+        self._logger = logging.getLogger('watchdog')
+        self._timeout = timeout
+        self._fn = fn
+        self._event = asyncio.Event()
+
+    async def run(self) -> None:
+        while True:
+            try:
+                await asyncio.wait_for(self._event.wait(), self._timeout)
+                await asyncio.sleep(0)
+            except asyncio.TimeoutError:
+                self._logger.critical('Watchdog has timed out.')
+                await self._fn()
+
+    def reset(self) -> None:
+        self._event.set()
+        self._event.clear()