Merge pull request #160 from python-astrodynamics/feature/ratelimit-callback-wait

Fix rate limit callback waits and support the predicate cache on trio

Author: RazerM

newsfragments/160.changed.2.rst

@@ -0,0 +1,2 @@
+Predicate caching for keyword validation now works when using :class:`~spacetrack.aio.AsyncSpaceTrackClient` with trio.
+Previously, the on-disk cache was skipped because the file locking implementation only supported asyncio.

newsfragments/160.changed.3.rst

@@ -0,0 +1 @@
+:class:`~spacetrack.aio.AsyncSpaceTrackClient` no longer performs blocking cache file reads and writes on the event loop; they now run in a worker thread.

newsfragments/160.changed.rst

@@ -0,0 +1,2 @@
+Rate limit waits now guarantee that the ``callback`` has finished before the request is retried, in all clients.
+As a result, an exception raised by the callback now propagates instead of being silently discarded: directly in :class:`~spacetrack.base.SpaceTrackClient`, and as an :class:`ExceptionGroup` on asyncio, matching the existing trio behaviour.

newsfragments/160.fixed.rst

@@ -0,0 +1 @@
+:class:`~spacetrack.aio.AsyncSpaceTrackClient` now accepts a plain function as the rate limit ``callback``, as shown in the documentation, in addition to a coroutine function.

pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
     "Programming Language :: Python :: 3.15",
 ]
 dependencies = [
+    "anyio>=4.10",
     "filelock>=3.17.0",
     "httpx2>=2.9.1",
     "logbook>=1.10.1",
@@ -31,7 +32,6 @@ dependencies = [
     "python-dateutil>=2.9.0.post0; python_version < '3.11'",
     "represent>=2.1",
     "rush>=2021.4.0",
-    "sniffio>=1.3.1",
 ]
 
 [project.urls]

src/spacetrack/aio.py

@@ -1,25 +1,26 @@
-import asyncio
+import inspect
 import time
 import weakref
+from functools import partial
 
+import anyio
 import httpx2
 import outcome
-import sniffio
-from filelock import AsyncFileLock
+from filelock import Timeout
 from httpx2 import USE_CLIENT_DEFAULT
 
 from .base import (
     BASE_URL,
-    AcquireLock,
+    AcquireFileLock,
     Event,
     IterContent,
     IterLines,
     NormalRequest,
     RateLimitWait,
     ReadResponse,
-    ReleaseLock,
+    ReleaseFileLock,
+    RunBlocking,
     SpaceTrackClient,
-    UnsupportedAsyncLibrary,
     logger,
 )
 
@@ -38,7 +39,6 @@ class AsyncSpaceTrackClient(SpaceTrackClient):
     be an ``httpx2.AsyncClient``.
     """
 
-    _file_lock_cls = AsyncFileLock
     _httpx_client_cls = httpx2.AsyncClient
 
     def __init__(
@@ -62,7 +62,6 @@ def __init__(
             additional_rate_limit=additional_rate_limit,
             cache_path=cache_path,
         )
-        self._ratelimit_tasks = set()
 
     def _setup_finalizer(self):
         self._finalizer = weakref.finalize(
@@ -90,15 +89,32 @@ async def _handle_event(self, event):
             return _iter_content_generator(event.response, event.decode)
         elif isinstance(event, RateLimitWait):
             await self._ratelimit_wait(event.duration)
-        elif isinstance(event, AcquireLock):
-            if (
-                isinstance(event.lock, AsyncFileLock)
-                and sniffio.current_async_library() != "asyncio"
-            ):
-                raise UnsupportedAsyncLibrary
-            await event.lock.acquire()
-        elif isinstance(event, ReleaseLock):
-            await event.lock.release()
+        elif isinstance(event, AcquireFileLock):
+            # Mirror filelock's AsyncFileLock structure with backend-agnostic
+            # primitives: non-blocking acquire attempts in a worker thread,
+            # with a cancellable sleep between attempts. Each attempt is
+            # shielded so that a cancellation cannot discard an acquired
+            # lock; cancellation is delivered at the sleep instead.
+            while True:
+                with anyio.CancelScope(shield=True):
+                    try:
+                        await anyio.to_thread.run_sync(
+                            partial(event.lock.acquire, blocking=False)
+                        )
+                    except Timeout:
+                        acquired = False
+                    else:
+                        acquired = True
+                if acquired:
+                    break
+                await anyio.sleep(0.05)
+        elif isinstance(event, ReleaseFileLock):
+            # Shielded so that a cancelled scope cannot skip the release and
+            # leak the lock.
+            with anyio.CancelScope(shield=True):
+                await anyio.to_thread.run_sync(event.lock.release)
+        elif isinstance(event, RunBlocking):
+            return await anyio.to_thread.run_sync(event.func)
         else:
             raise RuntimeError(f"Unknown event type: {type(event)}")
 
@@ -240,29 +256,15 @@ async def _ratelimit_callback(self, until):
         logger.info("Rate limit reached. Sleeping for {:d} seconds.", duration)
 
         if self.callback is not None:
-            await self.callback(until)
+            result = self.callback(until)
+            if inspect.isawaitable(result):
+                await result
 
     async def _ratelimit_wait(self, duration):
-        async_library = sniffio.current_async_library()
-        if async_library == "asyncio":
-            await self._ratelimit_wait_asyncio(duration)
-        elif async_library == "trio":
-            await self._ratelimit_wait_trio(duration)
-
-    async def _ratelimit_wait_asyncio(self, duration):
-        until = time.monotonic() + duration
-        task = asyncio.create_task(self._ratelimit_callback(until))
-        self._ratelimit_tasks.add(task)
-        task.add_done_callback(self._ratelimit_tasks.discard)
-        await asyncio.sleep(duration)
-
-    async def _ratelimit_wait_trio(self, duration):
-        import trio
-
         until = time.monotonic() + duration
-        async with trio.open_nursery() as nursery:
-            nursery.start_soon(self._ratelimit_callback, until)
-            nursery.start_soon(trio.sleep, duration)
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(self._ratelimit_callback, until)
+            tg.start_soon(anyio.sleep, duration)
 
     async def get_predicates(self, class_, controller=None):
         """Get full predicate information for given request class, and cache

src/spacetrack/base.py

@@ -7,16 +7,17 @@
 import warnings
 import weakref
 from collections import OrderedDict
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from datetime import datetime, timedelta, timezone
 from functools import partial
 from json import JSONDecodeError
 from pathlib import Path
+from typing import Any
 from urllib.parse import quote
 
-import attr
 import httpx2
 import outcome
+from attrs import define
 from filelock import FileLock
 from httpx2 import USE_CLIENT_DEFAULT
 from logbook import Logger
@@ -79,48 +80,47 @@ class Event:
     pass
 
 
-@attr.s(slots=True)
+@define
 class NormalRequest(Event):
-    request = attr.ib()
-    stream = attr.ib(default=False)
-    follow_redirects = attr.ib(default=False)
+    request: httpx2.Request
+    stream: bool = False
+    follow_redirects: bool = False
 
 
-@attr.s(slots=True)
+@define
 class ReadResponse(Event):
-    response = attr.ib()
+    response: httpx2.Response
 
 
-@attr.s(slots=True)
+@define
 class IterLines(Event):
-    response = attr.ib()
+    response: httpx2.Response
 
 
-@attr.s(slots=True)
+@define
 class IterContent(Event):
-    response = attr.ib()
-    decode = attr.ib()
+    response: httpx2.Response
+    decode: bool
 
 
-@attr.s(slots=True)
+@define
 class RateLimitWait(Event):
-    duration = attr.ib()
+    duration: float
 
 
-@attr.s(slots=True)
-class AcquireLock(Event):
-    lock = attr.ib()
+@define
+class AcquireFileLock(Event):
+    lock: FileLock
 
 
-@attr.s(slots=True)
-class ReleaseLock(Event):
-    lock = attr.ib()
+@define
+class ReleaseFileLock(Event):
+    lock: FileLock
 
 
-class UnsupportedAsyncLibrary(Exception):
-    """Raised internally when an event cannot be handled with the active async
-    library.
-    """
+@define
+class RunBlocking(Event):
+    func: Callable[[], Any]
 
 
 class Predicate(ReprHelperMixin):
@@ -299,7 +299,6 @@ class SpaceTrackClient:
         Predicate("favorites", "str"),
     }
 
