Extend consensus clients (#51)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sw-utils"
-version = "0.3.16"
+version = "0.3.17"
 description = "StakeWise Python utils"
 authors = ["StakeWise Labs <[email protected]>"]
 license = "GPL-3.0-or-later"

sw_utils/consensus.py

@@ -3,14 +3,18 @@
 from typing import TYPE_CHECKING, Any
 
 import aiohttp
-from eth_typing import URI, HexStr
+from aiohttp import ClientResponseError
+from eth_typing import URI, BlockNumber, HexStr
+from web3 import Web3
 from web3._utils.request import async_json_make_get_request
 from web3.beacon import AsyncBeacon
 from web3.beacon.api_endpoints import GET_VOLUNTARY_EXITS
+from web3.types import Timestamp
 
 from sw_utils.common import urljoin
 from sw_utils.decorators import retry_aiohttp_errors
 from sw_utils.exceptions import AiohttpRecoveredErrors
+from sw_utils.typings import ChainHead, ConsensusFork
 
 logger = logging.getLogger(__name__)
 
@@ -95,6 +99,38 @@ async def submit_voluntary_exit(
                     raise error
                 logger.error('%s: %s', url, repr(error))
 
+    async def get_chain_finalized_head(self, slots_per_epoch: int) -> ChainHead:
+        """Fetches the fork safe chain head."""
+        checkpoints = await self.get_finality_checkpoint()
+        epoch: int = int(checkpoints['data']['finalized']['epoch'])
+        last_slot_id: int = (epoch * slots_per_epoch) + slots_per_epoch - 1
+        for i in range(slots_per_epoch):
+            try:
+                slot = await self.get_block(str(last_slot_id - i))
+            except ClientResponseError as e:
+                if hasattr(e, 'status') and e.status == 404:
+                    # slot was not proposed, try the previous one
+                    continue
+                raise e
+
+            execution_payload = slot['data']['message']['body']['execution_payload']
+            return ChainHead(
+                epoch=epoch,
+                consensus_block=last_slot_id - i,
+                execution_block=BlockNumber(int(execution_payload['block_number'])),
+                execution_ts=Timestamp(int(execution_payload['timestamp'])),
+            )
+
+        raise RuntimeError(f'Failed to fetch slot for epoch {epoch}')
+
+    async def get_consensus_fork(self, state_id: str = 'head') -> ConsensusFork:
+        """Fetches current fork data."""
+        fork_data = (await self.get_fork_data(state_id))['data']
+        return ConsensusFork(
+            version=Web3.to_bytes(hexstr=fork_data['current_version']),
+            epoch=int(fork_data['epoch']),
+        )
+
     async def _async_make_get_request(self, endpoint_uri: str) -> dict[str, Any]:
         if self.retry_timeout:
 

sw_utils/decorators.py

@@ -3,7 +3,15 @@
 import typing
 
 import aiohttp
-from tenacity import retry, retry_if_exception, stop_after_delay, wait_exponential
+from tenacity import (
+    retry,
+    retry_if_exception,
+    retry_if_exception_type,
+    stop_after_delay,
+    wait_exponential,
+)
+
+from .ipfs import IpfsException
 
 default_logger = logging.getLogger(__name__)
 
@@ -40,3 +48,12 @@ def retry_aiohttp_errors(delay: int = 60, log_func=custom_before_log):
         stop=stop_after_delay(delay),
         before=log_func(default_logger, logging.INFO),
     )
+
+
+def retry_ipfs_exception(delay: int):
+    return retry(
+        retry=retry_if_exception_type(IpfsException),
+        wait=wait_exponential(multiplier=1, min=1, max=delay // 2),
+        stop=stop_after_delay(delay),
+        before=custom_before_log(default_logger, logging.INFO),
+    )

sw_utils/typings.py

@@ -1,10 +1,21 @@
 from dataclasses import dataclass
 from typing import NewType
 
+from eth_typing import BlockNumber
+from web3.types import Timestamp
+
 Bytes32 = NewType('Bytes32', bytes)
 
 
 @dataclass
 class ConsensusFork:
     version: bytes
     epoch: int
+
+
+@dataclass
+class ChainHead:
+    epoch: int
+    consensus_block: int
+    execution_block: BlockNumber
+    execution_ts: Timestamp