Defer spawning index datasources, add get_quotes method (#133)

* Defer spawning index datasources

* get_quotes

* Limit

* Logs, runtimeerror fixes

* Docs

* Set index sync levels early

* Fix head check in `process`

* Lint

Author: droserasprout

src/dipdup/datasources/tzkt/datasource.py

@@ -537,7 +537,7 @@ async def get_big_maps(
             big_maps.append(self.convert_big_map(bm))
         return big_maps
 
-    async def get_quotes(self, level: int) -> QuoteData:
+    async def get_quote(self, level: int) -> QuoteData:
         """Get quote for block"""
         self._logger.info('Fetching quotes for level %s', level)
         quote_json = await self._http.request(
@@ -548,6 +548,21 @@ async def get_quotes(self, level: int) -> QuoteData:
         )
         return self.convert_quote(quote_json[0])
 
+    async def get_quotes(self, from_level: int, to_level: int) -> List[QuoteData]:
+        """Get quotes for blocks"""
+        self._logger.info('Fetching quotes for levels %s-%s', from_level, to_level)
+        quotes_json = await self._http.request(
+            'get',
+            url='v1/quotes',
+            params={
+                "level.ge": from_level,
+                "level.lt": to_level,
+                "limit": self.request_limit,
+            },
+            cache=False,
+        )
+        return [self.convert_quote(quote) for quote in quotes_json]
+
     async def add_index(self, index_config: ResolvedIndexConfigT) -> None:
         """Register index config in internal mappings and matchers. Find and register subscriptions."""
 
@@ -680,9 +695,9 @@ async def _extract_message_data(self, channel: str, message: List[Any]) -> Async
 
             head_level = item['state']
             if self._level and head_level < self._level:
-                raise RuntimeError('Received data message from level lower than current: {head_level} < {self._level}')
+                raise RuntimeError(f'Received data message from level lower than current: {head_level} < {self._level}')
 
-            # NOTE: State messages will be replaced with negotiation some day
+            # NOTE: State messages will be replaced with WS negotiation some day
             if message_type == TzktMessageType.STATE:
                 if self._sync_level != head_level:
                     self._logger.info('Datasource level set to %s', head_level)
@@ -732,7 +747,6 @@ async def _on_head_message(self, message: List[Dict[str, Any]]) -> None:
             created = False
             if self._head is None:
                 self._head, created = await Head.get_or_create(
-                    # NOTE: It would be better to use datasource name but it's not available
                     name=self._http._url,
                     defaults=dict(
                         level=block.level,
@@ -748,6 +762,31 @@ async def _on_head_message(self, message: List[Dict[str, Any]]) -> None:
 
             self.emit_head(block)
 
+    # FIXME: I don't like this approach, too hacky.
+    async def set_head_from_http(self) -> None:
+        """Set block from `get_head_block` HTTP method for indexes to use the same level during initial sync"""
+        if self._head:
+            self._logger.warning('Attempt to set head twice')
+            return
+        block = await self.get_head_block()
+        self._head, created = await Head.get_or_create(
+            name=self._http._url,
+            defaults=dict(
+                level=block.level,
+                hash=block.hash,
+                timestamp=block.timestamp,
+            ),
+        )
+        if not created:
+            self._head.level = block.level  # type: ignore
+            self._head.hash = block.hash  # type: ignore
+            self._head.timestamp = block.timestamp  # type: ignore
+            await self._head.save()
+
+        self._logger.info('Datasource head set to block with level %s', self._head.level)
+
+        # NOTE: No need to emit?
+
     @classmethod
     def convert_operation(cls, operation_json: Dict[str, Any]) -> OperationData:
         """Convert raw operation message from WS/REST into dataclass"""

src/dipdup/dipdup.py

@@ -45,41 +45,45 @@ def __init__(self, ctx: DipDupContext) -> None:
         self._indexes: Dict[str, Index] = {}
         self._contracts: Set[ContractConfig] = set()
         self._stopped: bool = False
-        self._synchronized: bool = False
 
-    async def run(self) -> None:
+    async def run(self, spawn_datasources_event: Optional[Event]) -> None:
         self._logger.info('Starting index dispatcher')
-        await self._subscribe_to_datasources()
+        await self._subscribe_to_datasource_events()
+        await self._set_datasource_heads()
         await self._load_index_states()
 
         while not self._stopped:
-            with suppress(IndexError):
-                while index := pending_indexes.pop():
-                    self._indexes[index._config.name] = index
-
             tasks = [index.process() for index in self._indexes.values()]
             async with slowdown(1.0):
                 await gather(*tasks)
 
             await self._check_states()
 
+            indexes_spawned = False
+            with suppress(IndexError):
+                while index := pending_indexes.popleft():
+                    self._indexes[index._config.name] = index
+                    indexes_spawned = True
+            if not indexes_spawned:
+                await self._check_states()
+
+                if spawn_datasources_event:
+                    spawn_datasources_event.set()
+
     def stop(self) -> None:
         self._stopped = True
 
     async def _check_states(self) -> None:
         statuses = [i.state.status for i in self._indexes.values()]
 
-        def _have_no(status: IndexStatus) -> bool:
+        def _every_index_is(status: IndexStatus) -> bool:
             nonlocal statuses
-            return bool(tuple(filter(partial(ne, status), statuses)))
+            return bool(statuses) and not bool(tuple(filter(partial(ne, status), statuses)))
 
-        synching_indexes = _have_no(IndexStatus.REALTIME)
-        if not synching_indexes:
-            # TODO: `on_synchronized` hook? Not sure if we need it.
-            ...
+        # TODO: `on_synchronized` hook? Not sure if we need it.
+        # if _every_index_is(IndexStatus.REALTIME): ...
 
-        pending_oneshot_indexes = _have_no(IndexStatus.ONESHOT)
-        if not pending_oneshot_indexes:
+        if _every_index_is(IndexStatus.ONESHOT):
             self.stop()
 
     async def _fetch_contracts(self) -> None:
@@ -93,14 +97,19 @@ async def _fetch_contracts(self) -> None:
                 self._ctx.config.contracts[contract.name] = contract_config
         self._ctx.config.pre_initialize()
 
-    async def _subscribe_to_datasources(self) -> None:
+    async def _subscribe_to_datasource_events(self) -> None:
         for datasource in self._ctx.datasources.values():
             if not isinstance(datasource, IndexDatasource):
                 continue
             datasource.on_operations(self._dispatch_operations)
             datasource.on_big_maps(self._dispatch_big_maps)
             datasource.on_rollback(self._rollback)
 
+    async def _set_datasource_heads(self) -> None:
+        for datasource in self._ctx.datasources.values():
+            if isinstance(datasource, TzktDatasource):
+                await datasource.set_head_from_http()
+
     async def _load_index_states(self) -> None:
         await self._fetch_contracts()
         index_states = await IndexState.filter().all()
@@ -207,14 +216,15 @@ async def run(self, reindex: bool, oneshot: bool) -> None:
             await self._initialize_schema()
             await self._set_up_hasura(stack, tasks)
 
+            spawn_datasources_event: Optional[Event] = None
             if not oneshot:
                 await self._set_up_scheduler(stack, tasks)
-                await self._spawn_datasources(tasks)
+                spawn_datasources_event = await self._spawn_datasources(tasks)
 
             for name in self._config.indexes:
                 await self._ctx._spawn_index(name)
 
-            await self._set_up_index_dispatcher(tasks)
+            await self._set_up_index_dispatcher(tasks, spawn_datasources_event)
 
             await gather(*tasks)
 
@@ -313,19 +323,26 @@ async def _set_up_hasura(self, stack: AsyncExitStack, tasks: Set[Task]) -> None:
         tasks.add(create_task(hasura_gateway.configure()))
 
     async def _set_up_datasources(self, stack: AsyncExitStack) -> None:
-        # FIXME: Find a better way to do this
-        # if self._datasources:
-        #     raise RuntimeError
         await self._create_datasources()
         for datasource in self._datasources.values():
             await stack.enter_async_context(datasource)
 
-    async def _set_up_index_dispatcher(self, tasks: Set[Task]) -> None:
+    async def _set_up_index_dispatcher(self, tasks: Set[Task], spawn_datasources_event: Optional[Event]) -> None:
         index_dispatcher = IndexDispatcher(self._ctx)
-        tasks.add(create_task(index_dispatcher.run()))
+        tasks.add(create_task(index_dispatcher.run(spawn_datasources_event)))
+
+    async def _spawn_datasources(self, tasks: Set[Task]) -> Event:
+        event = Event()
+
+        async def _wrapper():
+            self._logger.info('Waiting for IndexDispatcher to spawn datasources')
+            await event.wait()
+            self._logger.info('Spawning datasources')
+            _tasks = [create_task(d.run()) for d in self._datasources.values()]
+            await gather(*_tasks)
 
-    async def _spawn_datasources(self, tasks: Set[Task]) -> None:
-        tasks.update(create_task(d.run()) for d in self._datasources.values())
+        tasks.add(create_task(_wrapper()))
+        return event
 
     async def _set_up_scheduler(self, stack: AsyncExitStack, tasks: Set[Task]) -> None:
         job_failed = Event()

src/dipdup/index.py

@@ -20,7 +20,7 @@
 )
 from dipdup.context import DipDupContext
 from dipdup.datasources.tzkt.datasource import BigMapFetcher, OperationFetcher, TzktDatasource
-from dipdup.exceptions import InvalidDataError
+from dipdup.exceptions import ConfigInitializationException, InvalidDataError
 from dipdup.models import BigMapData, BigMapDiff, HeadBlockData
 from dipdup.models import Index as IndexState
 from dipdup.models import IndexStatus, OperationData, Origination, Transaction
@@ -86,7 +86,11 @@ async def process(self) -> None:
 
         elif self._datasource.sync_level is None:
             self._logger.info('Datasource is not active, sync to the latest block')
-            last_level = (await self._datasource.get_head_block()).level
+            # NOTE: Late establishing connection to the WebSocket
+            if self.datasource.head:
+                last_level = self.datasource.head.level
+            else:
+                last_level = (await self._datasource.get_head_block()).level
             await self._synchronize(last_level)
 
         elif self._datasource.sync_level > self.state.level:
@@ -154,6 +158,7 @@ async def single_level_rollback(self, from_level: int) -> None:
             raise RuntimeError('Index level is higher than rollback level')
 
     async def _process_queue(self) -> None:
+        """Process WebSocket queue"""
         if not self._queue:
             return
         self._logger.info('Processing websocket queue')
@@ -195,8 +200,8 @@ async def _synchronize(self, last_level: int, cache: bool = False) -> None:
         await self._exit_sync_state(last_level)
 
     async def _process_level_operations(self, level: int, operations: List[OperationData], block: Optional[HeadBlockData] = None) -> None:
-        if level < self.state.level:
-            raise RuntimeError(f'Level of operation batch is lower than index state level: {level} < {self.state.level}')
+        if level <= self.state.level:
+            raise RuntimeError(f'Level of operation batch must be higher than index state level: {level} <= {self.state.level}')
 
         if self._rollback_level:
             levels = {
@@ -214,7 +219,7 @@ async def _process_level_operations(self, level: int, operations: List[Operation
             received_hashes = set([op.hash for op in operations])
             reused_hashes = received_hashes & expected_hashes
             if reused_hashes != expected_hashes:
-                await self._ctx.reindex(reason='attempted a single level rollback, but arrived block differs from processed one')
+                await self._ctx.reindex(reason='attempted a single level rollback, but arrived block has additional transactions')
 
             self._rollback_level = None
             self._last_hashes = set()
@@ -324,7 +329,7 @@ async def _on_match(
         """Prepare handler arguments, parse parameter and storage. Schedule callback in executor."""
         self._logger.info('%s: `%s` handler matched!', operation_subgroup.hash, handler_config.callback)
         if not handler_config.parent:
-            raise RuntimeError('Handler must have a parent')
+            raise ConfigInitializationException
 
         args: List[Optional[Union[Transaction, Origination, OperationData]]] = []
         for pattern_config, operation in zip(handler_config.pattern, matched_operations):
@@ -415,7 +420,8 @@ def __init__(self, ctx: DipDupContext, config: BigMapIndexConfig, datasource: Tz
     def push(self, level: int, big_maps: List[BigMapData], block: Optional[HeadBlockData] = None):
         self._queue.append((level, big_maps, block))
 
-    async def _process_queue(self):
+    async def _process_queue(self) -> None:
+        """Process WebSocket queue"""
         if not self._queue:
             return
         self._logger.info('Processing websocket queue')
@@ -450,8 +456,8 @@ async def _synchronize(self, last_level: int, cache: bool = False) -> None:
         await self._exit_sync_state(last_level)
 
     async def _process_level_big_maps(self, level: int, big_maps: List[BigMapData], block: Optional[HeadBlockData] = None):
-        if level < self.state.level:
-            raise RuntimeError(f'Level of operation batch is lower than index state level: {level} < {self.state.level}')
+        if level <= self.state.level:
+            raise RuntimeError(f'Level of big map batch must be higher than index state level: {level} <= {self.state.level}')
 
         async with in_global_transaction():
             self._logger.info('Processing %s big map diffs of level %s', len(big_maps), level)
@@ -479,7 +485,7 @@ async def _on_match(
         """Prepare handler arguments, parse key and value. Schedule callback in executor."""
         self._logger.info('%s: `%s` handler matched!', matched_big_map.operation_id, handler_config.callback)
         if not handler_config.parent:
-            raise RuntimeError('Handler must have a parent')
+            raise ConfigInitializationException
 
         if matched_big_map.action.has_key:
             key_type = handler_config.key_type_cls
@@ -525,14 +531,14 @@ async def _process_big_maps(self, big_maps: List[BigMapData]) -> None:
                     await self._on_match(handler_config, big_map)
 
     async def _get_big_map_addresses(self) -> Set[str]:
-        """Get addresses to fetch transactions from during initial synchronization"""
+        """Get addresses to fetch big map diffs from during initial synchronization"""
         addresses = set()
         for handler_config in self._config.handlers:
             addresses.add(cast(ContractConfig, handler_config.contract).address)
         return addresses
 
     async def _get_big_map_paths(self) -> Set[str]:
-        """Get addresses to fetch transactions from during initial synchronization"""
+        """Get addresses to fetch big map diffs from during initial synchronization"""
         paths = set()
         for handler_config in self._config.handlers:
             paths.add(handler_config.path)