-    _file_lock_cls = FileLock
     _httpx_client_cls = httpx2.Client
 
     def __init__(
@@ -403,10 +402,12 @@ def _handle_event(self, event):
             return _iter_content_generator(event.response, event.decode)
         elif isinstance(event, RateLimitWait):
             self._ratelimit_wait(event.duration)
-        elif isinstance(event, AcquireLock):
+        elif isinstance(event, AcquireFileLock):
             event.lock.acquire()
-        elif isinstance(event, ReleaseLock):
+        elif isinstance(event, ReleaseFileLock):
             event.lock.release()
+        elif isinstance(event, RunBlocking):
+            return event.func()
         else:
             raise RuntimeError(f"Unknown event type: {type(event)}")
 
@@ -779,10 +780,20 @@ def _ratelimit_callback(self, until):
 
     def _ratelimit_wait(self, duration):
         until = time.monotonic() + duration
-        t = threading.Thread(target=self._ratelimit_callback, args=(until,))
+        callback_outcome = None
+
+        def run_callback():
+            nonlocal callback_outcome
+            callback_outcome = outcome.capture(self._ratelimit_callback, until)
+
+        t = threading.Thread(target=run_callback)
         t.daemon = True
         t.start()
         time.sleep(duration)
+        t.join()
+        # Match the async client, where a callback exception propagates and
+        # aborts the request.
+        callback_outcome.unwrap()
 
     def __getattr__(self, attr):
         if attr in self.request_controllers:
@@ -867,41 +878,33 @@ def _get_predicates_generator(self, class_, controller, *, force=False):
             hasher.update(key.encode())
             hashkey = hasher.hexdigest()[:16]
             cache_file = self._cache_path / f"predicates-{hashkey}.json"
-            predicates_data = self._read_cache_file(
-                cache_file, PREDICATE_CACHE_EXPIRY_TIME
+            read_cache = partial(
+                self._read_cache_file, cache_file, PREDICATE_CACHE_EXPIRY_TIME
             )
+            predicates_data = yield RunBlocking(read_cache)
 
             if predicates_data is None:
-                self._cache_path.mkdir(parents=True, exist_ok=True)
+                yield RunBlocking(
+                    partial(self._cache_path.mkdir, parents=True, exist_ok=True)
+                )
 
                 lock_file = cache_file.with_name(cache_file.name + ".lock")
-                lock = self._file_lock_cls(lock_file)
-                try:
-                    yield AcquireLock(lock)
-                except UnsupportedAsyncLibrary:
-                    if not force:
-                        # The file lock doesn't support Trio, skip predicate
-                        # checking by setting None
-                        self._predicates[key] = None
-                        return self._predicates[key]
-                    lock_acquired = False
-                else:
-                    lock_acquired = True
+                # thread_local=False because the async client acquires and
+                # releases the lock from different worker threads.
+                lock = FileLock(lock_file, thread_local=False)
+                yield AcquireFileLock(lock)
 
                 try:
-                    if lock_acquired:
-                        predicates_data = self._read_cache_file(
-                            cache_file, PREDICATE_CACHE_EXPIRY_TIME
-                        )
+                    predicates_data = yield RunBlocking(read_cache)
                     if predicates_data is None:
                         predicates_data = yield from self._download_predicate_data_generator(
                             class_, controller
                         )
-                        if lock_acquired:
-                            self._write_cache_file(cache_file, predicates_data)
+                        yield RunBlocking(
+                            partial(self._write_cache_file, cache_file, predicates_data)
+                        )
                 finally:
-                    if lock_acquired:
-                        yield ReleaseLock(lock)
+                    yield ReleaseFileLock(lock)
 
             predicate_objects = self._parse_predicates_data(predicates_data)