diff --git a/ape-config.yaml b/ape-config.yaml new file mode 100644 index 00000000..4b16339d --- /dev/null +++ b/ape-config.yaml @@ -0,0 +1,4 @@ +# Ensure we are configured for fork mode correct for example +ethereum: + mainnet-fork: + default_provider: foundry diff --git a/silverback/main.py b/silverback/main.py index 015ec1a5..eccedd5e 100644 --- a/silverback/main.py +++ b/silverback/main.py @@ -1,7 +1,9 @@ import atexit +import inspect from collections import defaultdict from datetime import timedelta -from typing import Any, Callable +from functools import wraps +from typing import Any, Awaitable, Callable from ape.api.networks import LOCAL_NETWORK_NAME from ape.contracts import ContractEvent, ContractInstance @@ -108,17 +110,18 @@ def __init__(self, settings: Settings | None = None): provider_context = settings.get_provider_context() # NOTE: This allows using connected ape methods e.g. `Contract` - provider = provider_context.__enter__() + self.provider = provider_context.__enter__() self.identifier = SilverbackID( name=settings.BOT_NAME, - network=provider.network.name, - ecosystem=provider.network.ecosystem.name, + network=self.provider.network.name, + ecosystem=self.provider.network.ecosystem.name, ) # Adjust defaults from connection if settings.NEW_BLOCK_TIMEOUT is None and ( - provider.network.name.endswith("-fork") or provider.network.name == LOCAL_NETWORK_NAME + self.provider.network.name.endswith("-fork") + or self.provider.network.name == LOCAL_NETWORK_NAME ): settings.NEW_BLOCK_TIMEOUT = int(timedelta(days=1).total_seconds()) @@ -138,6 +141,7 @@ def __init__(self, settings: Settings | None = None): self.signer = settings.get_signer() self.new_block_timeout = settings.NEW_BLOCK_TIMEOUT + self.use_fork = settings.FORK_MODE and not self.provider.network.name.endswith("-fork") signer_str = f"\n SIGNER={repr(self.signer)}" new_block_timeout_str = ( @@ -146,7 +150,9 @@ def __init__(self, settings: Settings | None = None): network_choice = f"{self.identifier.ecosystem}:{self.identifier.network}" logger.success( - f'Loaded Silverback Bot:\n NETWORK="{network_choice}"' + "Loaded Silverback Bot:\n" + f' NETWORK="{network_choice}"\n' + f" FORK_MODE={self.use_fork}" f"{signer_str}{new_block_timeout_str}" ) @@ -225,6 +231,25 @@ async def __create_snapshot_handler( last_block_processed=self.state.get("system:last_block_processed", -1), ) + # To ensure we don't have too many forks at once + # HACK: Until `NetworkManager.fork` (and `ProviderContextManager`) allows concurrency + + def _with_fork_decorator(self, handler: Callable) -> Callable: + # Trigger worker-side handling using fork network by wrapping handler + fork_context = self.provider.network_manager.fork + + @wraps(handler) + async def fork_handler(*args, **kwargs): + with fork_context(): + result = handler(*args, **kwargs) + + if inspect.isawaitable(result): + return await result + + return result + + return fork_handler + def broker_task_decorator( self, task_type: TaskType, @@ -266,7 +291,9 @@ def broker_task_decorator( raise ContainerTypeMismatchError(task_type, container) # Register user function as task handler with our broker - def add_taskiq_task(handler: Callable) -> AsyncTaskiqDecoratedTask: + def add_taskiq_task( + handler: Callable[..., Any | Awaitable[Any]] + ) -> AsyncTaskiqDecoratedTask: labels = {"task_type": str(task_type)} # NOTE: Do *not* do `if container` because that does a `len(container)` call, @@ -276,7 +303,7 @@ def add_taskiq_task(handler: Callable) -> AsyncTaskiqDecoratedTask: # Address is almost a certainty if the container is being used as a filter here. if not (contract_address := getattr(container.contract, "address", None)): raise InvalidContainerTypeError( - "Please provider a contract event from a valid contract instance." + "Please provide a contract event from a valid contract instance." ) labels["contract_address"] = contract_address @@ -284,6 +311,9 @@ def add_taskiq_task(handler: Callable) -> AsyncTaskiqDecoratedTask: self.tasks[task_type].append(TaskData(name=handler.__name__, labels=labels)) + if self.use_fork: + handler = self._with_fork_decorator(handler) + return self.broker.register_task( handler, task_name=handler.__name__, diff --git a/silverback/settings.py b/silverback/settings.py index 7b756484..acd2b6d4 100644 --- a/silverback/settings.py +++ b/silverback/settings.py @@ -28,6 +28,10 @@ class Settings(BaseSettings, ManagerAccessMixin): # A unique identifier for this silverback instance BOT_NAME: str = "bot" + # Execute every handler using an independent fork context + # NOTE: Requires fork-able provider installed and configured for network + FORK_MODE: bool = False + BROKER_CLASS: str = "taskiq:InMemoryBroker" BROKER_URI: str = "" # To be deprecated in 0.6 BROKER_KWARGS: dict[str, Any] = dict()