From 134ba44aff3800b20b08cd596663bad8430602f5 Mon Sep 17 00:00:00 2001
From: Alexandr Rutkowski <kitanoyoru@protonmail.com>
Date: Sun, 8 Oct 2023 00:08:38 +0300
Subject: [PATCH] dev(auth): health check

---
 .../src/services/grpc/health_service/check.py |  9 ++----
 apps/auth/src/utils/deadline.py               | 31 ++++++++++++++++++
 apps/auth/src/utils/wrapper.py                | 32 ++++++++++---------
 3 files changed, 51 insertions(+), 21 deletions(-)
 create mode 100644 apps/auth/src/utils/deadline.py

diff --git a/apps/auth/src/services/grpc/health_service/check.py b/apps/auth/src/services/grpc/health_service/check.py
index f70604b..e502d30 100644
--- a/apps/auth/src/services/grpc/health_service/check.py
+++ b/apps/auth/src/services/grpc/health_service/check.py
@@ -5,6 +5,8 @@
 from typing import Awaitable, Callable, Optional, Set
 
 from src.constants import Constants
+from src.utils.deadline import Deadline
+from src.utils.wrapper import DeadlineWrapper
 
 log = logging.getLogger(__name__)
 
@@ -77,9 +79,7 @@ async def __check__(self) -> _Status:
         try:
             deadline = Deadline.from_timeout(self._check_timeout)
             with self._check_wrapper.start(deadline):
-                value = await self._check_func()
-            if value is not None and not isinstance(value, bool):
-                raise TypeError(f"Invalid health check func status type: {value}")
+                await self._check_func()
         except asyncio.CancelledError:
             raise
         except Exception:
@@ -160,6 +160,3 @@ async def __subscribe__(self) -> asyncio.Event:
     async def __unsubscribe__(self, event: asyncio.Event):
         self._events.discard(event)
 
-
-class DeadlineWrapper:
-    pass
diff --git a/apps/auth/src/utils/deadline.py b/apps/auth/src/utils/deadline.py
new file mode 100644
index 0000000..3b70462
--- /dev/null
+++ b/apps/auth/src/utils/deadline.py
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import time
+
+
+class Deadline:
+    def __init__(self, *, timestamp: float):
+        self._timestamp = timestamp
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, Deadline):
+            raise TypeError(
+                f"Comparison is not supported between {type(self).__name__} and {type(other).__name__}"
+            )
+
+        return self._timestamp < other._timestamp
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Deadline):
+            raise TypeError(
+                f"Comparison is not supported between {type(self).__name__} and {type(other).__name__}"
+            )
+
+        return self._timestamp == other._timestamp
+
+    @classmethod
+    def from_timeout(cls, timeout: float) -> Deadline:
+        return cls(timestamp=time.monotonic() + timeout)
+
+    def time_remaining(self) -> float:
+        return max(0, self._timestamp - time.monotonic())
diff --git a/apps/auth/src/utils/wrapper.py b/apps/auth/src/utils/wrapper.py
index 3dcb5a8..759281a 100644
--- a/apps/auth/src/utils/wrapper.py
+++ b/apps/auth/src/utils/wrapper.py
@@ -3,7 +3,7 @@
 from types import TracebackType
 from typing import Any, ContextManager, Iterator, Optional, Set, Type
 
-_current_task = asyncio.current_task
+from src.utils.deadline import Deadline
 
 
 class Wrapper(ContextManager[None]):
@@ -12,14 +12,14 @@ class Wrapper(ContextManager[None]):
     cancelled: Optional[bool] = None
     cancel_failed: Optional[bool] = None
 
-    def __init__(self) -> None:
-        self._tasks: Set["asyncio.Task[Any]"] = set()
+    def __init__(self):
+        self._tasks: Set[asyncio.Task[Any]] = set()
 
     def __enter__(self):
         if self._error is not None:
             raise self._error
 
-        task = _current_task()
+        task = asyncio.current_task()
         if task is None:
             raise RuntimeError("Called not inside a task")
 
@@ -28,17 +28,19 @@ def __enter__(self):
     def __exit__(
         self,
         exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        task = _current_task()
+        _exc_value: Optional[BaseException],
+        _exc_traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        task = asyncio.current_task()
         assert task
+
         self._tasks.discard(task)
+
         if self._error is not None:
             self.cancel_failed = exc_type is not asyncio.CancelledError
             raise self._error
 
-    def cancel(self, error: Exception) -> None:
+    def cancel(self, error: Exception):
         self._error = error
         for task in self._tasks:
             task.cancel()
@@ -47,16 +49,16 @@ def cancel(self, error: Exception) -> None:
 
 class DeadlineWrapper(Wrapper):
     @contextmanager
-    def start(self, deadline: "Deadline") -> Iterator[None]:
+    def start(self, deadline: Deadline) -> Iterator[None]:
         timeout = deadline.time_remaining()
         if not timeout:
-            raise asyncio.TimeoutError("Deadline exceeded")
-
-        def callback() -> None:
-            self.cancel(asyncio.TimeoutError("Deadline exceeded"))
+            raise asyncio.TimeoutError("Deadline exceed")
 
         loop = asyncio.get_event_loop()
-        timer = loop.call_later(timeout, callback)
+        timer = loop.call_later(
+            timeout, lambda: self.cancel(asyncio.TimeoutError("Deadline exceed"))
+        )
+
         try:
             yield
         finally: