yellowstone-fumarole-client 0.2.2__tar.gz → 0.3.1__tar.gz

{yellowstone_fumarole_client-0.2.2 → yellowstone_fumarole_client-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: yellowstone-fumarole-client
-Version: 0.2.2
+Version: 0.3.1
 Summary: Yellowstone Fumarole Python Client
 Home-page: https://github.com/rpcpool/yellowstone-fumarole
 Author: Louis-Vincent

{yellowstone_fumarole_client-0.2.2 → yellowstone_fumarole_client-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "yellowstone-fumarole-client"
-version = "0.2.2"
+version = "0.3.1"
 homepage = "https://github.com/rpcpool/yellowstone-fumarole"
 repository = "https://github.com/rpcpool/yellowstone-fumarole"
 description = "Yellowstone Fumarole Python Client"

{yellowstone_fumarole_client-0.2.2 → yellowstone_fumarole_client-0.3.1}/yellowstone_fumarole_client/__init__.py

@@ -1,5 +1,7 @@
 import asyncio
 import logging
+from contextlib import suppress
+
 from yellowstone_fumarole_client.grpc_connectivity import (
     FumaroleGrpcConnector,
 )
@@ -33,6 +35,7 @@ from yellowstone_fumarole_proto.fumarole_pb2 import (
     CreateConsumerGroupResponse,
 )
 from yellowstone_fumarole_proto.fumarole_pb2_grpc import FumaroleStub
+from yellowstone_fumarole_client.error import SubscribeError
 import grpc
 
 from yellowstone_fumarole_client import config
@@ -54,6 +57,7 @@ DEFAULT_COMMIT_INTERVAL = 5.0  # seconds
 DEFAULT_MAX_SLOT_DOWNLOAD_ATTEMPT = 3
 DEFAULT_CONCURRENT_DOWNLOAD_LIMIT_PER_TCP = 10
 
+LOGGER = logging.getLogger(__name__)
 # Error classes
 
 
@@ -262,18 +266,28 @@ class FumaroleClient:
             max_concurrent_download=config.concurrent_download_limit,
         )
 
-        async def rt_run(rt):
-            async with rt as rt:
-                await rt.run()
-
-        rt_task = asyncio.create_task(rt_run(rt))
-
+        rt_task = asyncio.create_task(rt.run())
+        rt_task.set_name("rt_task")
+        control_plane_src_task.set_name("control_plane_src_task")
         async def fumarole_overseer():
-            done, pending = await asyncio.wait(
-                [rt_task, control_plane_src_task], return_when=asyncio.FIRST_COMPLETED
-            )
-            for t in pending:
-                t.cancel()
+            try:
+                done, pending = await asyncio.wait(
+                    [rt_task, control_plane_src_task], return_when=asyncio.FIRST_COMPLETED
+                )
+                for t in done:
+                    try:
+                        exc = t.exception()
+                        LOGGER.error(f"Fumarole task '{t.get_name()}' failed: {exc}")
+                    except asyncio.CancelledError:
+                        pass
+                    if exc is not None:
+                        with suppress(asyncio.QueueShutDown):
+                            await dragonsmouth_outlet.put(exc)
+                for t in pending:
+                    t.cancel()
+            finally:
+                await rt.aclose()
+
 
         fumarole_handle = asyncio.create_task(fumarole_overseer())
 
@@ -281,6 +295,8 @@ class FumaroleClient:
             try:
                 while True:
                     update = await dragonsmouth_outlet.get()
+                    if isinstance(update, Exception):
+                        raise update
                     yield update
             except (asyncio.CancelledError, asyncio.QueueShutDown):
                 pass

yellowstone_fumarole_client-0.3.1/yellowstone_fumarole_client/error.py

@@ -0,0 +1,19 @@
+from typing import Any, Mapping
+
+
+class FumaroleClientError(Exception):
+    """Base class for all Fumarole-related exceptions."""
+    pass
+
+
+class SubscribeError(FumaroleClientError):
+    """Base class for all errors related to subscription."""
+    pass
+
+
+class DownloadSlotError(SubscribeError):
+    """Exception raised for errors in the download slot process."""
+    
+    def __init__(self, message: str, ctx: Mapping[str, Any]):
+        super().__init__(message)
+        self.ctx = ctx

{yellowstone_fumarole_client-0.2.2 → yellowstone_fumarole_client-0.3.1}/yellowstone_fumarole_client/runtime/aio.py

@@ -1,8 +1,11 @@
 # DataPlaneConn
 from abc import abstractmethod, ABC
 import asyncio
+from contextlib import suppress
+from typing import Any
+
 import grpc
-from typing import Optional
+from typing import Literal, Mapping, Optional
 from collections import deque
 from dataclasses import dataclass
 import time
@@ -16,7 +19,6 @@ from yellowstone_fumarole_proto.geyser_pb2 import (
     SubscribeRequest,
     SubscribeUpdate,
     SubscribeUpdateSlot,
-    CommitmentLevel as ProtoCommitmentLevel,
 )
 from yellowstone_fumarole_proto.fumarole_pb2 import (
     ControlCommand,
@@ -30,7 +32,7 @@ from yellowstone_fumarole_proto.fumarole_pb2_grpc import (
     FumaroleStub,
 )
 from yellowstone_fumarole_client.utils.aio import Interval
-from yellowstone_fumarole_client.grpc_connectivity import FumaroleGrpcConnector
+from yellowstone_fumarole_client.error import SubscribeError, DownloadSlotError
 import logging
 
 
@@ -51,13 +53,28 @@ class CompletedDownloadBlockTask:
     total_event_downloaded: int
 
 
+DownloadBlockErrorKind = Literal[
+    "Disconnected",
+    "OutletDisconnected",
+    "BlockShardNotFound",
+    "FailedDownload",
+    "Fatal",
+    "SessionDeprecated"
+]
+
 @dataclass
 class DownloadBlockError:
     """Represents an error that occurred during the download of a block."""
-
-    kind: str  # 'Disconnected', 'OutletDisconnected', 'BlockShardNotFound', 'FailedDownload', 'Fatal'
+    kind: DownloadBlockErrorKind
     message: str
+    ctx: Mapping[str, Any]
 
+    def into_subscribe_error(self) -> SubscribeError:
+        """Convert the error into a SubscribeError."""
+        return DownloadSlotError(
+            f"{self.kind}: {self.message}",
+            ctx=self.ctx
+        )
 
 @dataclass
 class DownloadTaskResult:
@@ -139,6 +156,8 @@ class AsyncioFumeDragonsmouthRuntime:
         # holds metadata about the download task
         self.download_tasks = dict()
         self.inflight_tasks = dict()
+        self.successful_download_cnt = 0
+        self.is_closed = False
 
     async def __aenter__(self):
         return self
@@ -227,18 +246,24 @@ class AsyncioFumeDragonsmouthRuntime:
             self.inflight_tasks[download_task] = DOWNLOAD_TASK_TYPE_MARKER
             LOGGER.debug(f"Scheduling download task for slot {download_request.slot}")
 
-    def _handle_download_result(self, download_result: DownloadTaskResult):
+    async def _handle_download_result(self, download_result: DownloadTaskResult):
         """Handles the result of a download task."""
         if download_result.kind == "Ok":
+            self.successful_download_cnt += 1
             completed = download_result.completed
             LOGGER.debug(
-                f"Download completed for slot {completed.slot}, shard {completed.shard_idx}, {completed.total_event_downloaded} total events"
+                f"Download({self.successful_download_cnt}) completed for slot {completed.slot}, shard {completed.shard_idx}, {completed.total_event_downloaded} total events"
             )
             self.sm.make_slot_download_progress(completed.slot, completed.shard_idx)
         else:
             slot = download_result.slot
-            err = download_result.err
-            raise RuntimeError(f"Failed to download slot {slot}: {err.message}")
+            err_kind = download_result.err.kind
+            LOGGER.error(f"Download error for slot {slot}: {download_result.err}")
+            # If the client queue is disconnected, we don't do anything, next run iteration will close anyway.
+            self.is_closed = True
+            if err_kind != "OutletDisconnected":
+                with suppress(asyncio.QueueShutDown):
+                    await self.dragonsmouth_outlet.put(download_result.err.into_subscribe_error())
 
     async def _force_commit_offset(self):
         LOGGER.debug(f"Force committing offset {self.sm.committable_offset}")
@@ -328,7 +353,7 @@ class AsyncioFumeDragonsmouthRuntime:
             ): COMMIT_TICK_TYPE_MARKER,
         }
 
-        while self.inflight_tasks:
+        while self.inflight_tasks and not self.is_closed:
             ticks += 1
             LOGGER.debug(f"Runtime loop tick")
             if ticks % self.gc_interval == 0:
@@ -370,7 +395,7 @@ class AsyncioFumeDragonsmouthRuntime:
                 elif sigcode == DOWNLOAD_TASK_TYPE_MARKER:
                     LOGGER.debug("Download task result received")
                     assert self.download_tasks.pop(t)
-                    self._handle_download_result(result)
+                    await self._handle_download_result(result)
                 elif sigcode == COMMIT_TICK_TYPE_MARKER:
                     LOGGER.debug("Commit tick reached")
                     await self._commit_offset()
@@ -382,10 +407,7 @@ class AsyncioFumeDragonsmouthRuntime:
                     raise RuntimeError(f"Unexpected task name: {sigcode}")
 
             await self._drain_slot_status()
-
         await self.aclose()
-        LOGGER.debug("Fumarole runtime exiting")
-
 
 # DownloadTaskRunnerChannels
 @dataclass
@@ -456,15 +478,21 @@ class GrpcDownloadBlockTaskRun:
         self.dragonsmouth_oulet = dragonsmouth_oulet
 
     def map_tonic_error_code_to_download_block_error(
-        self, e: grpc.aio.AioRpcError
+        self,
+        e: grpc.aio.AioRpcError
     ) -> DownloadBlockError:
         code = e.code()
