Feature/async checkpoints (#22)

* Fix healthcheck port.

* Fix import slipstream without installing extra cache dependency.

* Allow generic parameters for AsyncCallable.

* Added helper to handle potential awaitable/non-awaitables.

* Add test.

Author: Menziess

docker-compose.yml

@@ -26,7 +26,7 @@ services:
           "kafka-topics",
           "--list",
           "--bootstrap-server",
-          "localhost:9091",
+          "localhost:29091",
         ]
       interval: 5s
       timeout: 10s

slipstream/caching.py

@@ -6,12 +6,7 @@
 from contextlib import asynccontextmanager
 from pathlib import Path
 from types import TracebackType
-from typing import (
-    Any,
-    TypeVar,
-)
-
-from rocksdict import WriteBatch
+from typing import Any, TypeVar
 
 from slipstream.interfaces import ICache, Key
 
@@ -49,6 +44,7 @@
         RdictIter,
         ReadOptions,
         Snapshot,
+        WriteBatch,
         WriteOptions,
     )
     from rocksdict.rocksdict import (

slipstream/checkpointing.py

@@ -3,13 +3,11 @@
 import logging
 from collections.abc import AsyncIterable, Callable, Generator
 from datetime import datetime, timedelta
-from typing import (
-    Any,
-)
+from typing import Any
 
 from slipstream.core import Conf, Signal
 from slipstream.interfaces import ICache
-from slipstream.utils import iscoroutinecallable
+from slipstream.utils import AsyncCallable, awaitable, iscoroutinecallable
 
 _logger = logging.getLogger(__name__)
 
@@ -35,12 +33,16 @@ class Dependency:
     """
 
     @property
-    def downtime_check(self) -> Callable[['Checkpoint', 'Dependency'], Any]:
+    def downtime_check(
+        self,
+    ) -> AsyncCallable[['Checkpoint', 'Dependency'], Any]:
         """Is called when downtime is detected."""
         return self._downtime_check
 
     @property
-    def recovery_check(self) -> Callable[['Checkpoint', 'Dependency'], bool]:
+    def recovery_check(
+        self,
+    ) -> AsyncCallable[['Checkpoint', 'Dependency'], bool]:
         """Is called when downtime is resolved."""
         return self._recovery_check
 
@@ -49,9 +51,9 @@ def __init__(
         name: str,
         dependency: AsyncIterable[Any],
         downtime_threshold: Any = timedelta(minutes=10),
-        downtime_check: Callable[['Checkpoint', 'Dependency'], Any]
+        downtime_check: AsyncCallable[['Checkpoint', 'Dependency'], Any]
         | None = None,
-        recovery_check: Callable[['Checkpoint', 'Dependency'], bool]
+        recovery_check: AsyncCallable[['Checkpoint', 'Dependency'], bool]
         | None = None,
     ) -> None:
         """Initialize dependency for checkpointing."""
@@ -286,7 +288,7 @@ async def heartbeat(
         self._save_checkpoint(dependency, self.state, marker)
 
         if dependency.is_down:
-            if dependency.recovery_check(self, dependency):
+            if await awaitable(dependency.recovery_check(self, dependency)):
                 dependency.is_down = False
 
             if not any(_.is_down for _ in self.dependencies.values()):
@@ -342,7 +344,9 @@ async def check_pulse(
 
             # Trigger on the first dependency that is down and
             # pause the dependent stream
-            if downtime := dependency.downtime_check(self, dependency):
+            if downtime := await awaitable(
+                dependency.downtime_check(self, dependency)
+            ):
                 log_msg = (
                     f'Downtime of dependency "{dependency.name}" detected'
                 )

slipstream/utils.py

@@ -7,13 +7,15 @@
 from typing import (
     Any,
     ClassVar,
+    ParamSpec,
     TypeAlias,
     TypeVar,
 )
 
 T = TypeVar('T')
+P = ParamSpec('P')
 
-AsyncCallable: TypeAlias = Callable[..., Awaitable[Any]] | Callable[..., Any]
+AsyncCallable: TypeAlias = Callable[P, T | Awaitable[T]]
 Pipe: TypeAlias = Callable[[AsyncIterable[Any]], AsyncIterable[Any]]
 
 
@@ -31,6 +33,11 @@ class Signal(Enum):
     STOP = 3
 
 
+async def awaitable(x: Any) -> Any:
+    """Convert into awaitable."""
+    return await x if isinstance(x, Awaitable) else x
+
+
 def iscoroutinecallable(o: Any) -> bool:
     """Check whether object is coroutine."""
     call = o.__call__ if callable(o) else None  # type: ignore[attr-defined]

tests/test_checkpointing.py

@@ -268,6 +268,58 @@ async def test_custom_callbacks(is_async, checkpoint, mocker):
     )
 
 
+@pytest.mark.parametrize('is_async', [True, False])
+@pytest.mark.asyncio
+async def test_custom_checks(is_async, mock_cache, mocker):
+    """Check custom check functions called."""
+    if is_async:
+        downtime_check = mocker.AsyncMock(return_value=timedelta(hours=1))
+        recovery_check = mocker.AsyncMock(return_value=timedelta(hours=1))
+    else:
+        downtime_check = mocker.Mock(return_value=timedelta(hours=1))
+        recovery_check = mocker.Mock(return_value=timedelta(hours=1))
+
+    async def messages():
+        yield {
+            'event_timestamp': datetime(2025, 1, 1, 10, tzinfo=UTC),
+        }
+
+    dependency = Dependency(
+        'dependency',
+        messages(),
+        downtime_check=downtime_check,
+        recovery_check=recovery_check,
+    )
+
+    async def dependent():
+        yield {
+            'event_timestamp': datetime(2025, 1, 1, 10, tzinfo=UTC),
+        }
+
+    checkpoint = Checkpoint(
+        'test', dependent(), [dependency], cache=mock_cache
+    )
+
+    # Trigger downtime
+    await checkpoint.check_pulse(
+        datetime(2025, 1, 1, 10, tzinfo=UTC),
+        state={'offset': 0},
+    )
+    assert dependency.is_down is True
+    await checkpoint.check_pulse(
+        datetime(2025, 1, 1, 11, tzinfo=UTC),
+        state={'offset': 1},
+    )
+    downtime_check.assert_called()
+
+    # Trigger recovery
+    await checkpoint.heartbeat(
+        datetime(2025, 1, 1, 11, 1, tzinfo=UTC),
+    )
+    recovery_check.assert_called()
+    assert dependency.is_down is False
+
+
 def test_repr(checkpoint):
     """Should print representation without crashing."""
     assert str(checkpoint)

tests/test_core.py

@@ -127,6 +127,12 @@ async def test_conf(mocker: MockerFixture):
     assert c.__getattr__('group.id') == 'test'
     assert c.iterables == {}
 
+    # Missing prop
+    with pytest.raises(
+        AttributeError, match='object has no attribute "missing_prop"'
+    ):
+        assert c.missing_prop
+
     # Register iterable
     iterable = emoji()
     iterable_key = str(id(iterable))