diff --git a/ldclient/impl/aio/concurrency.py b/ldclient/impl/aio/concurrency.py index fdc51403..97a2b5c2 100644 --- a/ldclient/impl/aio/concurrency.py +++ b/ldclient/impl/aio/concurrency.py @@ -223,6 +223,16 @@ def stop(self): if task is not None and task is not asyncio.current_task(): task.cancel() + async def wait_stopped(self): + """Waits for the task to finish unwinding after ``stop()``. A no-op if + the task never started or is the current task.""" + task = self.__task + if task is not None and task is not asyncio.current_task(): + try: + await task + except asyncio.CancelledError: + pass + async def _run(self): try: if self.__initial_delay > 0: diff --git a/ldclient/impl/datasource/async_feature_requester.py b/ldclient/impl/datasource/async_feature_requester.py new file mode 100644 index 00000000..a9b0e038 --- /dev/null +++ b/ldclient/impl/datasource/async_feature_requester.py @@ -0,0 +1,56 @@ +""" +Default implementation of feature flag polling requests. +""" + +import json +from collections import namedtuple +from typing import Optional +from urllib import parse + +from ldclient.impl.aio.transport import AsyncHTTPTransport +from ldclient.impl.datasource.datasource_common import FDV1_POLLING_ENDPOINT +from ldclient.impl.util import _headers, log, throw_if_unsuccessful_response +from ldclient.interfaces import FeatureRequester +from ldclient.versioned_data_kind import FEATURES, SEGMENTS + +CacheEntry = namedtuple('CacheEntry', ['data', 'etag']) + + +class AsyncFeatureRequesterImpl(FeatureRequester): + def __init__(self, config, transport: Optional[AsyncHTTPTransport] = None): + self._cache: dict = dict() + # Only close the transport on shutdown if we created it; an injected + # transport is owned by the caller. + self._owns_transport = transport is None + self._transport = transport if transport is not None else AsyncHTTPTransport(config) + self._config = config + self._poll_uri = config.base_uri + FDV1_POLLING_ENDPOINT + if config.payload_filter_key is not None: + self._poll_uri += '?%s' % parse.urlencode({'filter': config.payload_filter_key}) + + async def get_all_data(self): + uri = self._poll_uri + hdrs = _headers(self._config) + cache_entry = self._cache.get(uri) + hdrs['Accept-Encoding'] = 'gzip' + if cache_entry is not None: + hdrs['If-None-Match'] = cache_entry.etag + r = await self._transport.request('GET', uri, headers=hdrs) + throw_if_unsuccessful_response(r) + if r.status == 304 and cache_entry is not None: + data = cache_entry.data + etag = cache_entry.etag + from_cache = True + else: + data = json.loads(r.body) + etag = r.headers.get('ETag') + from_cache = False + if etag is not None: + self._cache[uri] = CacheEntry(data=data, etag=etag) + log.debug("%s response status:[%d] From cache? [%s] ETag:[%s]", uri, r.status, from_cache, etag) + + return {FEATURES: data['flags'], SEGMENTS: data['segments']} + + async def close(self): + if self._owns_transport: + await self._transport.close() diff --git a/ldclient/impl/datasource/async_polling.py b/ldclient/impl/datasource/async_polling.py new file mode 100644 index 00000000..773fb9c0 --- /dev/null +++ b/ldclient/impl/datasource/async_polling.py @@ -0,0 +1,91 @@ +""" +Default implementation of the polling component. +""" + +# currently excluded from documentation - see docs/README.md + +import time +from typing import Optional + +from ldclient.async_config import AsyncConfig +from ldclient.impl.aio.concurrency import AsyncEvent, AsyncRepeatingTask +from ldclient.impl.datasource.async_feature_requester import ( + AsyncFeatureRequesterImpl +) +from ldclient.impl.datasource.datasource_common import sink_or_store +from ldclient.impl.util import ( + UnsuccessfulResponseException, + http_error_message, + is_http_error_recoverable, + log +) +from ldclient.interfaces import ( + AsyncFeatureStore, + AsyncUpdateProcessor, + DataSourceErrorInfo, + DataSourceErrorKind, + DataSourceState +) + + +class AsyncPollingUpdateProcessor(AsyncUpdateProcessor): + def __init__(self, config: AsyncConfig, requester: AsyncFeatureRequesterImpl, store: AsyncFeatureStore, ready: AsyncEvent): + self._config = config + self._data_source_update_sink = config.data_source_update_sink + self._requester = requester + self._store = store + self._ready = ready + self._task = AsyncRepeatingTask("ldclient.datasource.polling", config.poll_interval, 0, self._fetch_and_store) + + def start(self): + log.info("Starting AsyncPollingUpdateProcessor with request interval: " + str(self._config.poll_interval)) + self._task.start() + + def initialized(self): + return self._ready.is_set() is True and self._store.initialized is True + + async def stop(self): + self.__stop_with_error_info(None) + # Wait for the in-flight poll to finish unwinding before closing the + # transport, so awaiting stop() guarantees background work has stopped + # and we don't close the HTTP transport out from under a live request. + await self._task.wait_stopped() + await self._requester.close() + + def __stop_with_error_info(self, error: Optional[DataSourceErrorInfo]): + log.info("Stopping AsyncPollingUpdateProcessor") + self._task.stop() + + if self._data_source_update_sink is None: + return + + self._data_source_update_sink.update_status(DataSourceState.OFF, error) + + async def _fetch_and_store(self): + try: + all_data = await self._requester.get_all_data() + await sink_or_store(self._data_source_update_sink, self._store).init(all_data) + if not self._ready.is_set() and self._store.initialized: + log.info("AsyncPollingUpdateProcessor initialized ok") + self._ready.set() + + # Signal VALID once the store is populated. + if self._store.initialized and self._data_source_update_sink is not None: + self._data_source_update_sink.update_status(DataSourceState.VALID, None) + except UnsuccessfulResponseException as e: + error_info = DataSourceErrorInfo(DataSourceErrorKind.ERROR_RESPONSE, e.status, time.time(), str(e)) + + http_error_message_result = http_error_message(e.status, "polling request") + if not is_http_error_recoverable(e.status): + log.error(http_error_message_result) + self._ready.set() # if client is initializing, make it stop waiting; has no effect if already inited + self.__stop_with_error_info(error_info) + else: + log.warning(http_error_message_result) + + if self._data_source_update_sink is not None: + self._data_source_update_sink.update_status(DataSourceState.INTERRUPTED, error_info) + except Exception as e: + log.exception('Error: Exception encountered when updating flags. %s' % e) + if self._data_source_update_sink is not None: + self._data_source_update_sink.update_status(DataSourceState.INTERRUPTED, DataSourceErrorInfo(DataSourceErrorKind.UNKNOWN, 0, time.time(), str(e))) diff --git a/ldclient/testing/impl/datasource/test_async_polling.py b/ldclient/testing/impl/datasource/test_async_polling.py new file mode 100644 index 00000000..0a330fbd --- /dev/null +++ b/ldclient/testing/impl/datasource/test_async_polling.py @@ -0,0 +1,439 @@ +""" +Tests for AsyncFeatureRequesterImpl and AsyncPollingUpdateProcessor. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from ldclient.config import Config +from ldclient.impl.aio.transport_types import TransportResponse +from ldclient.impl.datasource.async_feature_requester import ( + AsyncFeatureRequesterImpl +) +from ldclient.impl.datasource.async_polling import AsyncPollingUpdateProcessor +from ldclient.impl.util import UnsuccessfulResponseException +from ldclient.interfaces import DataSourceErrorKind, DataSourceState +from ldclient.testing.mock_async_components import MockAsyncFeatureStore +from ldclient.versioned_data_kind import FEATURES, SEGMENTS + +# Sample data returned by a successful poll +SAMPLE_FLAGS = {'flagkey': {'key': 'flagkey', 'version': 1, 'deleted': False}} +SAMPLE_SEGMENTS = {'segkey': {'key': 'segkey', 'version': 1, 'deleted': False}} +SAMPLE_DATA = {FEATURES: SAMPLE_FLAGS, SEGMENTS: SAMPLE_SEGMENTS} + + +def make_config(**kwargs): + """Create a Config with a short poll_interval for tests. + + Config enforces a minimum poll_interval of 30s, so we patch the property. + """ + return Config('SDK_KEY', **kwargs) + + +def make_processor(config=None, store=None, ready=None, requester=None): + if config is None: + config = make_config() + if store is None: + store = MockAsyncFeatureStore() + if ready is None: + ready = asyncio.Event() + if requester is None: + requester = MagicMock() + requester.close = AsyncMock() + return AsyncPollingUpdateProcessor( + config=config, + requester=requester, + store=store, + ready=ready, + ) + + +def make_transport(*responses: TransportResponse): + """Create a stub transport whose request() returns the given responses in order.""" + transport = MagicMock() + transport.request = AsyncMock(side_effect=list(responses)) + return transport + + +class TestAsyncFeatureRequesterImpl: + @pytest.mark.asyncio + async def test_successful_response_returns_flags_and_segments(self): + import json + config = make_config() + transport = make_transport( + TransportResponse(200, {}, json.dumps({'flags': SAMPLE_FLAGS, 'segments': SAMPLE_SEGMENTS})) + ) + requester = AsyncFeatureRequesterImpl(config, transport) + + data = await requester.get_all_data() + + assert data[FEATURES] == SAMPLE_FLAGS + assert data[SEGMENTS] == SAMPLE_SEGMENTS + + @pytest.mark.asyncio + async def test_304_not_modified_returns_cached_data(self): + from ldclient.impl.datasource.async_feature_requester import CacheEntry + + config = make_config() + transport = make_transport(TransportResponse(304, {}, '')) + requester = AsyncFeatureRequesterImpl(config, transport) + + # Pre-populate the cache with a known etag and data + cached_data = {'flags': SAMPLE_FLAGS, 'segments': SAMPLE_SEGMENTS} + requester._cache[requester._poll_uri] = CacheEntry(data=cached_data, etag='"abc"') + + data = await requester.get_all_data() + + # 304 returns the cached data rather than None + assert data[FEATURES] == SAMPLE_FLAGS + assert data[SEGMENTS] == SAMPLE_SEGMENTS + # The cached etag is sent as If-None-Match + headers = transport.request.call_args.kwargs['headers'] + assert headers['If-None-Match'] == '"abc"' + + @pytest.mark.asyncio + async def test_etag_and_data_stored_after_successful_response(self): + config = make_config() + transport = make_transport( + TransportResponse(200, {'ETag': '"v1"'}, '{"flags": {}, "segments": {}}') + ) + requester = AsyncFeatureRequesterImpl(config, transport) + + await requester.get_all_data() + + cache_entry = requester._cache.get(requester._poll_uri) + assert cache_entry is not None + assert cache_entry.etag == '"v1"' + + @pytest.mark.asyncio + async def test_http_error_raises_unsuccessful_response_exception(self): + config = make_config() + transport = make_transport(TransportResponse(401, {}, '')) + requester = AsyncFeatureRequesterImpl(config, transport) + + with pytest.raises(UnsuccessfulResponseException) as exc_info: + await requester.get_all_data() + + assert exc_info.value.status == 401 + + @pytest.mark.asyncio + async def test_payload_filter_key_appended_to_uri(self): + config = Config('SDK_KEY', payload_filter_key='my-filter') + requester = AsyncFeatureRequesterImpl(config, MagicMock()) + + assert 'filter=my-filter' in requester._poll_uri + + @pytest.mark.asyncio + async def test_close_closes_owned_transport(self): + with patch('ldclient.impl.datasource.async_feature_requester.AsyncHTTPTransport') as MockTransport: + MockTransport.return_value.close = AsyncMock() + requester = AsyncFeatureRequesterImpl(make_config()) # no transport -> owns one + + await requester.close() + + MockTransport.return_value.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_close_does_not_close_injected_transport(self): + transport = MagicMock() + transport.close = AsyncMock() + requester = AsyncFeatureRequesterImpl(make_config(), transport) + + await requester.close() + + transport.close.assert_not_called() + + +class TestAsyncPollingUpdateProcessor: + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_successful_poll_initializes_store_and_sets_ready(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + processor = make_processor(config=config, store=store, ready=ready) + + processor._requester.get_all_data = AsyncMock(return_value=SAMPLE_DATA) + + processor.start() + await asyncio.wait_for(ready.wait(), timeout=2.0) + + assert ready.is_set() + assert store.initialized + assert processor.initialized() + assert len(store.inits) >= 1 + assert store.inits[0] == SAMPLE_DATA + + await processor.stop() + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_unrecoverable_http_error_stops_polling_and_sets_ready(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + processor = make_processor(config=config, store=store, ready=ready) + + mock_requester = AsyncMock(side_effect=UnsuccessfulResponseException(401)) + processor._requester.get_all_data = mock_requester + + processor.start() + await asyncio.wait_for(ready.wait(), timeout=2.0) + + assert ready.is_set() + assert not processor.initialized() + + # The polling task must have stopped itself: no further polls occur. + await asyncio.sleep(0.05) + snapshot = mock_requester.call_count + await asyncio.sleep(0.05) + assert mock_requester.call_count == snapshot + + await processor.stop() + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_recoverable_http_error_continues_polling(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + processor = make_processor(config=config, store=store, ready=ready) + + call_count = 0 + + async def get_all_data(): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise UnsuccessfulResponseException(500) + return SAMPLE_DATA + + processor._requester.get_all_data = get_all_data + + processor.start() + await asyncio.wait_for(ready.wait(), timeout=2.0) + + assert ready.is_set() + assert processor.initialized() + assert call_count >= 3 + + await processor.stop() + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_general_exception_does_not_stop_polling(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + processor = make_processor(config=config, store=store, ready=ready) + + call_count = 0 + + async def get_all_data(): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise RuntimeError("transient error") + return SAMPLE_DATA + + processor._requester.get_all_data = get_all_data + + processor.start() + + # A generic error must NOT release start_wait; _ready stays unset until a + # later poll succeeds. Wait until the store is actually initialized + # (call_count reaches 3) to verify polling kept running. + async def wait_for_initialized(): + while not store.initialized: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_initialized(), timeout=2.0) + + assert ready.is_set() + assert call_count >= 3 + + await processor.stop() + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_general_exception_does_not_set_ready(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + processor = make_processor(config=config, store=store, ready=ready) + + async def get_all_data(): + raise RuntimeError("transient error") + + processor._requester.get_all_data = get_all_data + + processor.start() + # Let several polls run; all raise. A transient error must not release + # start_wait (mirrors sync, which leaves _ready unset here). + await asyncio.sleep(0.05) + + assert not ready.is_set() + assert not processor.initialized() + + await processor.stop() + + @pytest.mark.asyncio + async def test_stop_closes_requester(self): + processor = make_processor() + + await processor.stop() + + processor._requester.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_stop_awaits_poll_before_closing_transport(self): + # The in-flight poll must finish unwinding before the transport is + # closed, so we never close it out from under a live request. + order = [] + started = asyncio.Event() + + async def slow_poll(): + started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + order.append('poll_done') + raise + + async def close(): + order.append('transport_closed') + + requester = MagicMock() + requester.get_all_data = slow_poll + requester.close = close + + processor = make_processor(requester=requester) + processor.start() + await asyncio.wait_for(started.wait(), timeout=1.0) + + await processor.stop() + + assert order == ['poll_done', 'transport_closed'] + + @pytest.mark.asyncio + async def test_stop_cancels_polling_task_cleanly(self): + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + processor = make_processor(config=config, store=store, ready=ready) + + poll_count = 0 + + async def slow_poll(): + nonlocal poll_count + poll_count += 1 + await asyncio.sleep(60) # Would block indefinitely without cancel + + processor._requester.get_all_data = slow_poll + + processor.start() + await asyncio.sleep(0.05) # Let the task start + + # stop() should return promptly even though the poll is "sleeping" + await asyncio.wait_for(processor.stop(), timeout=1.0) + + # No further polls occur after stopping + await asyncio.sleep(0.05) + assert poll_count == 1 + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_unrecoverable_error_updates_sink_to_off(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + + sink = AsyncMock() + config._data_source_update_sink = sink + + processor = make_processor(config=config, store=store, ready=ready) + processor._data_source_update_sink = sink + + processor._requester.get_all_data = AsyncMock( + side_effect=UnsuccessfulResponseException(403) + ) + + processor.start() + await asyncio.wait_for(ready.wait(), timeout=2.0) + + # Verify the sink was told to go OFF + calls = [call for call in sink.update_status.call_args_list if call.args[0] == DataSourceState.OFF] + assert len(calls) >= 1 + error_info = calls[0].args[1] + assert error_info.kind == DataSourceErrorKind.ERROR_RESPONSE + assert error_info.status_code == 403 + + await processor.stop() + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_successful_poll_updates_sink_to_valid(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + store = MockAsyncFeatureStore() + ready = asyncio.Event() + config = make_config() + + # The processor checks self._store.initialized before sending VALID. + # Use an AsyncMock sink whose init() also marks the underlying store as initialized + # so that both the sink-path and the initialized check work correctly. + sink = AsyncMock() + + async def _init_and_store(data): + await store.init(data) + + sink.init = _init_and_store + config._data_source_update_sink = sink + + processor = make_processor(config=config, store=store, ready=ready) + processor._data_source_update_sink = sink + + processor._requester.get_all_data = AsyncMock(return_value=SAMPLE_DATA) + + processor.start() + await asyncio.wait_for(ready.wait(), timeout=2.0) + + valid_calls = [c for c in sink.update_status.call_args_list if c.args[0] == DataSourceState.VALID] + assert len(valid_calls) >= 1 + + await processor.stop() + + @pytest.mark.asyncio + async def test_initialized_returns_false_before_first_poll(self): + processor = make_processor() + assert not processor.initialized() + + @pytest.mark.asyncio + @patch('ldclient.config.Config.poll_interval', new_callable=MagicMock) + async def test_second_start_call_raises(self, mock_interval): + mock_interval.__get__ = MagicMock(return_value=0) + + processor = make_processor() + processor._requester.get_all_data = AsyncMock(return_value=SAMPLE_DATA) + + processor.start() + # Like a thread, the polling task can only be started once + with pytest.raises(RuntimeError): + processor.start() + + await processor.stop()