+        ctx = {
+            "grpc_status_code": code,
+            "slot": self.download_request.slot,
+            "block_uid": self.download_request.block_uid,
+        }
         if code == grpc.StatusCode.NOT_FOUND:
             return DownloadBlockError(
-                kind="BlockShardNotFound", message="Block shard not found"
+                kind="BlockShardNotFound", message="Block shard not found", ctx=ctx,
             )
         elif code == grpc.StatusCode.UNAVAILABLE:
-            return DownloadBlockError(kind="Disconnected", message="Disconnected")
+            return DownloadBlockError(kind="Disconnected", message="Disconnected", ctx=ctx)
         elif code in (
             grpc.StatusCode.INTERNAL,
             grpc.StatusCode.ABORTED,
@@ -474,11 +502,13 @@ class GrpcDownloadBlockTaskRun:
             grpc.StatusCode.CANCELLED,
             grpc.StatusCode.DEADLINE_EXCEEDED,
         ):
-            return DownloadBlockError(kind="FailedDownload", message="Failed download")
+            return DownloadBlockError(kind="FailedDownload", message="Failed download", ctx=ctx)
         elif code == grpc.StatusCode.INVALID_ARGUMENT:
             raise ValueError("Invalid argument")
+        elif code == grpc.StatusCode.FAILED_PRECONDITION:
+            return DownloadBlockError(kind="SessionDeprecated", message="Session is deprecated", ctx=ctx)
         else:
-            return DownloadBlockError(kind="Fatal", message=f"Unknown error: {code}")
+            return DownloadBlockError(kind="Fatal", message=f"Unknown error: {code}", ctx=ctx)
 
     async def run(self) -> DownloadTaskResult:
         request = DownloadBlockShard(
@@ -519,6 +549,10 @@ class GrpcDownloadBlockTaskRun:
                                 err=DownloadBlockError(
                                     kind="OutletDisconnected",
                                     message="Outlet disconnected",
+                                    ctx = {
+                                        "slot": self.download_request.slot,
+                                        "block_uid": self.download_request.block_uid,
+                                    }
                                 ),
                             )
                     case "block_shard_download_finish":
@@ -535,17 +569,26 @@ class GrpcDownloadBlockTaskRun:
                             ),
                         )
                     case unknown:
-                        raise RuntimeError("Unexpected response kind: {unknown}")
+                        raise RuntimeError(f"Unexpected response kind: {unknown}")
+
         except grpc.aio.AioRpcError as e:
-            LOGGER.error(f"Download block error: {e}")
+            e2 = self.map_tonic_error_code_to_download_block_error(e)
+            LOGGER.error(f"Download block error: {e}, {e2}")
             return DownloadTaskResult(
                 kind="Err",
                 slot=self.download_request.slot,
-                err=self.map_tonic_error_code_to_download_block_error(e),
+                err=e2,
             )
 
         return DownloadTaskResult(
             kind="Err",
             slot=self.download_request.slot,
-            err=DownloadBlockError(kind="FailedDownload", message="Failed download"),
+            err=DownloadBlockError(
+                kind="FailedDownload",
+                message="Failed download",
+                ctx={
+                    "slot": self.download_request.slot,
+                    "block_uid": self.download_request.block_uid,
+                },
+            ),
         )