Add SubscriptionManager class to avoid subscription storm, add merge_subscriptions flag (#170)

* wip

* Useful cli options

* on_error

* Bump

* Drop heavily outdated datasource tests

* Docs

* Fixes

* Fix tests, switch dependencies, fix lost subs

* TConfigure action on reindexing triggered

* Bump deps

* Initial SubscriptionManager impl attempt

* Fix subs shitstorm

* Less logs

* Logs

* Fix false alarm on duplicate sub

* Bump

* Convinience cli commands

* You gotta go fast

* Bump

* Subscription sync_levels WIP

* Typo

* Fix retrier

* Lint cli

* Bump, docs

* Typo

* Fix levels one more time

* Bump

* Fix sync levels

* Filter out operation subgroups, fix merge_subscriptions flag

* daemon flag, fix subgroups filtering

* slowdoooown, fix crash on empty operation batches, increase maxsize

* Fix bencchmarks

* Prefiilter operations during subgroup extraction

* Bump

* Fix applying filters to indexdispatcher

* Bump

Author: droserasprout

CHANGELOG.md

@@ -1,21 +1,31 @@
 # Changelog
 
-## 3.2.0 - [unreleased]
+## 4.0.0 - [unreleased]
 
 ### Fixed
 
 * tzkt: Realtime connection watchdog is no longer triggered after TzKT outtage.
-* codegen: Fixed missing imports in handlers generated during init.
-* coinbase: Candles caching disabled.
+* tzkt: Adding many indexes in runtime no longer leads to sending useless subscription requests.
+* tzkt: Fixed `get_originated_contracts` and `get_similar_contracts` methods whose output was limited to `batch_size` tunable.
 * index: Fixed head index callback invocation.
+* coinbase: Candles caching disabled.
+
+### Changed
+
+* cli: `run --oneshot` option is deprecated and will be removed in the next major release. Oneshot mode applied automatically when `first_level` and `last_level` fields are set in index config.  
+* codegen: Fixed missing imports in handlers generated during init.
 
 ### Added
 
-* cli: `dipdip run --early-realtime` flag to establish a real-time connection before all indexes are synchronized.
-* cli: `dipdup run --skip-hasura` flag to skip updating Hasura metadata.
+* cli: New flag `dipdip run --early-realtime` to establish a real-time connection before all indexes are synchronized.
+* cli: New flag `dipdup run --skip-hasura` to skip updating Hasura metadata.
+* cli: New command `dipdup status` to print the current status of indexes from database
+* cli: New command `dipdup config` to print config after resolving all links and variables
+* config: Added optional fields `first_level` and `last_level` to `TemplateIndexConfig`. These limits are applied after ones from the template itself.
 
 ### Improved
 
+* index: Time required to initialize indexes presented in database reduced by ~25%.
 * tzkt: Replaced `aiosignalrcore` library with `pysignalr`.
 
 ## 3.1.3 - 2021-11-15

poetry.lock

@@ -32,7 +32,7 @@ datamodel-code-generator = "^0.11.1"
 fcache = "^0.4.7"
 pydantic = "^1.8.1"
 pyhumps = "^3.0.2"
-pysignalr = {git = "https://github.com/baking-bad/pysignalr.git", branch = "initial-release"}
+pysignalr = "^0.1.0"
 pytezos = {version = "^3.2.4", optional = true}
 python-dotenv = "^0.18.0"
 "ruamel.yaml" = "^0.17.2"

pyproject.toml

@@ -1,5 +1,5 @@
 for i in config index; do
     rm tests/benchmarks/$i.latest.json;
-    python tests/benchmarks/$i.py -o tests/benchmarks/$i.latest.json;
+    python tests/benchmarks/$i.py -o tests/benchmarks/$i.latest.json --quiet;
     python -m pyperf compare_to --table tests/benchmarks/$i.json tests/benchmarks/$i.latest.json;
 done;

scripts/run_benchmarks.sh

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import logging
 import os
 import signal
@@ -12,6 +13,7 @@
 import sentry_sdk
 from dotenv import load_dotenv
 from fcache.cache import FileCache  # type: ignore
+from pydantic.json import pydantic_encoder
 from sentry_sdk.integrations.aiohttp import AioHttpIntegration
 from sentry_sdk.integrations.logging import LoggingIntegration
 from tortoise import Tortoise
@@ -145,6 +147,7 @@ async def cli(ctx, config: List[str], env_file: List[str], logging_config: str):
 @click.option('--postpone-jobs', is_flag=True, help='Do not start job scheduler until all indexes are synchronized')
 @click.option('--skip-hasura', is_flag=True, help='Do not update Hasura metadata')
 @click.option('--early-realtime', is_flag=True, help='Establish a realtime connection before all indexes are synchronized')
+@click.option('--merge-subscriptions', is_flag=True, help='Subscribe to all updates instead of individual contracts')
 @click.pass_context
 @cli_wrapper
 async def run(
@@ -153,17 +156,21 @@ async def run(
     postpone_jobs: bool,
     skip_hasura: bool,
     early_realtime: bool,
+    merge_subscriptions: bool,
 ) -> None:
+    if oneshot:
+        _logger.warning('`oneshot` argument is deprecated: use `first_level` and `last_level` fields of index config instead')
     config: DipDupConfig = ctx.obj.config
     config.initialize()
+    config.advanced.postpone_jobs |= postpone_jobs
+    config.advanced.skip_hasura |= skip_hasura
+    config.advanced.early_realtime |= early_realtime
+    config.advanced.merge_subscriptions |= merge_subscriptions
+
     set_decimal_context(config.package)
+
     dipdup = DipDup(config)
-    await dipdup.run(
-        oneshot=oneshot,
-        postpone_jobs=postpone_jobs or config.advanced.postpone_jobs,
-        skip_hasura=skip_hasura or config.advanced.skip_hasura,
-        early_realtime=early_realtime or config.advanced.early_realtime,
-    )
+    await dipdup.run()
 
 
 @cli.command(help='Generate missing callbacks and types')
@@ -187,6 +194,39 @@ async def migrate(ctx):
     await migrations.migrate()
 
 
+from dipdup.models import Index
+
+
+@cli.command(help='Show current status of indexes in database')
+@click.pass_context
+@cli_wrapper
+async def status(ctx):
+    config: DipDupConfig = ctx.obj.config
+    url = config.database.connection_string
+    models = f'{config.package}.models'
+    async with tortoise_wrapper(url, models):
+        # TODO: Formatting
+        print('_' * 80)
+        async for index in Index.all():
+            print(f'{index.name}\t{index.status.value}\t{index.level}')
+        print('_' * 80)
+
+
+# TODO: Docs, `--unsafe` argument to resolve env variables, default to not doing it
+@cli.command(help='Show config')
+@click.pass_context
+@cli_wrapper
+async def config(ctx):
+    import ruamel.yaml as yaml
+
+    config: DipDupConfig = ctx.obj.config
+    config_json = json.dumps(config, default=pydantic_encoder)
+    config_yaml = yaml.dump(yaml.safe_load(config_json), indent=2, default_flow_style=False)
+    print('_' * 80)
+    print(config_yaml)
+    print('_' * 80)
+
+
 # TODO: "cache clear"?
 @cli.command(help='Clear development request cache')
 @click.pass_context

src/dipdup/cli.py

@@ -12,7 +12,7 @@
 from copy import copy
 from dataclasses import field
 from enum import Enum
-from functools import reduce
+from functools import cached_property, reduce
 from os import environ as env
 from os.path import dirname
 from pydoc import locate
@@ -25,6 +25,7 @@
 from ruamel.yaml import YAML
 from typing_extensions import Literal
 
+from dipdup.datasources.subscription import BigMapSubscription, OriginationSubscription, Subscription, TransactionSubscription
 from dipdup.enums import ReindexingAction, ReindexingReasonC
 from dipdup.exceptions import ConfigInitializationException, ConfigurationError
 from dipdup.utils import import_from, pascal_to_snake, snake_to_pascal
@@ -590,20 +591,33 @@ def template_values(self, value: Dict[str, str]) -> None:
         self._template_values = value
 
 
+@dataclass
+class SubscriptionsMixin:
+    def __post_init_post_parse__(self) -> None:
+        self._subscriptions: Set[Subscription] = set()
+
+    @property
+    def subscriptions(self) -> Set[Subscription]:
+        return self._subscriptions
+
+
 @dataclass
 class IndexTemplateConfig(NameMixin):
     kind = 'template'
     template: str
     values: Dict[str, str]
+    first_level: int = 0
+    last_level: int = 0
 
 
 @dataclass
-class IndexConfig(TemplateValuesMixin, NameMixin, ParentMixin['ResolvedIndexConfigT']):
+class IndexConfig(TemplateValuesMixin, NameMixin, SubscriptionsMixin, ParentMixin['ResolvedIndexConfigT']):
     datasource: Union[str, TzktDatasourceConfig]
 
     def __post_init_post_parse__(self) -> None:
         TemplateValuesMixin.__post_init_post_parse__(self)
         NameMixin.__post_init_post_parse__(self)
+        SubscriptionsMixin.__post_init_post_parse__(self)
         ParentMixin.__post_init_post_parse__(self)
 
     @property
@@ -642,8 +656,8 @@ class OperationIndexConfig(IndexConfig):
 
     :param datasource: Alias of index datasource in `datasources` section
     :param contracts: Aliases of contracts being indexed in `contracts` section
-    :param first_level: First block to process (use with `--oneshot` run argument)
-    :param last_level: Last block to process (use with `--oneshot` run argument)
+    :param first_level: First block to process (one time sync)
+    :param last_level: Last block to process (one time sync)
     :param handlers: List of indexer handlers
     """
 
@@ -655,6 +669,33 @@ class OperationIndexConfig(IndexConfig):
     first_level: int = 0
     last_level: int = 0
 
+    @cached_property
+    def entrypoint_filter(self) -> Set[Optional[str]]:
+        entrypoints = set()
+        for handler_config in self.handlers:
+            for pattern_config in handler_config.pattern:
+                if isinstance(pattern_config, OperationHandlerTransactionPatternConfig):
+                    entrypoints.add(pattern_config.entrypoint)
+        return set(entrypoints)
+
+    @cached_property
+    def address_filter(self) -> Set[str]:
+        addresses = set()
+        for handler_config in self.handlers:
+            for pattern_config in handler_config.pattern:
+                if isinstance(pattern_config, OperationHandlerTransactionPatternConfig):
+                    if isinstance(pattern_config.source, ContractConfig):
+                        addresses.add(pattern_config.source.address)
+                    elif isinstance(pattern_config.source, str):
+                        raise ConfigInitializationException
+
+                    if isinstance(pattern_config.destination, ContractConfig):
+                        addresses.add(pattern_config.destination.address)
+                    elif isinstance(pattern_config.destination, str):
+                        raise ConfigInitializationException
+
+        return addresses
+
 
 @dataclass
 class BigMapHandlerConfig(HandlerConfig, kind='handler'):
@@ -809,11 +850,15 @@ class JobConfig(NameMixin):
     hook: Union[str, 'HookConfig']
     crontab: Optional[str] = None
     interval: Optional[int] = None
+    daemon: bool = False
     args: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init_post_parse__(self):
-        if int(bool(self.crontab)) + int(bool(self.interval)) != 1:
-            raise ConfigurationError('Either `interval` or `crontab` field must be specified')
+        if self.crontab and self.interval:
+            raise ConfigurationError('Only one of `crontab` and `interval` can be specified')
+        elif not (self.crontab or self.interval or self.daemon):
+            raise ConfigurationError('One of `crontab`, `interval` or `daemon` must be specified')
+
         NameMixin.__post_init_post_parse__(self)
 
     @property
@@ -878,10 +923,12 @@ def _args_with_context(self) -> Dict[str, str]:
 
 @dataclass
 class AdvancedConfig:
+    reindex: Dict[ReindexingReasonC, ReindexingAction] = Field(default_factory=dict)
+    oneshot: bool = False
     postpone_jobs: bool = False
     skip_hasura: bool = False
     early_realtime: bool = False
-    reindex: Dict[ReindexingReasonC, ReindexingAction] = Field(default_factory=dict)
+    merge_subscriptions: bool = False
 
 
 @dataclass
@@ -945,6 +992,17 @@ def package_path(self, value: str):
             raise ConfigInitializationException
         self._package_path = value
 
+    @property
+    def oneshot(self) -> bool:
+        syncable_indexes = tuple(c for c in self.indexes.values() if not isinstance(c, HeadIndexConfig))
+        oneshot_indexes = tuple(c for c in syncable_indexes if c.last_level)
+        if not oneshot_indexes:
+            return False
+        elif len(oneshot_indexes) == len(syncable_indexes):
+            return True
+        else:
+            raise ConfigurationError('Either all or none of indexes can have `last_level` field set')
+
     @classmethod
     def load(
         cls,
@@ -962,7 +1020,7 @@ def load(
             with open(filename) as file:
                 raw_config = file.read()
 
-            _logger.info('Substituting environment variables')
+            _logger.debug('Substituting environment variables')
             for match in re.finditer(ENV_VARIABLE_REGEX, raw_config):
                 variable, default_value = match.group(1), match.group(2)
                 config_environment[variable] = default_value
@@ -1034,7 +1092,7 @@ def initialize(self, skip_imports: bool = False) -> None:
             if index_config.name in self._imports_resolved:
                 continue
 
-            _logger.info('Loading callbacks and typeclasses of index `%s`', index_config.name)
+            _logger.debug('Loading callbacks and typeclasses of index `%s`', index_config.name)
 
             if isinstance(index_config, IndexTemplateConfig):
                 raise ConfigInitializationException
@@ -1107,7 +1165,7 @@ def get_pattern_type(pattern_items: Sequence[HandlerPatternConfigT]) -> str:
                 )
 
     def _resolve_template(self, template_config: IndexTemplateConfig) -> None:
-        _logger.info('Resolving index config `%s` from template `%s`', template_config.name, template_config.template)
+        _logger.debug('Resolving index config `%s` from template `%s`', template_config.name, template_config.template)
 
         template = self.get_template(template_config.template)
         raw_template = json.dumps(template, default=pydantic_encoder)
@@ -1124,6 +1182,9 @@ def _resolve_template(self, template_config: IndexTemplateConfig) -> None:
         new_index_config.template_values = template_config.values
         new_index_config.parent = template
         new_index_config.name = template_config.name
+        if not isinstance(new_index_config, HeadIndexConfig):
+            new_index_config.first_level |= template_config.first_level
+            new_index_config.last_level |= template_config.last_level
         self.indexes[template_config.name] = new_index_config
 
     def _resolve_templates(self) -> None:
@@ -1136,12 +1197,49 @@ def _resolve_links(self) -> None:
             if name in self._links_resolved:
                 continue
             self._resolve_index_links(index_config)
+            # TODO: Not exactly link resolving, move somewhere else
+            self._resolve_index_subscriptions(index_config)
             self._links_resolved.add(index_config.name)
 
         for job_config in self.jobs.values():
             if isinstance(job_config.hook, str):
                 job_config.hook = self.hooks[job_config.hook]
 
+    def _resolve_index_subscriptions(self, index_config: IndexConfigT) -> None:
+        if isinstance(index_config, IndexTemplateConfig):
+            return
+        if index_config.subscriptions:
+            return
+
+        if isinstance(index_config, OperationIndexConfig):
+            if self.advanced.merge_subscriptions:
+                index_config.subscriptions.add(TransactionSubscription())
+            else:
+                for contract_config in index_config.contracts or ():
+                    address = cast(ContractConfig, contract_config).address
+                    index_config.subscriptions.add(TransactionSubscription(address=address))
+
+            for handler_config in index_config.handlers:
+                for pattern_config in handler_config.pattern:
+                    if isinstance(pattern_config, OperationHandlerOriginationPatternConfig):
+                        index_config.subscriptions.add(OriginationSubscription())
+                        break
+
+        elif isinstance(index_config, BigMapIndexConfig):
+            if self.advanced.merge_subscriptions:
+                index_config.subscriptions.add(BigMapSubscription())
+            else:
+                for big_map_handler_config in index_config.handlers:
+                    address, path = big_map_handler_config.contract_config.address, big_map_handler_config.path
+                    index_config.subscriptions.add(BigMapSubscription(address=address, path=path))
+
+        # NOTE: HeadSubscription is always enabled
+        elif isinstance(index_config, HeadIndexConfig):
+            pass
+
+        else:
+            raise NotImplementedError(f'Index kind `{index_config.kind}` is not supported')
+
     def _resolve_index_links(self, index_config: IndexConfigT) -> None:
         """Resolve contract and datasource configs by aliases"""