fix: Make the async big-segment store status query awaitable

contract-tests/async_client_entity.py

@@ -185,7 +185,7 @@ def _context_response(self, c: Context) -> dict:
         return {"error": c.error}
 
     async def get_big_segment_store_status(self) -> dict:
-        status = self._client.big_segment_store_status_provider.status
+        status = await self._client.big_segment_store_status_provider.get_status()
         return {"available": status.available, "stale": status.stale}
 
     async def migration_variation(self, params: dict) -> dict:

ldclient/async_client.py

@@ -41,9 +41,9 @@
 from ldclient.impl.stubs import AsyncNullEventProcessor
 from ldclient.impl.util import log
 from ldclient.interfaces import (
+    AsyncBigSegmentStoreStatusProvider,
     AsyncFeatureStore,
     AsyncFlagTracker,
-    BigSegmentStoreStatusProvider,
     DataSourceStatusProvider,
     DataStoreStatusProvider
 )
@@ -691,13 +691,13 @@ async def __try_execute_stage(self, method: str, hook_name: str, block: Callable
             return {}
 
     @property
-    def big_segment_store_status_provider(self) -> BigSegmentStoreStatusProvider:
+    def big_segment_store_status_provider(self) -> AsyncBigSegmentStoreStatusProvider:
         """
         Returns an interface for tracking the status of a Big Segment store.
 
-        The :class:`ldclient.interfaces.BigSegmentStoreStatusProvider` has methods for checking
-        whether the Big Segment store is (as far as the SDK knows) currently operational and
-        tracking changes in this status.
+        The :class:`ldclient.interfaces.AsyncBigSegmentStoreStatusProvider` has methods for
+        checking whether the Big Segment store is (as far as the SDK knows) currently
+        operational and tracking changes in this status.
         """
         return self.__big_segment_store_manager.status_provider
 

ldclient/impl/async_big_segments.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple
+from typing import Awaitable, Callable, Optional, Tuple
 
 from expiringdict import ExpiringDict
 
@@ -7,17 +7,48 @@
 from ldclient.impl.aio.concurrency import AsyncRepeatingTask
 from ldclient.impl.big_segments_common import (
     EMPTY_MEMBERSHIP,
-    BigSegmentStoreStatusProviderImpl,
     _hash_for_user_key,
     is_stale
 )
+from ldclient.impl.listeners import Listeners
 from ldclient.impl.util import log
 from ldclient.interfaces import (
-    BigSegmentStoreStatus,
-    BigSegmentStoreStatusProvider
+    AsyncBigSegmentStoreStatusProvider,
+    BigSegmentStoreStatus
 )
 
 
+class AsyncBigSegmentStoreStatusProviderImpl(AsyncBigSegmentStoreStatusProvider):
+    """
+    Default implementation of the AsyncBigSegmentStoreStatusProvider interface.
+
+    Mirrors :class:`BigSegmentStoreStatusProviderImpl`, except the status getter passed in
+    is a coroutine, so :meth:`get_status` is a coroutine too and awaits it.
+    """
+
+    def __init__(self, status_getter: Callable[[], Awaitable[BigSegmentStoreStatus]]):
+        self.__status_getter = status_getter
+        self.__status_listeners = Listeners()
+        self.__last_status = None  # type: Optional[BigSegmentStoreStatus]
+
+    async def get_status(self) -> BigSegmentStoreStatus:
+        return await self.__status_getter()
+
+    def add_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
+        self.__status_listeners.add(listener)
+
+    def remove_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
+        self.__status_listeners.remove(listener)
+
+    def _update_status(self, new_status: BigSegmentStoreStatus):
+        last = self.__last_status
+        if last is None:
+            self.__last_status = new_status
+        elif new_status.available != last.available or new_status.stale != last.stale:
+            self.__last_status = new_status
+            self.__status_listeners.notify(new_status)
+
+
 class AsyncBigSegmentStoreManager:
     """
     Internal component that decorates the Big Segment store with caching behavior, and also polls the
@@ -28,7 +59,7 @@ def __init__(self, config: AsyncBigSegmentsConfig):
         self.__store = config.store
 
         self.__stale_after_millis = config.stale_after * 1000
-        self.__status_provider = BigSegmentStoreStatusProviderImpl(self.get_status)
+        self.__status_provider = AsyncBigSegmentStoreStatusProviderImpl(self.get_status)
         self.__last_status = None  # type: Optional[BigSegmentStoreStatus]
         self.__poll_task = None  # type: Optional[AsyncRepeatingTask]
 
@@ -49,7 +80,7 @@ async def stop(self):
             await self.__store.stop()
 
     @property
-    def status_provider(self) -> BigSegmentStoreStatusProvider:
+    def status_provider(self) -> AsyncBigSegmentStoreStatusProvider:
         return self.__status_provider
 
     async def get_user_membership(self, user_key: str) -> Tuple[Optional[dict], str]:
@@ -74,14 +105,16 @@ async def get_user_membership(self, user_key: str) -> Tuple[Optional[dict], str]
             return membership, BigSegmentsStatus.STORE_ERROR
         return membership, BigSegmentsStatus.STALE if status.stale else BigSegmentsStatus.HEALTHY
 
-    def get_status(self) -> BigSegmentStoreStatus:
+    async def get_status(self) -> BigSegmentStoreStatus:
         """Return the most recently polled status.
 
-        When no status has been cached yet, the sync variant polls the store
-        inline; the async variant (whose status getter cannot await) reports
-        the store as unavailable until the polling task has run.
+        When no status has been cached yet, poll the store inline and wait for the
+        result, so the status is accurate even if called immediately after start().
         """
-        return self.__last_status or BigSegmentStoreStatus(False, False)
+        status = self.__last_status
+        if status is None:
+            status = await self.poll_store_and_update_status()
+        return status
 
     async def poll_store_and_update_status(self) -> BigSegmentStoreStatus:
         new_status = BigSegmentStoreStatus(False, False)  # default to "unavailable" if we don't get a new status below

ldclient/impl/big_segments.py

@@ -1,15 +1,15 @@
-from typing import Optional, Tuple
+from typing import Callable, Optional, Tuple
 
 from expiringdict import ExpiringDict
 
 from ldclient.config import BigSegmentsConfig
 from ldclient.evaluation import BigSegmentsStatus
 from ldclient.impl.big_segments_common import (
     EMPTY_MEMBERSHIP,
-    BigSegmentStoreStatusProviderImpl,
     _hash_for_user_key,
     is_stale
 )
+from ldclient.impl.listeners import Listeners
 from ldclient.impl.repeating_task import RepeatingTask
 from ldclient.impl.util import log
 from ldclient.interfaces import (
@@ -18,6 +18,39 @@
 )
 
 
+class BigSegmentStoreStatusProviderImpl(BigSegmentStoreStatusProvider):
+    """
+    Default implementation of the BigSegmentStoreStatusProvider interface.
+
+    The real implementation of getting the status is in the big segment store manager - we pass in a lambda that
+    allows us to get the current status from that class. So this class provides a facade for that, and
+    also adds the listener mechanism.
+    """
+
+    def __init__(self, status_getter: Callable[[], BigSegmentStoreStatus]):
+        self.__status_getter = status_getter
+        self.__status_listeners = Listeners()
+        self.__last_status = None  # type: Optional[BigSegmentStoreStatus]
+
+    @property
+    def status(self) -> BigSegmentStoreStatus:
+        return self.__status_getter()
+
+    def add_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
+        self.__status_listeners.add(listener)
+
+    def remove_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
+        self.__status_listeners.remove(listener)
+
+    def _update_status(self, new_status: BigSegmentStoreStatus):
+        last = self.__last_status
+        if last is None:
+            self.__last_status = new_status
+        elif new_status.available != last.available or new_status.stale != last.stale:
+            self.__last_status = new_status
+            self.__status_listeners.notify(new_status)
+
+
 class BigSegmentStoreManager:
     """
     Internal component that decorates the Big Segment store with caching behavior, and also polls the

ldclient/impl/big_segments_common.py

@@ -1,21 +1,14 @@
 """
-Shared, I/O-free big-segments status provider used by both the sync
-:mod:`ldclient.impl.big_segments` and the async
-:mod:`ldclient.impl.async_big_segments`. It holds the last known status and
-notifies listeners; nothing here touches the store or network, so it is
-identical across the two managers.
+This module holds the shared, I/O-free helpers for big-segments status
+tracking. The helpers (`_hash_for_user_key`, `is_stale`, `EMPTY_MEMBERSHIP`)
+are shared by both the sync :mod:`ldclient.impl.big_segments` and the async
+:mod:`ldclient.impl.async_big_segments` managers; nothing here touches the
+store or network.
 """
 
 import base64
 import time
 from hashlib import sha256
-from typing import Callable, Optional
-
-from ldclient.impl.listeners import Listeners
-from ldclient.interfaces import (
-    BigSegmentStoreStatus,
-    BigSegmentStoreStatusProvider
-)
 
 # use EMPTY_MEMBERSHIP as a singleton whenever a membership query returns None; it's safe to reuse it
 # because we will never modify the membership properties after they're queried
@@ -28,36 +21,3 @@ def _hash_for_user_key(user_key: str) -> str:
 
 def is_stale(timestamp: int, stale_after_millis) -> bool:
     return (timestamp is None) or ((int(time.time() * 1000) - timestamp) >= stale_after_millis)
-
-
-class BigSegmentStoreStatusProviderImpl(BigSegmentStoreStatusProvider):
-    """
-    Default implementation of the BigSegmentStoreStatusProvider interface.
-
-    The real implementation of getting the status is in the big segment store manager - we pass in a lambda that
-    allows us to get the current status from that class. So this class provides a facade for that, and
-    also adds the listener mechanism.
-    """
-
-    def __init__(self, status_getter: Callable[[], BigSegmentStoreStatus]):
-        self.__status_getter = status_getter
-        self.__status_listeners = Listeners()
-        self.__last_status = None  # type: Optional[BigSegmentStoreStatus]
-
-    @property
-    def status(self) -> BigSegmentStoreStatus:
-        return self.__status_getter()
-
-    def add_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
-        self.__status_listeners.add(listener)
-
-    def remove_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
-        self.__status_listeners.remove(listener)
-
-    def _update_status(self, new_status: BigSegmentStoreStatus):
-        last = self.__last_status
-        if last is None:
-            self.__last_status = new_status
-        elif new_status.available != last.available or new_status.stale != last.stale:
-            self.__last_status = new_status
-            self.__status_listeners.notify(new_status)

ldclient/interfaces.py

@@ -894,9 +894,8 @@ def status(self) -> BigSegmentStoreStatus:
         """
         Gets the current status of the store.
 
-        Before the first poll completes, the synchronous SDK performs a blocking store query and
-        returns its result, while the async SDK (``AsyncLDClient``) returns the last polled status
-        without blocking -- ``available=False`` until the first background poll completes.
+        Before the first poll completes, this performs a blocking store query and returns its
+        result, so the status is accurate even if called immediately after the client starts.
 
         :return: the status
         """
@@ -925,6 +924,57 @@ def remove_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) ->
         pass
 
 
