Backport: Fix issue with processing rollbacks while sync is in progress (#832)

Author: droserasprout

CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog], and this project adheres to [Semantic Versioning].
 
+## [Unreleased]
+
+### Fixed
+
+- tzkt: Fixed issue with processing rollbacks while sync is in progress.
+
 ## [6.5.11] - 2023-09-02
 
 ### Fixed

src/dipdup/datasources/tzkt/datasource.py

@@ -43,6 +43,7 @@
 from dipdup.models import OperationData
 from dipdup.models import QuoteData
 from dipdup.models import TokenTransferData
+from dipdup.models import TzktRollbackMessage
 from dipdup.utils import FormattedLogger
 from dipdup.utils import split_by_chunks
 
@@ -119,7 +120,7 @@ class TzktMessageType(Enum):
     REORG = 2
 
 
-MessageData = dict[str, Any] | list[dict[str, Any]]
+MessageData = dict[str, Any] | list[dict[str, Any]] | TzktRollbackMessage
 
 
 class BufferedMessage(NamedTuple):

src/dipdup/dipdup.py

@@ -55,6 +55,7 @@
 from dipdup.models import OperationData
 from dipdup.models import Schema
 from dipdup.models import TokenTransferData
+from dipdup.models import TzktRollbackMessage
 from dipdup.prometheus import Metrics
 from dipdup.scheduler import SchedulerManager
 from dipdup.transactions import TransactionManager
