Ability to process realtime messages with lag (#275)

* WIP

* wip

* Bump hasura, POC buffer

* Fix config validators

* Fix emitted levels

* Better logging

* Changelog

* Docs

Author: droserasprout

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## [unreleased]
+
+### Added
+
+* tzkt: Ability to process realtime messages with lag.
+
 ## 4.2.7 - 2022-03-02
 
 ### Fixed

src/dipdup/config.py

@@ -211,11 +211,13 @@ class TzktDatasourceConfig(NameMixin):
     :param kind: always 'tzkt'
     :param url: Base API URL, e.g. https://api.tzkt.io/
     :param http: HTTP client configuration
+    :param buffer_size: Number of levels to keep in FIFO buffer before processing
     """
 
     kind: Literal['tzkt']
     url: str
     http: Optional[HTTPConfig] = None
+    buffer_size: int = 0
 
     def __hash__(self) -> int:
         return hash(self.kind + self.url)
@@ -737,8 +739,10 @@ def hash(self) -> str:
         # FIXME: How to convert pydantic dataclass into dict without json.dumps? asdict is not recursive.
         config_dict = json.loads(config_json)
 
-        # NOTE: We need to preserve datasource URL but remove it's HTTP tunables to avoid false-positives.
+        # NOTE: We need to preserve datasource URL but remove its HTTP tunables to avoid false-positives.
         config_dict['datasource'].pop('http', None)
+        # NOTE: TzKT tunable
+        config_dict['datasource'].pop('buffer_size', None)
         # NOTE: Same for BigMapIndex tunables
         config_dict.pop('skip_history', None)
 

src/dipdup/datasources/factory.py

@@ -32,6 +32,7 @@ def _build_datasource(cls, name: str, config: DipDupConfig) -> Datasource:
                 url=datasource_config.url,
                 http_config=datasource_config.http,
                 merge_subscriptions=config.advanced.merge_subscriptions,
+                buffer_size=datasource_config.buffer_size,
             )
 
         if isinstance(datasource_config, CoinbaseDatasourceConfig):

src/dipdup/datasources/tzkt/datasource.py

@@ -274,6 +274,7 @@ def __init__(
         http_config: Optional[HTTPConfig] = None,
         watchdog: Optional[Watchdog] = None,
         merge_subscriptions: bool = False,
+        buffer_size: int = 0,
     ) -> None:
         super().__init__(
             url=url,
@@ -282,6 +283,8 @@ def __init__(
         )
         self._logger = logging.getLogger('dipdup.tzkt')
         self._watchdog = watchdog
+        self._buffer_size = buffer_size
+        self._buffer: DefaultDict[int, List[Tuple[MessageType, Dict[str, Any]]]] = defaultdict(list)
 
         self._ws_client: Optional[SignalRClient] = None
         self._level: DefaultDict[MessageType, Optional[int]] = defaultdict(lambda: None)
@@ -785,20 +788,28 @@ async def _on_error(self, message: CompletionMessage) -> NoReturn:
 
     async def _extract_message_data(self, type_: MessageType, message: List[Any]) -> AsyncGenerator[Dict, None]:
         """Parse message received from Websocket, ensure it's correct in the current context and yield data."""
+        # NOTE: Parse messages and either buffer or yield data
         for item in message:
             tzkt_type = TzktMessageType(item['type'])
+            # NOTE: Legacy, sync level returned by TzKT during negotiation
             if tzkt_type == TzktMessageType.STATE:
                 continue
 
-            level, current_level = item['state'], self._level[type_]
-            self._level[type_] = level
-            self._logger.info('Realtime message received: %s, %s, %s -> %s', type_.value, tzkt_type.name, current_level, level)
+            message_level, current_level = item['state'], self._level[type_]
+            self._level[type_] = message_level
+            self._logger.info(
+                'Realtime message received: %s, %s, %s -> %s',
+                type_.value,
+                tzkt_type.name,
+                current_level,
+                message_level,
+            )
 
-            # NOTE: Just yield data
+            # NOTE: Put data messages to buffer by level
             if tzkt_type == TzktMessageType.DATA:
-                yield item['data']
+                self._buffer[message_level].append((type_, item['data']))
 
-            # NOTE: Emit rollback, but not on `head` message
+            # NOTE: Try to process rollback automatically, emit if failed
             elif tzkt_type == TzktMessageType.REORG:
                 # NOTE: operation/big_map channels have their own levels
                 if type_ == MessageType.head:
@@ -812,15 +823,39 @@ async def _extract_message_data(self, type_: MessageType, message: List[Any]) ->
                         raise RuntimeError('Reorg message received, but neither current nor sync level is known')
 
                     # NOTE: This rollback does not affect us, so we can safely ignore it
-                    if current_level <= level:
+                    if current_level <= message_level:
                         return
 
-                self._logger.info('Emitting rollback from %s to %s', current_level, level)
-                await self.emit_rollback(current_level, level)
+                self._logger.info('Rollback requested from %s to %s', current_level, message_level)
+
+                # NOTE: Drop buffered messages in reversed order while possible
+                rolled_back_levels = range(current_level, message_level, -1)
+                for rolled_back_level in rolled_back_levels:
+                    if self._buffer.pop(rolled_back_level, None):
+                        self._logger.info('Level %s is buffered', rolled_back_level)
+                    else:
+                        self._logger.info('Level %s is not buffered, emitting rollback', rolled_back_level)
+                        await self.emit_rollback(current_level, message_level)
+                        return
+
+                self._logger.info('Rollback is not required, continuing')
 
             else:
                 raise NotImplementedError
 
+        # NOTE: Yield extensive data from buffer
+        buffered_levels = sorted(self._buffer.keys())
+        emitted_levels = buffered_levels[: len(buffered_levels) - self._buffer_size]
+
+        for level in emitted_levels:
+            for idx, level_data in enumerate(self._buffer[level]):
+                level_message_type, level_message = level_data
+                if level_message_type == type_:
+                    yield level_message
+                    self._buffer[level].pop(idx)
+            if not self._buffer[level]:
+                self._buffer.pop(level)
+
     async def _on_operations_message(self, message: List[Dict[str, Any]]) -> None:
         """Parse and emit raw operations from WS"""
         level_operations: DefaultDict[int, Deque[OperationData]] = defaultdict(deque)