+class AsyncBigSegmentStoreStatusProvider(ABC):
+    """
+    Async interface for querying the status of a Big Segment store, for use with the async
+    client (:class:`ldclient.async_client.AsyncLDClient`). It mirrors
+    :class:`BigSegmentStoreStatusProvider`, except the current status is read with the
+    coroutine :meth:`get_status` instead of a synchronous ``status`` property (a property
+    cannot await).
+
+    .. caution::
+        This feature is experimental and should NOT be considered ready for production
+        use. It may change or be removed without notice and is not subject to backwards
+        compatibility guarantees.
+
+    An implementation of this abstract class is returned by
+    :meth:`ldclient.async_client.AsyncLDClient.big_segment_store_status_provider`. Application
+    code never needs to implement this interface.
+    """
+
+    @abstractmethod
+    async def get_status(self) -> BigSegmentStoreStatus:
+        """
+        Gets the current status of the store.
+
+        Before the first poll completes, this awaits a store query and returns its result, so
+        the status is accurate even if called immediately after the client starts.
+
+        :return: the status
+        """
+        pass
+
+    @abstractmethod
+    def add_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
+        """
+        Subscribes for notifications of status changes. Behaves identically to
+        :meth:`BigSegmentStoreStatusProvider.add_listener`.
+
+        :param listener: the listener to add
+        """
+        pass
+
+    @abstractmethod
+    def remove_listener(self, listener: Callable[[BigSegmentStoreStatus], None]) -> None:
+        """
+        Unsubscribes from notifications of status changes.
+
+        :param listener: a listener that was previously added with :func:`add_listener()`; if it was not,
+            this method does nothing
+        """
+        pass
+
+
 class DataSourceState(Enum):
     """
     Enumeration representing the states a data source can be in at any given time.

ldclient/testing/impl/test_async_big_segments.py

@@ -302,17 +302,49 @@ async def test_stop_stops_store():
 
 
 @pytest.mark.asyncio
-async def test_status_provider_status_is_synchronous():
-    """BigSegmentStoreStatusProvider.status must be readable synchronously without an await."""
+async def test_status_provider_get_status_after_poll():
+    """AsyncBigSegmentStoreStatusProvider.get_status() reflects the last polled status."""
     store = MockAsyncBigSegmentStore()
     store.setup_metadata_always_up_to_date()
 
     manager = await make_started_manager(store)
     try:
         # Poll once to populate __last_status
         await manager.poll_store_and_update_status()
-        # status property is sync — calling it should not raise
-        status = manager.status_provider.status
+        status = await manager.status_provider.get_status()
         assert status.available is True
     finally:
         await manager.stop()
+
+
+@pytest.mark.asyncio
+async def test_get_status_polls_inline_before_background_poll_runs():
+    """
+    get_status() must not race the background AsyncRepeatingTask: calling it immediately
+    after start() -- before the poll task has had a chance to run -- should still report an
+    accurate, available status by polling inline, mirroring the sync manager's behavior.
+    """
+    store = MockAsyncBigSegmentStore()
+    store.setup_metadata_always_up_to_date()
+
+    config = AsyncBigSegmentsConfig(store=store)
+    manager = AsyncBigSegmentStoreManager(config)
+    manager.start()
+    try:
+        status = await manager.get_status()
+        assert status.available is True
+        assert status.stale is False
+    finally:
+        await manager.stop()
+
+
+@pytest.mark.asyncio
+async def test_get_status_with_no_store_configured():
+    """With no store configured, get_status() should report unavailable rather than hang or error."""
+    config = AsyncBigSegmentsConfig(store=None)
+    manager = AsyncBigSegmentStoreManager(config)
+    try:
+        status = await manager.get_status()
+        assert status.available is False
+    finally:
+        await manager.stop()