@@ -300,7 +301,11 @@ async def _on_events(self, datasource: IndexDatasource, events: tuple[EventData,
                 index.push_events(events)
 
     async def _on_rollback(
-        self, datasource: IndexDatasource, type_: MessageType, from_level: int, to_level: int
+        self,
+        datasource: IndexDatasource,
+        type_: MessageType,
+        from_level: int,
+        to_level: int,
     ) -> None:
         """Call `on_index_rollback` hook for each index that is affected by rollback"""
         if from_level <= to_level:
@@ -312,8 +317,6 @@ async def _on_rollback(
             Metrics.set_datasource_rollback(datasource.name)
 
         # NOTE: Choose action for each index
-        affected_indexes: set[str] = set()
-
         for index_name, index in self._indexes.items():
             index_level = index.state.level
 
@@ -328,19 +331,11 @@ async def _on_rollback(
 
             else:
                 self._logger.debug('%s: affected', index_name)
-                affected_indexes.add(index_name)
-
-        hook_name = 'on_index_rollback'
-        for index_name in affected_indexes:
-            self._logger.warning('`%s` index is affected by rollback; firing `%s` hook', index_name, hook_name)
-            await self._ctx.fire_hook(
-                hook_name,
-                index=self._indexes[index_name],
-                from_level=from_level,
-                to_level=to_level,
-            )
+                index.push_realtime_message(
+                    TzktRollbackMessage(from_level, to_level),
+                )
 
-        self._logger.info('`%s` rollback complete', channel)
+        self._logger.info('`%s` rollback processed', channel)
 
 
 class DipDup:

src/dipdup/index.py

@@ -196,3 +196,18 @@ async def _update_state(
         state.status = status or state.status
         state.level = level or state.level
         await state.save()
+
+    # TODO: Move to TezosTzktIndex
+    async def _tzkt_rollback(
+        self,
+        from_level: int,
+        to_level: int,
+    ) -> None:
+        hook_name = 'on_index_rollback'
+        self._logger.warning('Affected by rollback; firing `%s` hook', self.name, hook_name)
+        await self._ctx.fire_hook(
+            name=hook_name,
+            index=self,
+            from_level=from_level,
+            to_level=to_level,
+        )

src/dipdup/indexes/big_map/index.py

@@ -17,9 +17,10 @@
 from dipdup.models import BigMapAction
 from dipdup.models import BigMapData
 from dipdup.models import BigMapDiff
+from dipdup.models import TzktRollbackMessage
 from dipdup.prometheus import Metrics
 
-BigMapQueueItem = tuple[BigMapData, ...]
+BigMapQueueItem = tuple[BigMapData, ...] | TzktRollbackMessage
 
 
 class BigMapIndex(
@@ -35,16 +36,20 @@ async def _process_queue(self) -> None:
         if self._queue:
             self._logger.debug('Processing websocket queue')
         while self._queue:
-            big_maps = self._queue.popleft()
-            message_level = big_maps[0].level
+            message = self._queue.popleft()
+            if isinstance(message, TzktRollbackMessage):
+                await self._tzkt_rollback(message.from_level, message.to_level)
+                continue
+
+            message_level = message[0].level
             if message_level <= self.state.level:
                 self._logger.debug('Skipping outdated message: %s <= %s', message_level, self.state.level)
                 continue
 
             with ExitStack() as stack:
                 if Metrics.enabled:
                     stack.enter_context(Metrics.measure_level_realtime_duration())
-                await self._process_level_big_maps(big_maps, message_level)
+                await self._process_level_big_maps(message, message_level)
 
     async def _synchronize(self, sync_level: int) -> None:
         """Fetch operations via Fetcher and pass to message callback"""

src/dipdup/indexes/event/index.py

@@ -13,10 +13,11 @@
 from dipdup.indexes.event.matcher import match_events
 from dipdup.models import Event
 from dipdup.models import EventData
+from dipdup.models import TzktRollbackMessage
 from dipdup.models import UnknownEvent
 from dipdup.prometheus import Metrics
 
-EventQueueItem = tuple[EventData, ...]
+EventQueueItem = tuple[EventData, ...] | TzktRollbackMessage
 
 
 class EventIndex(
@@ -31,16 +32,20 @@ async def _process_queue(self) -> None:
         if self._queue:
             self._logger.debug('Processing websocket queue')
         while self._queue:
-            events = self._queue.popleft()
-            message_level = events[0].level
+            message = self._queue.popleft()
+            if isinstance(message, TzktRollbackMessage):
+                await self._tzkt_rollback(message.from_level, message.to_level)
+                continue
+
+            message_level = message[0].level
             if message_level <= self.state.level:
                 self._logger.debug('Skipping outdated message: %s <= %s', message_level, self.state.level)
                 continue
 
             with ExitStack() as stack:
                 if Metrics.enabled:
                     stack.enter_context(Metrics.measure_level_realtime_duration())
-                await self._process_level_events(events, message_level)
+                await self._process_level_events(message, message_level)
 
     def _create_fetcher(self, first_level: int, last_level: int) -> EventFetcher:
         event_addresses = self._get_event_addresses()

src/dipdup/indexes/head/index.py

@@ -7,8 +7,9 @@
 from dipdup.exceptions import FrameworkException
 from dipdup.index import Index
 from dipdup.models import HeadBlockData
+from dipdup.models import TzktRollbackMessage
 
-HeadQueueItem = HeadBlockData
+HeadQueueItem = HeadBlockData | TzktRollbackMessage
 
 
 class HeadIndex(
@@ -24,23 +25,27 @@ async def _synchronize(self, sync_level: int) -> None:
 
     async def _process_queue(self) -> None:
         while self._queue:
-            head = self._queue.popleft()
-            message_level = head.level
+            message = self._queue.popleft()
+            if isinstance(message, TzktRollbackMessage):
+                await self._tzkt_rollback(message.from_level, message.to_level)
+                continue
+
+            message_level = message.level
             if message_level <= self.state.level:
                 self._logger.debug('Skipping outdated message: %s <= %s', message_level, self.state.level)
                 continue
 
             self._logger.debug('Processing head realtime message, %s left in queue', len(self._queue))
 
-            batch_level = head.level
+            batch_level = message.level
             index_level = self.state.level
             if batch_level <= index_level:
                 raise FrameworkException(f'Batch level is lower than index level: {batch_level} <= {index_level}')
 
             async with self._ctx._transactions.in_transaction(batch_level, message_level, self.name):
                 self._logger.debug('Processing head info of level %s', batch_level)
                 for handler_config in self._config.handlers:
-                    await self._call_matched_handler(handler_config, head)
+                    await self._call_matched_handler(handler_config, message)
                 await self._update_state(level=batch_level)
 
     async def _call_matched_handler(self, handler_config: HeadHandlerConfig, head: HeadBlockData) -> None:

src/dipdup/indexes/operation/index.py

@@ -27,11 +27,12 @@
 from dipdup.indexes.operation.matcher import match_operation_subgroup
 from dipdup.indexes.operation.matcher import match_operation_unfiltered_subgroup
 from dipdup.models import OperationData
+from dipdup.models import TzktRollbackMessage
 from dipdup.prometheus import Metrics
 
 _logger = logging.getLogger('dipdup.matcher')
 
-OperationQueueItem = tuple[OperationSubgroup, ...]
+OperationQueueItem = tuple[OperationSubgroup, ...] | TzktRollbackMessage
 
 
 def entrypoint_filter(handlers: tuple[OperationHandlerConfig, ...]) -> set[str | None]:
@@ -190,12 +191,12 @@ async def _process_queue(self) -> None:
 
         while self._queue:
             message = self._queue.popleft()
-            messages_left = len(self._queue)
-
-            if not message:
-                raise FrameworkException('Got empty message from realtime queue')
+            if isinstance(message, TzktRollbackMessage):
+                await self._tzkt_rollback(message.from_level, message.to_level)
+                continue
 
             if Metrics.enabled:
+                messages_left = len(self._queue)
                 Metrics.set_levels_to_realtime(self._config.name, messages_left)
 
             message_level = message[0].operations[0].level

src/dipdup/indexes/token_transfer/index.py

@@ -10,9 +10,10 @@
 from dipdup.indexes.token_transfer.fetcher import TokenTransferFetcher
 from dipdup.indexes.token_transfer.matcher import match_token_transfers
 from dipdup.models import TokenTransferData
+from dipdup.models import TzktRollbackMessage
 from dipdup.prometheus import Metrics
 
-TokenTransferQueueItem = tuple[TokenTransferData, ...]
+TokenTransferQueueItem = tuple[TokenTransferData, ...] | TzktRollbackMessage
 
 
 class TokenTransferIndex(
@@ -115,13 +116,17 @@ async def _process_queue(self) -> None:
         if self._queue:
             self._logger.debug('Processing websocket queue')
         while self._queue:
-            token_transfers = self._queue.popleft()
-            message_level = token_transfers[0].level
+            message = self._queue.popleft()
+            if isinstance(message, TzktRollbackMessage):
+                await self._tzkt_rollback(message.from_level, message.to_level)
+                continue
+
+            message_level = message[0].level
             if message_level <= self.state.level:
                 self._logger.debug('Skipping outdated message: %s <= %s', message_level, self.state.level)
                 continue
 
             with ExitStack() as stack:
                 if Metrics.enabled:
                     stack.enter_context(Metrics.measure_level_realtime_duration())
-                await self._process_level_token_transfers(token_transfers, message_level)
+                await self._process_level_token_transfers(message, message_level)

src/dipdup/models.py

@@ -264,6 +264,12 @@ class UnknownEvent:
     payload: Any | None
 
 
+@dataclass(frozen=True)
+class TzktRollbackMessage:
+    from_level: int
+    to_level: int
+
+
 # ===> Model Versioning