diff --git a/osprey_async_worker/src/osprey/async_worker/executor.py b/osprey_async_worker/src/osprey/async_worker/executor.py index a29221a9..ee55193d 100644 --- a/osprey_async_worker/src/osprey/async_worker/executor.py +++ b/osprey_async_worker/src/osprey/async_worker/executor.py @@ -47,6 +47,9 @@ logger = get_logger(__name__) +# Shared placeholder for the "make mypy happy" default below; always overwritten before use. +_UNSET_RESULT: NodeResult = Err(None) + _DEFAULT_MAX_ASYNC_PER_EXECUTION = 12 @@ -119,7 +122,7 @@ def _execute_sync( error_info_: list[NodeErrorInfo], ) -> NodeResult: """Execute a sync UDF inline. For pure computation only — no I/O.""" - execution_result: NodeResult = Err(None) + execution_result: NodeResult = _UNSET_RESULT try: execution_result = Ok(chain.executor.execute(execution_context=context)) except Exception as e: @@ -148,7 +151,7 @@ def _execute_legacy_sync( call_node: CallExecutor = chain.executor metric_tags += [f'udf:{call_node._udf.__class__.__name__}'] - execution_result: NodeResult = Err(None) + execution_result: NodeResult = _UNSET_RESULT try: with metrics.timed('udf_execution_duration', tags=metric_tags, sample_rate=0.01): execution_result = Ok(chain.executor.execute(execution_context=context)) @@ -227,17 +230,27 @@ async def _execute_async_udf( chain: DependencyChain, context: ExecutionContext, error_info_: list[NodeErrorInfo], + pre_resolved_arguments: Any | None = None, ) -> NodeResult: - """Execute a native async UDF. Awaited directly on the event loop.""" + """Execute a native async UDF. Awaited directly on the event loop. + + `pre_resolved_arguments`, when given, is the Arguments already computed for this same + chain (e.g. by `_enqueue_batches` while checking whether a batch would form), so this + skips a redundant `resolve_arguments` call for the same message. + """ async with semaphore: call_executor: CallExecutor = chain.executor # type: ignore udf: AsyncUDFBase[Any, Any] = call_executor._udf # type: ignore metric_tags = _get_metric_tags(context) + [f'udf:{udf.__class__.__name__}'] caught_exception: Exception | None = None - execution_result: NodeResult = Err(None) + execution_result: NodeResult = _UNSET_RESULT try: - resolved_arguments = udf.resolve_arguments(context, call_executor) + resolved_arguments = ( + pre_resolved_arguments + if pre_resolved_arguments is not None + else udf.resolve_arguments(context, call_executor) + ) with metrics.timed('udf_execution_duration', tags=metric_tags, sample_rate=0.01): result = await udf.async_execute(context, resolved_arguments) execution_result = Ok(udf.check_result_type(result)) @@ -320,12 +333,15 @@ async def _enqueue_batches( context: ExecutionContext, error_infos: list[NodeErrorInfo], ready_async: Sequence[DependencyChain], -) -> tuple[Sequence[DependencyChain], dict[asyncio.Task[Sequence[NodeResult]], Sequence[DependencyChain]]]: - """Collect batchable async chains and launch them as tasks. +) -> tuple[ + Sequence[tuple[DependencyChain, Any | None]], + dict[asyncio.Task[Sequence[NodeResult]], Sequence[DependencyChain]], +]: + """Launch batches and return the chains that remain. - Returns (remaining non-batched chains, dict of batch tasks -> chains). + A native async chain can include arguments that this function resolved. Legacy chains include `None`. """ - batch_chains: dict[tuple[type, str], list[tuple[DependencyChain, Any]]] = defaultdict(list) + batch_chains: dict[tuple[type, str], list[tuple[DependencyChain, Any, Any]]] = defaultdict(list) chains_to_remove: list[DependencyChain] = [] for async_chain in ready_async: @@ -341,7 +357,7 @@ async def _enqueue_batches( resolved_arguments = udf.resolve_arguments(context, call_executor) batchable_arguments = udf.get_batchable_arguments(resolved_arguments) routing_key = udf.get_batch_routing_key(batchable_arguments) - batch_chains[(batch_type, routing_key)].append((async_chain, batchable_arguments)) + batch_chains[(batch_type, routing_key)].append((async_chain, resolved_arguments, batchable_arguments)) except Exception as e: if not isinstance(e, NodeFailurePropagationException): error_infos.append(NodeErrorInfo(e, call_executor.node)) @@ -349,12 +365,17 @@ async def _enqueue_batches( context.set_resolved_value(async_chain, Err(None)) new_batch_tasks: dict[asyncio.Task[Sequence[NodeResult]], Sequence[DependencyChain]] = {} + pre_resolved_by_chain: dict[DependencyChain, Any] = {} for _, chains_and_args in batch_chains.items(): if len(chains_and_args) < 2: + # Reuse native async arguments when a batch does not form. + for chain, resolved_arguments, _ in chains_and_args: + if isinstance(chain.executor, CallExecutor) and isinstance(chain.executor._udf, AsyncBatchableUDFBase): + pre_resolved_by_chain[chain] = resolved_arguments continue - chains, args = zip(*chains_and_args) + chains, _resolved_args, args = zip(*chains_and_args) chains_to_remove.extend(chains) batch_udfs = [chain.executor._udf for chain in chains] @@ -375,7 +396,7 @@ async def _run_legacy_batch(s, u, n, a, c, e): ) new_batch_tasks[task] = chains - remaining = [chain for chain in ready_async if chain not in chains_to_remove] + remaining = [(chain, pre_resolved_by_chain.get(chain)) for chain in ready_async if chain not in chains_to_remove] return remaining, new_batch_tasks @@ -449,12 +470,14 @@ async def execute( ) in_progress_batches.update(new_batch_tasks) - for async_chain in remaining_ready_async: + for async_chain, pre_resolved_arguments in remaining_ready_async: # Native async UDF → await on event loop if isinstance(async_chain.executor, CallExecutor) and isinstance( async_chain.executor._udf, (AsyncUDFBase, AsyncBatchableUDFBase) ): - task = asyncio.create_task(_execute_async_udf(semaphore, async_chain, context, error_infos)) + task = asyncio.create_task( + _execute_async_udf(semaphore, async_chain, context, error_infos, pre_resolved_arguments) + ) else: # Legacy sync UDF with execute_async=True → thread pool fallback task = asyncio.create_task( diff --git a/osprey_async_worker/src/osprey/async_worker/lib/external_service.py b/osprey_async_worker/src/osprey/async_worker/lib/external_service.py index 8fb5049f..8411c143 100644 --- a/osprey_async_worker/src/osprey/async_worker/lib/external_service.py +++ b/osprey_async_worker/src/osprey/async_worker/lib/external_service.py @@ -1,7 +1,7 @@ """Async external service utilities for the async worker. Port of osprey.engine.executor.external_service_utils with asyncio instead of gevent. -Uses asyncio.Future instead of gevent.event.AsyncResult for cache entries. +Uses asyncio.Task instead of gevent.event.AsyncResult for single-value cache entries. """ import asyncio @@ -59,8 +59,9 @@ class ExternalServiceAccessor(Generic[KeyT, ValueT]): def __init__(self, service: AsyncExternalService[KeyT, ValueT]): self._service = service - # Key -> Tuple[ Future[ValueT], Expiration datetime ] - self._cache: dict[KeyT, tuple[asyncio.Future[ValueT], datetime | None]] = {} + # Key -> (Future[ValueT], expiration datetime, count-error-once eligibility) + self._cache: dict[KeyT, tuple[asyncio.Future[ValueT], datetime | None, bool]] = {} + self._active_batch_loaders: set[asyncio.Task[None]] = set() def _is_past_cache_expiration(self, cache_expiration: datetime | None) -> bool: """ @@ -77,77 +78,100 @@ def _get_cache_expiration_datetime(self) -> datetime | None: ttl = self._service.cache_ttl() return datetime.now() + ttl if ttl is not None else None + def _make_task(self, key: KeyT) -> asyncio.Task[ValueT]: + task = asyncio.create_task(self._service.get_from_service(key)) + task.add_done_callback(self._consume_future_exception) + return task + + @staticmethod + def _consume_future_exception(future: asyncio.Future[ValueT]) -> None: + if not future.cancelled(): + future.exception() + def _make_future(self) -> asyncio.Future[ValueT]: - """Create a new Future on the running event loop.""" - return asyncio.get_running_loop().create_future() + future: asyncio.Future[ValueT] = asyncio.get_running_loop().create_future() + future.add_done_callback(self._consume_future_exception) + return future + + def _make_batch_loader(self, keys: Sequence[KeyT], futures: Sequence[asyncio.Future[ValueT]]) -> asyncio.Task[None]: + loader = asyncio.create_task(self._load_batch(keys, futures)) + self._active_batch_loaders.add(loader) + loader.add_done_callback(self._active_batch_loaders.discard) + return loader + + async def _load_batch(self, keys: Sequence[KeyT], futures: Sequence[asyncio.Future[ValueT]]) -> None: + try: + results = await self._service.batch_get_from_service(keys) + if len(results) != len(keys): + raise ValueError(f'batch service returned {len(results)} results for {len(keys)} keys') + for future, result in zip(futures, results): + if result.is_ok(): + future.set_result(result.unwrap()) + else: + future.set_exception(cast(BaseException, result.value)) + except asyncio.CancelledError: + for key, future in zip(keys, futures): + cache_entry = self._cache.get(key) + if cache_entry is not None and cache_entry[0] is future: + del self._cache[key] + if not future.done(): + future.cancel() + raise + except Exception as error: + for future in futures: + if not future.done(): + future.set_exception(error) async def get_without_cache(self, key: KeyT) -> ValueT: """ Ignores any cached values and performs a read-through `get` to the external service. The new value is then used to update the cache entry for subsequent `get` calls. """ - future: asyncio.Future[ValueT] = self._make_future() - cache_entry: tuple[asyncio.Future[ValueT], datetime | None] = ( - future, + task = self._make_task(key) + cache_entry: tuple[asyncio.Future[ValueT], datetime | None, bool] = ( + task, self._get_cache_expiration_datetime(), + False, ) self._cache[key] = cache_entry - try: - result = await self._service.get_from_service(key) - future.set_result(result) - except Exception as e: - future.set_exception(e) - - return await future + return await asyncio.shield(task) async def get(self, key: KeyT) -> ValueT: cache_entry = self._cache.get(key) - if cache_entry is not None and not self._is_past_cache_expiration(cache_entry[1]): - # Cache hit — await the existing future (may still be in-flight from another caller) - return await cache_entry[0] - - future: asyncio.Future[ValueT] = self._make_future() - cache_entry = (future, self._get_cache_expiration_datetime()) - self._cache[key] = cache_entry + is_creator = cache_entry is None or self._is_past_cache_expiration(cache_entry[1]) + if is_creator: + task = self._make_task(key) + cache_entry = (task, self._get_cache_expiration_datetime(), True) + self._cache[key] = cache_entry + assert cache_entry is not None try: - result = await self._service.get_from_service(key) - future.set_result(result) - except Exception as e: - if self._service.count_error_once(): - future.set_result(cast(ValueT, None)) - else: - future.set_exception(e) - future.exception() + return await asyncio.shield(cache_entry[0]) + except Exception: + if self._service.count_error_once() and cache_entry[2] and not is_creator: + return cast(ValueT, None) raise - return await future - async def batch_get(self, keys: Sequence[KeyT]) -> Sequence[Result[ValueT, Exception]]: - cached_entries = [self._cache.get(key) for key in keys] - non_cached_keys = [ - key - for key, cache_entry in zip(keys, cached_entries) - if cache_entry is None or self._is_past_cache_expiration(cache_entry[1]) - ] + non_cached_keys = [] + for key in dict.fromkeys(keys): + cache_entry = self._cache.get(key) + if cache_entry is None or self._is_past_cache_expiration(cache_entry[1]): + non_cached_keys.append(key) if non_cached_keys: + futures = [] for key in non_cached_keys: - self._cache[key] = (self._make_future(), self._get_cache_expiration_datetime()) - try: - result = await self._service.batch_get_from_service(non_cached_keys) - for i, key in enumerate(non_cached_keys): - if result[i].is_ok(): - self._cache[key][0].set_result(result[i].unwrap()) - else: - self._cache[key][0].set_exception(cast(BaseException, result[i].value)) - except Exception as e: - for key in non_cached_keys: - self._cache[key][0].set_exception(e) + future = self._make_future() + self._cache[key] = (future, self._get_cache_expiration_datetime(), False) + futures.append(future) + futures_by_key = {key: self._cache[key][0] for key in keys} + if non_cached_keys: + loader = self._make_batch_loader(non_cached_keys, futures) + await asyncio.shield(loader) results: list[Result[ValueT, Exception]] = [] for key in keys: - future = self._cache[key][0] try: - value = await future + value = await asyncio.shield(futures_by_key[key]) results.append(Ok(value)) except Exception as e: results.append(Err(e)) diff --git a/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py b/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py index 0bc1d6bc..7cf1327c 100644 --- a/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py +++ b/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py @@ -4,7 +4,187 @@ executor for stdlib UDFs (pure computation, no I/O). """ +import asyncio +from collections import defaultdict +from dataclasses import dataclass +from datetime import datetime, timezone +from textwrap import dedent +from typing import ClassVar, Sequence + import pytest +from osprey.async_worker.adaptor.interfaces import AsyncBatchableUDFBase, AsyncUDFBase +from osprey.async_worker.executor import execute +from osprey.engine.ast.grammar import Source +from osprey.engine.ast.sources import Sources +from osprey.engine.ast_validator import validate_sources +from osprey.engine.ast_validator.validator_registry import ValidatorRegistry +from osprey.engine.executor.execution_context import Action, ExecutionContext +from osprey.engine.executor.execution_graph import ExecutionGraph, compile_execution_graph +from osprey.engine.executor.execution_plan import ExecutionPlanState +from osprey.engine.executor.udf_execution_helpers import UDFHelpers +from osprey.engine.stdlib import get_config_registry +from osprey.engine.udf.arguments import ArgumentsBase +from osprey.engine.udf.registry import UDFRegistry +from result import Ok + + +class CountingBatchableArguments(ArgumentsBase): + key: str + value: str + + +@dataclass +class CountingBatchableArgs: + key: str + value: str + + +class CountingBatchableUdf(AsyncBatchableUDFBase[CountingBatchableArguments, str, CountingBatchableArgs]): + """A batchable async UDF that records argument resolution calls. + + The explicit type getters avoid unsubstituted TypeVars for this test class. + """ + + resolve_call_count: ClassVar[int] = 0 + raise_on_resolve: ClassVar[bool] = False + + @classmethod + def get_arguments_type(cls): + return CountingBatchableArguments + + @classmethod + def get_rvalue_type(cls): + return str + + @classmethod + def get_batchable_arguments_type(cls): + return CountingBatchableArgs + + def resolve_arguments(self, execution_context, call_executor) -> CountingBatchableArguments: + type(self).resolve_call_count += 1 + if type(self).raise_on_resolve: + raise RuntimeError('resolve boom') + return super().resolve_arguments(execution_context, call_executor) + + def get_batchable_arguments(self, arguments: CountingBatchableArguments) -> CountingBatchableArgs: + return CountingBatchableArgs(key=arguments.key, value=arguments.value) + + def get_batch_routing_key(self, arguments: CountingBatchableArgs) -> str: + return arguments.key + + async def async_execute(self, execution_context: ExecutionContext, arguments: CountingBatchableArguments) -> str: + return arguments.value + + async def async_execute_batch(self, execution_context, udfs, arguments: Sequence[CountingBatchableArgs]): + return [Ok(a.value) for a in arguments] + + +class GatedArguments(ArgumentsBase): + value: str + + +class GatedAsyncUdf(AsyncUDFBase[GatedArguments, str]): + entered = 0 + both_entered: asyncio.Event + release: asyncio.Event + + async def async_execute(self, execution_context: ExecutionContext, arguments: GatedArguments) -> str: + type(self).entered += 1 + if type(self).entered == 2: + type(self).both_entered.set() + await type(self).release.wait() + return arguments.value + + +@pytest.mark.asyncio +async def test_planned_execute_activates_each_imported_or_required_source_once( + async_execute_fn, monkeypatch: pytest.MonkeyPatch +) -> None: + activated_sources: list[str] = [] + activate_source = ExecutionPlanState.activate_source + + def record_activation(self: ExecutionPlanState, source: Source) -> None: + activated_sources.append(source.path) + activate_source(self, source) + + monkeypatch.setattr(ExecutionPlanState, 'activate_source', record_activation) + + result = await async_execute_fn( + { + 'main.sml': "Import(rules=['branch.sml', 'shared.sml'])", + 'branch.sml': "Require(rule='shared.sml')", + 'shared.sml': 'Shared = 1 + 0', + } + ) + + assert result == {'Shared': 1} + assert activated_sources == ['main.sml', 'branch.sml', 'shared.sml'] + + +@pytest.mark.asyncio +async def test_concurrent_actions_isolate_dynamic_source_activation( + stdlib_udf_registry: UDFRegistry, monkeypatch: pytest.MonkeyPatch +) -> None: + registry = stdlib_udf_registry + registry.register(GatedAsyncUdf) + sources = Sources.from_dict( + { + 'main.sml': dedent( + """ + ActionName: str = JsonData(path="$.action_name", coerce_type=True) + Require(rule=f"actions/{ActionName}.sml") + """ + ), + 'actions/a.sml': 'A = GatedAsyncUdf(value="a")', + 'actions/b.sml': 'B = GatedAsyncUdf(value="b")', + } + ) + validator_registry = ValidatorRegistry.get_instance().instance_with_additional_validators( + get_config_registry().get_validator() + ) + graph = compile_execution_graph(validate_sources(sources, registry, validator_registry)) + sources_by_action: dict[str, list[str]] = defaultdict(list) + original_enqueue_source = ExecutionContext.enqueue_source + + def record_enqueue_source(context: ExecutionContext, source: Source) -> None: + sources_by_action[context.get_action_name()].append(source.path) + original_enqueue_source(context, source) + + monkeypatch.setattr(ExecutionContext, 'enqueue_source', record_enqueue_source) + GatedAsyncUdf.entered = 0 + GatedAsyncUdf.both_entered = asyncio.Event() + GatedAsyncUdf.release = asyncio.Event() + timestamp = datetime(2026, 8, 4, tzinfo=timezone.utc) + tasks = [ + asyncio.create_task( + execute( + graph, + UDFHelpers(), + Action(action_id=index, action_name=name, data={'action_name': name}, timestamp=timestamp), + ) + ) + for index, name in enumerate(('a', 'b'), start=1) + ] + + try: + await asyncio.wait_for(GatedAsyncUdf.both_entered.wait(), timeout=5) + assert not any(task.done() for task in tasks) + GatedAsyncUdf.release.set() + result_a, result_b = await asyncio.gather(*tasks) + finally: + GatedAsyncUdf.release.set() + await asyncio.gather(*tasks, return_exceptions=True) + + assert result_a.extracted_features['A'] == 'a' + assert 'B' not in result_a.extracted_features + assert result_b.extracted_features['B'] == 'b' + assert 'A' not in result_b.extracted_features + assert not result_a.error_infos + assert not result_b.error_infos + assert sources_by_action == { + 'a': ['main.sml', 'actions/a.sml'], + 'b': ['main.sml', 'actions/b.sml'], + } @pytest.mark.asyncio @@ -188,3 +368,108 @@ async def test_parity_complex_graph(async_execute_fn): assert result['BUpper'] == 'HI' assert result['RuleA'] is True assert result['RuleB'] is False + + +@pytest.fixture() +def counting_batchable_udf(): + """Registers CountingBatchableUdf and resets its call-count/raise state around the test.""" + CountingBatchableUdf.resolve_call_count = 0 + CountingBatchableUdf.raise_on_resolve = False + yield CountingBatchableUdf + CountingBatchableUdf.resolve_call_count = 0 + CountingBatchableUdf.raise_on_resolve = False + + +@pytest.mark.asyncio +async def test_batch_of_one_reuses_resolved_arguments( + async_execute_with_result, stdlib_udf_registry: UDFRegistry, counting_batchable_udf +): + """A batchable UDF alone (batch group size 1) falls through to the singleton async path. + + Before the fix, this chain's arguments were resolved once to compute the routing key in + `_enqueue_batches`, then resolved again in `_execute_async_udf` -- a duplicate Arguments + construction. resolve_arguments should now only be called once. + """ + stdlib_udf_registry.register(counting_batchable_udf) + result = await async_execute_with_result( + 'A = CountingBatchableUdf(key="solo", value="a")', udf_registry=stdlib_udf_registry + ) + assert result.extracted_features['A'] == 'a' + assert not result.error_infos + assert counting_batchable_udf.resolve_call_count == 1 + + +@pytest.mark.asyncio +async def test_batch_of_two_resolves_once_per_chain( + async_execute_with_result, stdlib_udf_registry: UDFRegistry, counting_batchable_udf +): + """When a batch group forms (>=2 chains sharing a routing key), each chain's arguments are + resolved exactly once -- the batch execution path never re-resolves, so this is unchanged + by the fallthrough fix.""" + stdlib_udf_registry.register(counting_batchable_udf) + result = await async_execute_with_result( + """ + A = CountingBatchableUdf(key="shared", value="a") + B = CountingBatchableUdf(key="shared", value="b") + """, + udf_registry=stdlib_udf_registry, + ) + assert result.extracted_features['A'] == 'a' + assert result.extracted_features['B'] == 'b' + assert not result.error_infos + assert counting_batchable_udf.resolve_call_count == 2 + + +@pytest.mark.asyncio +async def test_plan_matches_legacy_native_async_batch( + async_execute_with_result, + stdlib_udf_registry: UDFRegistry, + counting_batchable_udf, + monkeypatch: pytest.MonkeyPatch, +): + stdlib_udf_registry.register(counting_batchable_udf) + sources = """ + A = CountingBatchableUdf(key="shared", value="a") + B = CountingBatchableUdf(key="shared", value="b") + """ + action_time = datetime(2026, 8, 4, tzinfo=timezone.utc) + + with monkeypatch.context() as legacy: + legacy.setattr(ExecutionGraph, 'get_execution_plan', lambda _graph: None) + legacy_result = await async_execute_with_result( + sources, + udf_registry=stdlib_udf_registry, + action_time=action_time, + ) + legacy_resolve_count = counting_batchable_udf.resolve_call_count + counting_batchable_udf.resolve_call_count = 0 + + planned_result = await async_execute_with_result( + sources, + udf_registry=stdlib_udf_registry, + action_time=action_time, + ) + + assert planned_result.extracted_features == legacy_result.extracted_features + assert planned_result.error_infos == legacy_result.error_infos == [] + assert counting_batchable_udf.resolve_call_count == legacy_resolve_count == 2 + + +@pytest.mark.asyncio +async def test_batch_of_one_resolve_failure_surfaces_once( + async_execute_with_result, stdlib_udf_registry: UDFRegistry, counting_batchable_udf +): + """If resolve_arguments raises while computing the routing key (batch group size 1), the + failure is fully handled inside _enqueue_batches: the chain is resolved to Err(None) there + and never falls through to _execute_async_udf. resolve_arguments is therefore still only + attempted once, and the same exception surfaces as before the fallthrough fix.""" + counting_batchable_udf.raise_on_resolve = True + stdlib_udf_registry.register(counting_batchable_udf) + result = await async_execute_with_result( + 'A = CountingBatchableUdf(key="solo", value="a")', udf_registry=stdlib_udf_registry + ) + assert result.extracted_features.get('A') is None + assert len(result.error_infos) == 1 + assert isinstance(result.error_infos[0].error, RuntimeError) + assert str(result.error_infos[0].error) == 'resolve boom' + assert counting_batchable_udf.resolve_call_count == 1 diff --git a/osprey_async_worker/src/osprey/async_worker/tests/test_external_service.py b/osprey_async_worker/src/osprey/async_worker/tests/test_external_service.py index 5c553f1a..789c2406 100644 --- a/osprey_async_worker/src/osprey/async_worker/tests/test_external_service.py +++ b/osprey_async_worker/src/osprey/async_worker/tests/test_external_service.py @@ -1,17 +1,16 @@ """Tests for the async external service cache.""" import asyncio +import gc from collections.abc import Sequence from datetime import timedelta import pytest from osprey.async_worker.lib.external_service import AsyncExternalService, ExternalServiceAccessor -from result import Ok, Result +from result import Err, Ok, Result class FakeService(AsyncExternalService[str, str]): - """Test service that records calls and returns predictable results.""" - def __init__(self, delay: float = 0.0): self.call_count = 0 self.delay = delay @@ -24,25 +23,61 @@ async def get_from_service(self, key: str) -> str: class FailingService(AsyncExternalService[str, str]): - """Test service that always raises.""" + def __init__(self): + self.call_count = 0 async def get_from_service(self, key: str) -> str: + self.call_count += 1 raise ValueError(f'service error for {key}') -class FailOnceService(AsyncExternalService[str, str | None]): - """Raises on first call per key, succeeds after.""" - +class CountErrorOnceGatedService(AsyncExternalService[str, str | None]): def __init__(self): - self.seen = set() + self.call_count = 0 + self.started = asyncio.Event() + self.release = asyncio.Event() def count_error_once(self) -> bool: return True async def get_from_service(self, key: str) -> str | None: - if key not in self.seen: - self.seen.add(key) + self.call_count += 1 + self.started.set() + await self.release.wait() + raise ValueError('service fails') + + +class CancelOnceService(AsyncExternalService[str, str]): + def __init__(self): + self.call_count = 0 + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def get_from_service(self, key: str) -> str: + self.call_count += 1 + if self.call_count == 1: + self.started.set() + await self.release.wait() + return f'value_{key}' + + +class ReplacementService(AsyncExternalService[str, str]): + def __init__(self): + self.call_count = 0 + self.first_started = asyncio.Event() + self.first_release = asyncio.Event() + self.second_started = asyncio.Event() + self.second_release = asyncio.Event() + + async def get_from_service(self, key: str) -> str: + self.call_count += 1 + if self.call_count == 1: + self.first_started.set() + await self.first_release.wait() raise ValueError('first call fails') + if self.call_count == 2: + self.second_started.set() + await self.second_release.wait() return f'value_{key}' @@ -60,28 +95,81 @@ async def get_from_service(self, key: str) -> str: class BatchService(AsyncExternalService[str, str]): - """Test service that supports batch operations.""" + def __init__(self): + self.batch_call_count = 0 + + async def get_from_service(self, key: str) -> str: + return f'value_{key}' + + async def batch_get_from_service(self, keys: Sequence[str]) -> Sequence[Result[str, Exception]]: + self.batch_call_count += 1 + return [Ok(f'batch_{key}') for key in keys] + + +class FailOnceBatchService(AsyncExternalService[str, str]): + def __init__(self, raise_exception: bool): + self.raise_exception = raise_exception + self.batch_call_count = 0 + + async def get_from_service(self, key: str) -> str: + return f'value_{key}' + + async def batch_get_from_service(self, keys: Sequence[str]) -> Sequence[Result[str, Exception]]: + self.batch_call_count += 1 + if self.batch_call_count == 1: + if self.raise_exception: + raise ValueError('batch fails') + return [Err(ValueError('item fails')) for _ in keys] + return [Ok(f'batch_{key}') for key in keys] + + +class CountErrorOnceBatchService(AsyncExternalService[str, str | None]): + def count_error_once(self) -> bool: + return True + + async def get_from_service(self, key: str) -> str | None: + raise ValueError('single read fails') + + async def batch_get_from_service(self, keys: Sequence[str]) -> Sequence[Result[str | None, Exception]]: + return [Err(ValueError('batch read fails')) for _ in keys] + +class GatedBatchService(AsyncExternalService[str, str]): def __init__(self): + self.get_call_count = 0 self.batch_call_count = 0 + self.started = asyncio.Event() + self.release = asyncio.Event() async def get_from_service(self, key: str) -> str: + self.get_call_count += 1 return f'value_{key}' async def batch_get_from_service(self, keys: Sequence[str]) -> Sequence[Result[str, Exception]]: self.batch_call_count += 1 + self.started.set() + await self.release.wait() return [Ok(f'batch_{key}') for key in keys] -# --- Cache tests --- +class FailingGatedBatchService(AsyncExternalService[str, str]): + def __init__(self): + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def get_from_service(self, key: str) -> str: + raise ValueError('single read fails') + + async def batch_get_from_service(self, keys: Sequence[str]) -> Sequence[Result[str, Exception]]: + self.started.set() + await self.release.wait() + raise ValueError('batch fails') @pytest.mark.asyncio async def test_get_returns_value(): - service = FakeService() - accessor = ExternalServiceAccessor(service) - result = await accessor.get('foo') - assert result == 'value_foo' + accessor = ExternalServiceAccessor(FakeService()) + assert await accessor.get('foo') == 'value_foo' @pytest.mark.asyncio @@ -103,41 +191,71 @@ async def test_get_different_keys_not_cached(): @pytest.mark.asyncio -async def test_get_without_cache_bypasses(): +async def test_get_without_cache_updates_cache(): service = FakeService() accessor = ExternalServiceAccessor(service) - await accessor.get('foo') await accessor.get_without_cache('foo') - assert service.call_count == 2 + await accessor.get('foo') + assert service.call_count == 1 @pytest.mark.asyncio -async def test_get_without_cache_updates_cache(): - service = FakeService() +async def test_cancelling_get_without_cache_does_not_cancel_shared_get(): + service = CancelOnceService() accessor = ExternalServiceAccessor(service) - await accessor.get_without_cache('foo') - await accessor.get('foo') - assert service.call_count == 1 # Second get hits cache + owner = asyncio.create_task(accessor.get_without_cache('foo')) + await service.started.wait() + owner.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await owner + service.release.set() -# --- Concurrent access (future dedup) --- + assert await accessor.get('foo') == 'value_foo' + assert service.call_count == 1 @pytest.mark.asyncio async def test_concurrent_get_deduplicates(): - """Multiple concurrent gets for the same key should only call service once.""" service = FakeService(delay=0.05) accessor = ExternalServiceAccessor(service) - results = await asyncio.gather( - accessor.get('foo'), - accessor.get('foo'), - accessor.get('foo'), - ) - assert all(r == 'value_foo' for r in results) + results = await asyncio.gather(accessor.get('foo'), accessor.get('foo'), accessor.get('foo')) + assert results == ['value_foo', 'value_foo', 'value_foo'] assert service.call_count == 1 -# --- Error handling --- +@pytest.mark.asyncio +async def test_cancelling_waiter_does_not_cancel_shared_get(): + service = CancelOnceService() + accessor = ExternalServiceAccessor(service) + owner = asyncio.create_task(accessor.get('foo')) + await service.started.wait() + waiter = asyncio.create_task(accessor.get('foo')) + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await waiter + + service.release.set() + assert await owner == 'value_foo' + assert service.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancelling_owner_does_not_cancel_shared_get(): + service = CancelOnceService() + accessor = ExternalServiceAccessor(service) + owner = asyncio.create_task(accessor.get('foo')) + await service.started.wait() + survivor = asyncio.create_task(accessor.get('foo')) + + owner.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await owner + + service.release.set() + assert await survivor == 'value_foo' + assert service.call_count == 1 @pytest.mark.asyncio @@ -146,42 +264,78 @@ async def test_get_propagates_error(): accessor = ExternalServiceAccessor(service) with pytest.raises(ValueError, match='service error for foo'): await accessor.get('foo') + assert service.call_count == 1 @pytest.mark.asyncio -async def test_get_error_cached(): - """Errors are cached — second get raises the same error.""" +async def test_get_caches_failed_request(): service = FailingService() accessor = ExternalServiceAccessor(service) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='service error for foo'): await accessor.get('foo') - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='service error for foo'): await accessor.get('foo') + assert service.call_count == 1 @pytest.mark.asyncio -async def test_count_error_once(): - """With count_error_once, subsequent callers get None instead of the error.""" - service = FailOnceService() +async def test_failed_task_does_not_evict_replacement(): + service = ReplacementService() accessor = ExternalServiceAccessor(service) + first = asyncio.create_task(accessor.get('foo')) + await service.first_started.wait() + replacement = asyncio.create_task(accessor.get_without_cache('foo')) + await service.second_started.wait() + + service.first_release.set() with pytest.raises(ValueError): - await accessor.get('foo') - # Second get should return None (cached as None due to count_error_once) - result = await accessor.get('foo') - assert result is None + _ = await first + service.second_release.set() + assert await replacement == 'value_foo' + assert await accessor.get('foo') == 'value_foo' + assert service.call_count == 2 -# --- TTL --- + +@pytest.mark.asyncio +async def test_count_error_once_with_concurrent_waiter(): + service = CountErrorOnceGatedService() + accessor = ExternalServiceAccessor(service) + creator = asyncio.create_task(accessor.get('foo')) + await service.started.wait() + waiter = asyncio.create_task(accessor.get('foo')) + await asyncio.sleep(0) + service.release.set() + + with pytest.raises(ValueError): + _ = await creator + assert await waiter is None + assert await accessor.get('foo') is None + assert service.call_count == 1 + + +@pytest.mark.asyncio +async def test_count_error_once_does_not_apply_to_get_without_cache(): + service = CountErrorOnceGatedService() + accessor = ExternalServiceAccessor(service) + creator = asyncio.create_task(accessor.get_without_cache('foo')) + await service.started.wait() + service.release.set() + + with pytest.raises(ValueError, match='service fails'): + _ = await creator + with pytest.raises(ValueError, match='service fails'): + _ = await accessor.get('foo') + assert service.call_count == 1 @pytest.mark.asyncio async def test_ttl_expires_cache(): - """Expired TTL causes a re-fetch.""" - service = TTLService(ttl=timedelta(days=-1)) # Immediately expired + service = TTLService(ttl=timedelta(days=-1)) accessor = ExternalServiceAccessor(service) - r1 = await accessor.get('foo') - r2 = await accessor.get('foo') - assert r1 != r2 # Different values = two service calls + first = await accessor.get('foo') + second = await accessor.get('foo') + assert first != second assert service.call_count == 2 @@ -195,27 +349,138 @@ async def test_no_ttl_caches_forever(): assert service.call_count == 1 -# --- Batch --- +@pytest.mark.asyncio +async def test_batch_get_returns_values_and_uses_cache(): + service = BatchService() + accessor = ExternalServiceAccessor(service) + assert await accessor.batch_get(['a', 'b', 'c']) == [Ok('batch_a'), Ok('batch_b'), Ok('batch_c')] + assert await accessor.batch_get(['a', 'b', 'c']) == [Ok('batch_a'), Ok('batch_b'), Ok('batch_c')] + assert service.batch_call_count == 1 @pytest.mark.asyncio -async def test_batch_get(): +async def test_batch_get_deduplicates_duplicate_keys(): service = BatchService() accessor = ExternalServiceAccessor(service) - results = await accessor.batch_get(['a', 'b', 'c']) - assert len(results) == 3 - assert results[0] == Ok('batch_a') - assert results[1] == Ok('batch_b') - assert results[2] == Ok('batch_c') + assert await accessor.batch_get(['a', 'a']) == [Ok('batch_a'), Ok('batch_a')] assert service.batch_call_count == 1 @pytest.mark.asyncio -async def test_batch_get_uses_cache(): - service = BatchService() +async def test_cancelled_batch_loader_evicts_its_cache_entries(): + service = GatedBatchService() accessor = ExternalServiceAccessor(service) - await accessor.batch_get(['a', 'b']) - # Second batch with overlap — 'a' and 'b' cached, only 'c' fetched - results = await accessor.batch_get(['a', 'b', 'c']) - assert len(results) == 3 + batch = asyncio.create_task(accessor.batch_get(['a'])) + await service.started.wait() + + loader = next(iter(accessor._active_batch_loaders)) + loader.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await batch + service.release.set() + + assert await accessor.batch_get(['a']) == [Ok('batch_a')] assert service.batch_call_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize('raise_exception', [False, True]) +async def test_batch_error_is_cached(raise_exception: bool): + service = FailOnceBatchService(raise_exception) + accessor = ExternalServiceAccessor(service) + first = await accessor.batch_get(['a']) + second = await accessor.batch_get(['a']) + assert first[0].is_err() + assert second[0].is_err() + assert service.batch_call_count == 1 + + +@pytest.mark.asyncio +async def test_count_error_once_does_not_apply_to_batch_failure(): + accessor = ExternalServiceAccessor(CountErrorOnceBatchService()) + batch_result = await accessor.batch_get(['a']) + + assert batch_result[0].is_err() + with pytest.raises(ValueError, match='batch read fails'): + await accessor.get('a') + + +@pytest.mark.asyncio +async def test_cancelling_batch_owner_does_not_cancel_shared_get(): + service = GatedBatchService() + accessor = ExternalServiceAccessor(service) + batch = asyncio.create_task(accessor.batch_get(['a'])) + await service.started.wait() + survivor = asyncio.create_task(accessor.get('a')) + + batch.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await batch + service.release.set() + + assert await survivor == 'batch_a' + assert service.batch_call_count == 1 + assert service.get_call_count == 0 + + +@pytest.mark.asyncio +async def test_cancelled_batch_owner_keeps_loader_alive_through_garbage_collection(): + service = GatedBatchService() + accessor = ExternalServiceAccessor(service) + batch = asyncio.create_task(accessor.batch_get(['a'])) + await service.started.wait() + survivor = asyncio.create_task(accessor.get('a')) + + batch.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await batch + gc.collect() + service.release.set() + + assert await survivor == 'batch_a' + assert service.batch_call_count == 1 + + +@pytest.mark.asyncio +async def test_cancelled_failed_batch_consumes_future_exceptions(): + service = FailingGatedBatchService() + accessor = ExternalServiceAccessor(service) + loop = asyncio.get_running_loop() + contexts: list[dict[str, object]] = [] + previous_handler = loop.get_exception_handler() + loop.set_exception_handler(lambda _loop, context: contexts.append(context)) + try: + batch = asyncio.create_task(accessor.batch_get(['a'])) + await service.started.wait() + + batch.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await batch + service.release.set() + await asyncio.sleep(0) + await asyncio.sleep(0) + accessor._cache.clear() + gc.collect() + await asyncio.sleep(0) + finally: + loop.set_exception_handler(previous_handler) + + assert not any(context.get('message') == 'Future exception was never retrieved' for context in contexts) + + +@pytest.mark.asyncio +async def test_cancelling_batch_waiter_does_not_cancel_shared_get(): + service = CancelOnceService() + accessor = ExternalServiceAccessor(service) + owner = asyncio.create_task(accessor.get('a')) + await service.started.wait() + batch = asyncio.create_task(accessor.batch_get(['a'])) + await asyncio.sleep(0) + + batch.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await batch + service.release.set() + + assert await owner == 'value_a' + assert service.call_count == 1 diff --git a/osprey_coordinator/Cargo.toml b/osprey_coordinator/Cargo.toml index 33eb7395..162f56a8 100644 --- a/osprey_coordinator/Cargo.toml +++ b/osprey_coordinator/Cargo.toml @@ -85,6 +85,7 @@ prost-build = { version = "0.12" } tonic-build = "0.11" [dev-dependencies] +tokio = { version = "1.4", features = ["test-util"] } [[bin]] name = "osprey_coordinator" diff --git a/osprey_coordinator/src/pub_sub_streaming_pull/message_ack_queue.rs b/osprey_coordinator/src/pub_sub_streaming_pull/message_ack_queue.rs index a5a72b67..35bf8770 100644 --- a/osprey_coordinator/src/pub_sub_streaming_pull/message_ack_queue.rs +++ b/osprey_coordinator/src/pub_sub_streaming_pull/message_ack_queue.rs @@ -269,3 +269,96 @@ where self.next_flush = Instant::now() + self.pending_ack_flush_interval; } } + +#[cfg(test)] +mod tests { + use std::{collections::HashSet, time::Instant as StdInstant}; + + use tokio::time::{advance, Instant as TokioInstant}; + + use super::*; + + const LEASE_DELAY: Duration = Duration::from_secs(30); + + fn queue(capacity: usize, chunk_size: usize) -> MessageAckQueue { + MessageAckQueue::new(capacity, chunk_size, Duration::from_millis(100)) + } + + fn insert_due(queue: &mut MessageAckQueue, count: usize) -> Vec { + let renew_at = TokioInstant::now(); + (0..count) + .map(|index| { + queue.transform_and_store_ack_id( + format!("server-{index}"), + StdInstant::now(), + renew_at, + ) + }) + .collect() + } + + fn renew_one_tick(queue: &mut MessageAckQueue, chunk_size: usize) -> Vec { + queue + .collect_ack_ids_that_need_to_be_renewed(chunk_size, LEASE_DELAY) + .into_iter() + .map(|(_, ack_id)| ack_id) + .collect() + } + + #[tokio::test(start_paused = true)] + async fn renewal_tick_caps_work_and_preserves_remaining_messages() { + let mut queue = queue(8_000, 2_500); + insert_due(&mut queue, 5_001); + advance(Duration::from_secs(1)).await; + + let first_batch = renew_one_tick(&mut queue, 2_500); + + assert_eq!(first_batch.len(), 2_500); + advance(Duration::from_secs(1)).await; + assert_eq!(renew_one_tick(&mut queue, 2_500).len(), 2_500); + advance(Duration::from_secs(1)).await; + assert_eq!(renew_one_tick(&mut queue, 2_500).len(), 1); + } + + #[tokio::test(start_paused = true)] + async fn renewal_ticks_cover_due_messages_once_without_starvation() { + let mut queue = queue(8_000, 2_500); + insert_due(&mut queue, 5_001); + advance(Duration::from_secs(1)).await; + + let mut renewed = renew_one_tick(&mut queue, 2_500); + advance(Duration::from_secs(1)).await; + renewed.extend(renew_one_tick(&mut queue, 2_500)); + advance(Duration::from_secs(1)).await; + renewed.extend(renew_one_tick(&mut queue, 2_500)); + let unique: HashSet<_> = renewed.iter().collect(); + + assert_eq!(renewed.len(), 5_001); + assert_eq!(unique.len(), 5_001); + assert!(renew_one_tick(&mut queue, 2_500).is_empty()); + } + + #[tokio::test(start_paused = true)] + async fn acknowledged_messages_are_excluded_from_renewal() { + let mut queue = queue(10, 10); + let ack_ids = insert_due(&mut queue, 5); + queue.remove_ack_id(&ack_ids[2]); + advance(Duration::from_secs(1)).await; + + let renewed = renew_one_tick(&mut queue, 10); + + assert_eq!(renewed.len(), 4); + assert!(!renewed.iter().any(|ack_id| ack_id == "server-2")); + } + + #[tokio::test(start_paused = true)] + async fn renewal_recurs_after_each_lease_window() { + let mut queue = queue(10, 10); + insert_due(&mut queue, 3); + advance(Duration::from_secs(1)).await; + + assert_eq!(renew_one_tick(&mut queue, 10).len(), 3); + advance(Duration::from_secs(31)).await; + assert_eq!(renew_one_tick(&mut queue, 10).len(), 3); + } +} diff --git a/osprey_coordinator/src/pub_sub_streaming_pull/streaming_pull_manager.rs b/osprey_coordinator/src/pub_sub_streaming_pull/streaming_pull_manager.rs index 01d6ff2e..1e49b187 100644 --- a/osprey_coordinator/src/pub_sub_streaming_pull/streaming_pull_manager.rs +++ b/osprey_coordinator/src/pub_sub_streaming_pull/streaming_pull_manager.rs @@ -49,6 +49,51 @@ const METRICS_REPORTING_INTERVAL: Duration = Duration::from_secs(1); /// How often we should attempt to flush leases. const LEASE_RENEWAL_INTERVAL: Duration = Duration::from_secs(1); +/// The first lease renewal waits for messages to form a batch. +const INITIAL_LEASE_RENEWAL_DELAY: Duration = Duration::from_secs(1); + +/// The shortest renewal interval that yields to the manager loop. +const MIN_LEASE_RENEWAL_INTERVAL: Duration = Duration::from_millis(1); + +fn lease_renewal_interval( + flow_control: &FlowControl, + lease_renewal_deadline: Duration, +) -> Duration { + assert!( + flow_control.max_messages > 0, + "Flow control max_messages must be greater than zero" + ); + + let catch_up_window = lease_renewal_deadline + .checked_sub(INITIAL_LEASE_RENEWAL_DELAY) + .expect("Buffered lease deadline must exceed the initial renewal delay"); + let batch_count = flow_control.max_messages.div_ceil(ACK_IDS_MAX_BATCH_SIZE); + let batch_count_nanos = u128::try_from(batch_count).expect("Batch count must fit in u128"); + let minimum_catch_up_nanos = MIN_LEASE_RENEWAL_INTERVAL + .as_nanos() + .checked_mul(batch_count_nanos) + .expect("Lease renewal batch count exceeds the supported range"); + + assert!( + minimum_catch_up_nanos <= catch_up_window.as_nanos(), + "Flow control max_messages cannot renew all lease batches before the buffered deadline" + ); + + let interval_nanos = + (catch_up_window.as_nanos() / batch_count_nanos).min(LEASE_RENEWAL_INTERVAL.as_nanos()); + Duration::from_nanos( + u64::try_from(interval_nanos).expect("Lease renewal interval must fit in u64"), + ) +} + +fn buffered_lease_renewal_duration(lease_renewal_duration_secs: u32) -> Duration { + let buffer = Duration::from_millis( + ((((lease_renewal_duration_secs * 1000) as f64) * 0.2) as u64).min(5_000), + ); + + Duration::from_secs(lease_renewal_duration_secs as _) - buffer +} + type TonicSubscriberClient = SubscriberClient>; impl FlowControl { @@ -282,12 +327,7 @@ where /// Returns the maximum amount of time we should wait before renewing the lease on a message. fn get_buffered_lease_renewal_duration(&self) -> Duration { - // Allow for 20% buffer, or 5 seconds, which ever is greater. - let buffer = Duration::from_millis( - ((((self.message_lease_renewal_duration_secs * 1000) as f64) * 0.2) as u64).min(5_000), - ); - - Duration::from_secs(self.message_lease_renewal_duration_secs as _) - buffer + buffered_lease_renewal_duration(self.message_lease_renewal_duration_secs) } /// Recomputes `message_lease_renewal_duration_secs` based upon the p99 of message processing latency. @@ -346,6 +386,11 @@ where message_handler: M, metrics_client_builder: MetricsClientBuilder, ) -> StreamingPullManagerHandle { + let initial_lease_renewal_duration_secs = flow_control.min_duration_per_lease_extension; + lease_renewal_interval( + &flow_control, + buffered_lease_renewal_duration(initial_lease_renewal_duration_secs), + ); let client_id = uuid::Uuid::new_v4(); let metrics = StreamingPullManagerMetrics::new(); @@ -361,7 +406,7 @@ where (flow_control.max_duration_per_lease_extension * 1000) as _, 3 ).expect("invariant: histogram creation should never fail, as the bounds are checked in FlowControl"); - let message_lease_renewal_duration_secs = flow_control.min_duration_per_lease_extension; + let message_lease_renewal_duration_secs = initial_lease_renewal_duration_secs; let message_ack_queue = MessageAckQueue::new( flow_control.max_messages, std::cmp::min(ACK_IDS_MAX_BATCH_SIZE, flow_control.max_messages), @@ -427,7 +472,10 @@ where let mut metrics_reporting_interval = interval(METRICS_REPORTING_INTERVAL); metrics_reporting_interval.set_missed_tick_behavior(MissedTickBehavior::Skip); - let mut lease_renewal_interval = interval(LEASE_RENEWAL_INTERVAL); + let mut lease_renewal_interval = interval(lease_renewal_interval( + &self.state.flow_control, + self.state.get_buffered_lease_renewal_duration(), + )); lease_renewal_interval.set_missed_tick_behavior(MissedTickBehavior::Skip); while !self.should_stop() { @@ -598,7 +646,7 @@ where let received_at = Instant::now(); // The rationale of this being 1 second from now is that we'll delay these by 1 second so that we can build // up a batch of lease renewals within the message ack queue's lease delay queue. - let lease_renew_at = tokio::time::Instant::from(received_at + Duration::from_secs(1)); + let lease_renew_at = tokio::time::Instant::from(received_at + INITIAL_LEASE_RENEWAL_DELAY); streaming_pull_response .received_messages @@ -773,66 +821,56 @@ where && self.state.background_tasks.is_empty() } - /// Issues modack requests to pub-sub in order to renew leases for messages that are either in-flight or on hold. - /// - /// Returns true if a full batch was flushed, and the next lease renewal flush should perhaps be expedited. + /// Renews one lease batch per tick so that the manager can receive more messages. fn handle_renew_leases(&mut self) { - loop { - self.state.recompute_message_lease_renewal_duration_secs(); - let lease_renewal_duration = - Duration::from_secs(self.state.message_lease_renewal_duration_secs as _); - - let chunk = self - .state - .message_ack_queue - .collect_ack_ids_that_need_to_be_renewed( - ACK_IDS_MAX_BATCH_SIZE, - self.state.get_buffered_lease_renewal_duration(), - ); - - if chunk.is_empty() { - break; - } - - let chunk_is_full = chunk.len() == ACK_IDS_MAX_BATCH_SIZE; + self.state.recompute_message_lease_renewal_duration_secs(); + let lease_renewal_duration = + Duration::from_secs(self.state.message_lease_renewal_duration_secs as _); - tracing::debug!( - {subscription = %self.state.subscription, client_id = %self.state.client_id}, - "renewing leases for {} messages with a lease duration of {} seconds", - chunk.len(), - self.state.message_lease_renewal_duration_secs, + let chunk = self + .state + .message_ack_queue + .collect_ack_ids_that_need_to_be_renewed( + ACK_IDS_MAX_BATCH_SIZE, + self.state.get_buffered_lease_renewal_duration(), ); - self.state - .metrics - .message_leases_renewed - .incr_by(chunk.len() as _); + if chunk.is_empty() { + return; + } - let mut ack_ids = Vec::with_capacity(chunk.len()); + tracing::debug!( + {subscription = %self.state.subscription, client_id = %self.state.client_id}, + "renewing leases for {} messages with a lease duration of {} seconds", + chunk.len(), + self.state.message_lease_renewal_duration_secs, + ); - for (elapsed, ack_id) in chunk { - if elapsed > lease_renewal_duration { - self.state - .message_latency_histogram - .saturating_record(elapsed.as_millis() as _); - } - ack_ids.push(ack_id); - } + self.state + .metrics + .message_leases_renewed + .incr_by(chunk.len() as _); - self.state.perform_request_in_background_with_retries( - "renew_leases", - make_modack_request( - &self.state.subscription, - ack_ids, - self.state.message_lease_renewal_duration_secs, - ), - |mut c, r| async move { c.modify_ack_deadline(r).await }, - ); + let mut ack_ids = Vec::with_capacity(chunk.len()); - if !chunk_is_full { - break; + for (elapsed, ack_id) in chunk { + if elapsed > lease_renewal_duration { + self.state + .message_latency_histogram + .saturating_record(elapsed.as_millis() as _); } + ack_ids.push(ack_id); } + + self.state.perform_request_in_background_with_retries( + "renew_leases", + make_modack_request( + &self.state.subscription, + ack_ids, + self.state.message_lease_renewal_duration_secs, + ), + |mut c, r| async move { c.modify_ack_deadline(r).await }, + ); } /// Flushes a chunk of acks or nacks to pub-sub server. @@ -1046,3 +1084,166 @@ impl StreamingPullChannel { } } } + +#[cfg(test)] +mod tests { + use super::*; + + struct TestMessageHandler; + + impl MessageHandler for TestMessageHandler { + fn handle_messages(&mut self, _: Vec) {} + } + + #[derive(Clone)] + struct TestInterceptor; + + impl Interceptor for TestInterceptor { + fn call(&mut self, request: tonic::Request<()>) -> Result, Status> { + Ok(request) + } + } + + fn manager( + flow_control: FlowControl, + ) -> StreamingPullManager { + let channel = tonic::transport::Endpoint::from_static("http://[::]:50051").connect_lazy(); + let client = SubscriberClient::with_interceptor(channel, TestInterceptor); + let (sender, receiver) = unbounded_channel(); + let message_lease_renewal_duration_secs = flow_control.min_duration_per_lease_extension; + let message_latency_histogram = Histogram::new_with_max( + (flow_control.max_duration_per_lease_extension * 1000) as _, + 3, + ) + .unwrap(); + + StreamingPullManager { + receiver, + state: StreamingPullManagerState { + client, + flow_control, + message_lease_renewal_duration_secs, + messages_on_hold: Default::default(), + client_id: uuid::Uuid::nil(), + subscription: "projects/test/subscriptions/test".into(), + message_ack_queue: MessageAckQueue::new( + 25_001, + ACK_IDS_MAX_BATCH_SIZE, + Duration::ZERO, + ), + messages_in_flight: Default::default(), + graceful_stop_join_handles: None, + metrics: StreamingPullManagerMetrics::new(), + message_latency_histogram, + sender, + background_tasks: JoinSet::new(), + streaming_pull_channel_async_backoff: async_backoff(), + }, + channel: StreamingPullChannel::Closed, + message_handler: TestMessageHandler, + _metrics_emit_worker_abort_on_drop: tokio::spawn(async { + futures::future::pending::<()>().await; + }) + .into(), + } + } + + fn insert_due(manager: &mut StreamingPullManager) { + let renew_at = tokio::time::Instant::now(); + for index in 0..25_001 { + manager.state.message_ack_queue.transform_and_store_ack_id( + format!("server-{index}"), + Instant::now(), + renew_at, + ); + } + } + + #[test] + fn renewal_interval_covers_maximum_working_set_before_minimum_deadline() { + let buffered_deadline = buffered_lease_renewal_duration(10); + let catch_up_window = buffered_deadline - INITIAL_LEASE_RENEWAL_DELAY; + let batch_count = catch_up_window.as_nanos() / MIN_LEASE_RENEWAL_INTERVAL.as_nanos(); + let max_messages = usize::try_from(batch_count) + .unwrap() + .checked_mul(ACK_IDS_MAX_BATCH_SIZE) + .unwrap(); + let flow_control = FlowControl::new() + .set_max_messages(max_messages) + .set_duration_per_lease_extension(10, 10); + let interval = lease_renewal_interval(&flow_control, buffered_deadline); + let batches = flow_control.max_messages.div_ceil(ACK_IDS_MAX_BATCH_SIZE); + + assert!(interval < LEASE_RENEWAL_INTERVAL); + assert!( + interval.as_nanos() * u128::try_from(batches).unwrap() <= catch_up_window.as_nanos() + ); + } + + #[test] + #[should_panic(expected = "max_messages must be greater than zero")] + fn renewal_interval_rejects_zero_max_messages() { + lease_renewal_interval( + &FlowControl::new().set_max_messages(0), + Duration::from_secs(8), + ); + } + + #[test] + #[should_panic(expected = "cannot renew all lease batches")] + fn renewal_interval_rejects_one_batch_above_the_supported_limit() { + let max_messages = 7_001 * ACK_IDS_MAX_BATCH_SIZE; + let flow_control = FlowControl::new() + .set_max_messages(max_messages) + .set_duration_per_lease_extension(10, 10); + + lease_renewal_interval(&flow_control, buffered_lease_renewal_duration(10)); + } + + #[test] + #[should_panic(expected = "cannot renew all lease batches")] + fn renewal_interval_rejects_a_batch_count_above_u32() { + let batch_count = usize::try_from(u64::from(u32::MAX)) + .unwrap() + .checked_add(1) + .unwrap(); + let max_messages = batch_count.checked_mul(ACK_IDS_MAX_BATCH_SIZE).unwrap(); + let flow_control = FlowControl::new().set_max_messages(max_messages); + + lease_renewal_interval(&flow_control, Duration::from_secs(8)); + } + + #[test] + #[should_panic(expected = "cannot renew all lease batches")] + fn renewal_interval_rejects_usize_max_messages() { + lease_renewal_interval( + &FlowControl::new().set_max_messages(usize::MAX), + Duration::from_secs(8), + ); + } + + #[tokio::test(start_paused = true)] + async fn renewal_action_sends_one_request_and_later_actions_cover_due_leases() { + let flow_control = FlowControl::new() + .set_max_messages(25_001) + .set_max_processing_messages(0) + .set_duration_per_lease_extension(10, 10); + let mut manager = manager(flow_control); + insert_due(&mut manager); + tokio::time::advance(Duration::from_secs(1)).await; + + manager.handle_renew_leases(); + assert_eq!(manager.state.background_tasks.len(), 1); + + for _ in 1..11 { + manager.handle_renew_leases(); + } + + assert_eq!(manager.state.background_tasks.len(), 11); + assert!(manager + .state + .message_ack_queue + .collect_ack_ids_that_need_to_be_renewed(ACK_IDS_MAX_BATCH_SIZE, Duration::from_secs(8)) + .is_empty()); + } +} diff --git a/osprey_worker/src/osprey/engine/ast/grammar.py b/osprey_worker/src/osprey/engine/ast/grammar.py index 6db83666..8ec473c9 100644 --- a/osprey_worker/src/osprey/engine/ast/grammar.py +++ b/osprey_worker/src/osprey/engine/ast/grammar.py @@ -1,11 +1,12 @@ from __future__ import annotations from collections import defaultdict -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace from enum import Enum from pathlib import Path from threading import Lock +from types import MappingProxyType from typing import ClassVar, TypeVar # TODO: Uncomment logging when we have a logging system @@ -417,8 +418,12 @@ def find_argument(self, name: str) -> 'Keyword' | None: return None - def argument_dict(self) -> dict[str, Expression]: - return {arg.name: arg.value for arg in self.arguments} + def argument_dict(self) -> Mapping[str, Expression]: + return self._argument_dict_cached + + @cached_property + def _argument_dict_cached(self) -> Mapping[str, Expression]: + return MappingProxyType({arg.name: arg.value for arg in self.arguments}) @property def can_extract(self) -> bool: diff --git a/osprey_worker/src/osprey/engine/ast/tests/test_grammar.py b/osprey_worker/src/osprey/engine/ast/tests/test_grammar.py new file mode 100644 index 00000000..294bade4 --- /dev/null +++ b/osprey_worker/src/osprey/engine/ast/tests/test_grammar.py @@ -0,0 +1,24 @@ +import pytest +from osprey.engine.ast.ast_utils import filter_nodes +from osprey.engine.ast.grammar import Call, Source + + +def _parse_call(contents: str) -> Call: + source = Source(path='test.sml', contents=contents) + return next(iter(filter_nodes(source.ast_root, Call))) + + +def test_call_argument_dict_reuses_the_same_mapping() -> None: + call = _parse_call('Result = Function(a=1, b=2)\n') + + assert call.argument_dict() is call.argument_dict() + + +def test_call_argument_dict_cannot_be_mutated() -> None: + call = _parse_call('Result = Function(a=1, b=2)\n') + arguments = call.argument_dict() + + with pytest.raises(TypeError): + arguments['a'] = arguments['b'] # type: ignore[index] # This assignment verifies the runtime mutation guard. + + assert call.argument_dict()['a'] is arguments['a'] diff --git a/osprey_worker/src/osprey/engine/executor/dependency_chain.py b/osprey_worker/src/osprey/engine/executor/dependency_chain.py index 6666af88..e212cfd4 100644 --- a/osprey_worker/src/osprey/engine/executor/dependency_chain.py +++ b/osprey_worker/src/osprey/engine/executor/dependency_chain.py @@ -10,10 +10,16 @@ @add_slots -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class DependencyChain: """The dependency chain stores the requisite dependency chains that must be executed before the node executor - is able to be executed.""" + is able to be executed. + + eq=False: `executor` is a plain class with default identity equality, so the generated structural + __eq__/__hash__ already bottomed out at executor-object identity (dataclass field equality falls back + to `is` for any field without its own __eq__). Declaring eq=False makes that explicit and gives us + O(1) identity hash instead of an O(size-of-subtree) structural hash recomputed on every call, since + tuple.__hash__ does not cache its result the way str/frozenset do.""" executor: 'BaseNodeExecutor[ASTNode, object]' """The executor that we are holding dependencies for.""" diff --git a/osprey_worker/src/osprey/engine/executor/execution_context.py b/osprey_worker/src/osprey/engine/executor/execution_context.py index 56ff6149..f687dab6 100644 --- a/osprey_worker/src/osprey/engine/executor/execution_context.py +++ b/osprey_worker/src/osprey/engine/executor/execution_context.py @@ -23,6 +23,7 @@ ) from osprey.engine.executor.dependency_chain import DependencyChain from osprey.engine.executor.execution_graph import ExecutionGraph +from osprey.engine.executor.execution_plan import ExecutionPlanState from osprey.engine.executor.external_service_utils_base import ( ExternalService, KeyT, @@ -54,7 +55,7 @@ from osprey.engine.language_types.verdicts import VerdictEffect from osprey.engine.utils.types import add_slots, cached_property from osprey.rpc.common.v1.verdicts_pb2 import Verdicts -from result import Result, UnwrapError +from result import Ok, Result, UnwrapError if TYPE_CHECKING: from osprey.engine.ast_validator.validation_context import ValidatedSources @@ -129,15 +130,15 @@ class ExecutionContext: '_input_encoding', '_execution_graph', '_outputs', - '_pending_executions', '_resolved_node_values', - '_visited_executions', '_effects', '_udf_helpers', '_external_service_accessors_by_getter_id', '_async_external_service_accessors_by_getter_id', '_dependency_dag', + '_execution_plan_state', '_chain_by_id', + '_enqueued_sources', '_custom_extracted_features', '_rule_audit_entries', ) @@ -149,15 +150,16 @@ def __init__(self, execution_graph: ExecutionGraph, action: 'Action', helpers: U self._execution_graph = execution_graph self._udf_helpers: UDFHelpers = helpers self._outputs: dict[str, Any] = {} - self._pending_executions: set[DependencyChain] = set() self._resolved_node_values: dict[int, NodeResult] = {} - self._visited_executions: set[DependencyChain] = set() # a k/v store of effects, by effect type self._effects: defaultdict[Type[EffectBase], list[EffectBase]] = defaultdict(list) self._external_service_accessors_by_getter_id: dict[int, Any] = {} self._async_external_service_accessors_by_getter_id: dict[int, Any] = {} - self._dependency_dag = TopologicalSorter() + plan = execution_graph.get_execution_plan() + self._execution_plan_state = ExecutionPlanState(plan) if plan is not None else None + self._dependency_dag = TopologicalSorter() if plan is None else None self._chain_by_id: dict[int, DependencyChain] = {} + self._enqueued_sources: set[Source] = set() # feature name -> serializable feature self._custom_extracted_features: dict[str, Any] = {} self._rule_audit_entries: list[WhenRulesAuditEntry] = [] @@ -175,23 +177,29 @@ def resolved(self, node: ASTNode, return_none_for_failed_values: bool = False) - failed, will either raise a NodeFailurePropagationException (default) or return None (if return_none_for_failed_values is True). """ - # We need to check this on the original node, not (say) the assignment node if this is a Name. + try: + return self.resolved_result(node).unwrap() + except UnwrapError: + if return_none_for_failed_values: + return None + raise NodeFailurePropagationException() + + def resolved_result(self, node: ASTNode) -> NodeResult: + """Return the resolved result for a node after post-execution conversion.""" + # Check the original node before Name resolution. should_unwrap = self._execution_graph.should_unwrap(node) if isinstance(node, Name): node = self.get_name_node(node) - try: - value = self._resolved_node_values[id(node)].unwrap() - if should_unwrap: - assert isinstance(value, PostExecutionConvertible), (value, type(value)) - value = value.to_post_execution_value() - return value - except UnwrapError: - if return_none_for_failed_values: - return None - else: - raise NodeFailurePropagationException() + node_result = self._resolved_node_values[id(node)] + + if should_unwrap and node_result.is_ok(): + value = node_result.unwrap() + assert isinstance(value, PostExecutionConvertible), (value, type(value)) + return Ok(value.to_post_execution_value()) + + return node_result def get_name_node(self, name: Name) -> ASTNode: """Returns the node that is responsible for resolving a given Loaded name.""" @@ -201,7 +209,11 @@ def get_name_node(self, name: Name) -> ASTNode: def set_resolved_value(self, chain: DependencyChain, value: NodeResult) -> None: """Called by the main executor once a node has been resolved, to store its value for dependent executors.""" self._resolved_node_values[id(chain.executor.node)] = value - self._dependency_dag.done(id(chain)) + if self._execution_plan_state is not None: + self._execution_plan_state.done(chain) + else: + assert self._dependency_dag is not None + self._dependency_dag.done(id(chain)) def set_output_value(self, key: str, value: Any) -> None: """Called by the assignment node executor to store an output key/value pair.""" @@ -240,18 +252,34 @@ def get_action_time(self) -> datetime: return self._action.timestamp def enqueue_source(self, source: Source) -> None: + if source in self._enqueued_sources: + return + + if self._execution_plan_state is not None: + self._execution_plan_state.activate_source(source) + else: + self._enqueue_source_legacy(source) + self._enqueued_sources.add(source) + + def _enqueue_source_legacy(self, source: Source) -> None: + assert self._dependency_dag is not None sorted_dependency_chain = self._execution_graph.get_sorted_dependency_chain(source) for chain in sorted_dependency_chain: chainid = id(chain) if self._dependency_dag.already_added(chainid): continue - self._dependency_dag.add(chainid, *(id(pred) for pred in chain.dependent_on)) + predecessor_ids = [id(predecessor) for predecessor in chain.dependent_on] + self._dependency_dag.add(chainid, *predecessor_ids) self._chain_by_id[chainid] = chain self._dependency_dag.prepare() def get_ready_to_execute(self) -> Sequence[DependencyChain]: + if self._execution_plan_state is not None: + return self._execution_plan_state.get_ready() + + assert self._dependency_dag is not None ready_nodeids = self._dependency_dag.get_ready() ready = [] for nodeid in ready_nodeids: diff --git a/osprey_worker/src/osprey/engine/executor/execution_graph.py b/osprey_worker/src/osprey/engine/executor/execution_graph.py index fe506fd7..9e47f0ee 100644 --- a/osprey_worker/src/osprey/engine/executor/execution_graph.py +++ b/osprey_worker/src/osprey/engine/executor/execution_graph.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from collections.abc import Hashable, Iterator, Sequence from typing import TYPE_CHECKING, Any, TypeVar @@ -11,11 +12,13 @@ if TYPE_CHECKING: from osprey.engine.ast_validator.validation_context import ValidatedSources + from .execution_plan import ExecutionPlan from .node_executor._base_node_executor import BaseNodeExecutor from .node_executor_registry import NodeExecutorRegistry T = TypeVar('T', bound=Hashable) +logger = logging.getLogger(__name__) class ExecutionGraph: @@ -30,6 +33,7 @@ class ExecutionGraph: '_validated_sources', '_sorted_dependency_chains', '_nodes_to_unwrap', + '_execution_plan', ) _root_node_executor_mapping: dict[int, DependencyChain] @@ -51,6 +55,9 @@ class ExecutionGraph: _nodes_to_unwrap: set[int] """ID's for nodes that need to be unwrapped to its inner type when used.""" + _execution_plan: 'ExecutionPlan | None' + """An immutable schedule plan for this graph.""" + def __init__( self, node_executor_registry: 'NodeExecutorRegistry', sources: 'ValidatedSources', nodes_to_unwrap: set[int] ): @@ -60,6 +67,7 @@ def __init__( self._validated_sources = sources self._sorted_dependency_chains = {} self._nodes_to_unwrap = nodes_to_unwrap + self._execution_plan = None @property def validated_sources(self) -> 'ValidatedSources': @@ -87,6 +95,9 @@ def should_unwrap(self, node: ASTNode) -> bool: """Whether we need to unwrap the value that is represented by this node before using it.""" return id(node) in self._nodes_to_unwrap + def get_execution_plan(self) -> 'ExecutionPlan | None': + return self._execution_plan + def _get_executor_for(self, node: ASTNode) -> 'BaseNodeExecutor[Any, Any]': return self._node_executor_registry.construct_executor_for(node, validated_sources=self._validated_sources) @@ -96,7 +107,8 @@ def _build_dependency_chain(self, node: ASTNode) -> DependencyChain: return self.get_assignment_dependency_chain(node) executor = self._get_executor_for(node) - dependent_on = tuple(self._build_dependency_chain(node) for node in executor.get_dependent_nodes()) + dependent_chains = [self._build_dependency_chain(node) for node in executor.get_dependent_nodes()] + dependent_on = tuple(dependent_chains) return DependencyChain(executor=executor, dependent_on=dependent_on) def _add_validated_source(self, source: Source) -> None: @@ -119,6 +131,7 @@ def compile_execution_graph( from osprey.engine.ast_validator.validators.imports_must_not_have_cycles import ImportsMustNotHaveCycles from osprey.engine.ast_validator.validators.validate_static_types import ValidateStaticTypes + from .execution_plan import ExecutionPlan from .node_executor_registry import NodeExecutorRegistry node_executor_registry = node_executor_registry or NodeExecutorRegistry.get_instance() @@ -149,6 +162,12 @@ def compile_execution_graph( instance._add_sorted_dependency_chain(source, sorted_dependency_chain) maybe_periodic_yield() + execution_plan = ExecutionPlan.from_graph(instance) + plan_error = execution_plan.find_unclosed_source() + if plan_error is None: + instance._execution_plan = execution_plan + else: + logger.error('Execution plan is invalid. The graph scheduler will run: %s', plan_error) return instance diff --git a/osprey_worker/src/osprey/engine/executor/execution_plan.py b/osprey_worker/src/osprey/engine/executor/execution_plan.py new file mode 100644 index 00000000..2d0d1d66 --- /dev/null +++ b/osprey_worker/src/osprey/engine/executor/execution_plan.py @@ -0,0 +1,191 @@ +"""Immutable schedule data for a full execution graph. + +The graph shares one plan across actions. Each action creates separate schedule state. +""" + +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Mapping + +from osprey.engine.ast.grammar import Source +from osprey.engine.utils.periodic_execution_yielder import maybe_periodic_yield + +from .dependency_chain import DependencyChain + +if TYPE_CHECKING: + from .execution_graph import ExecutionGraph + + +@dataclass(frozen=True, slots=True, weakref_slot=True) +class ExecutionPlan: + """Store shared indices and edges for a full execution graph.""" + + chains: tuple[DependencyChain, ...] + index_by_chain_id: Mapping[int, int] + predecessors: tuple[tuple[int, ...], ...] + successors: tuple[tuple[int, ...], ...] + source_indices: Mapping[Source, tuple[int, ...]] + + @classmethod + def from_graph(cls, graph: 'ExecutionGraph') -> 'ExecutionPlan': + chains_by_id: dict[int, DependencyChain] = {} + source_chain_ids: dict[Source, tuple[int, ...]] = {} + for source in graph.validated_sources.sources: + chains = tuple(graph.get_sorted_dependency_chain(source)) + source_chain_ids[source] = tuple([id(chain) for chain in chains]) + for chain in chains: + chains_by_id[id(chain)] = chain + maybe_periodic_yield() + + chains = tuple(chains_by_id.values()) + del chains_by_id + + index_by_chain_id: dict[int, int] = {} + for index, chain in enumerate(chains): + index_by_chain_id[id(chain)] = index + maybe_periodic_yield() + + source_indices: dict[Source, tuple[int, ...]] = {} + for source, chain_ids in source_chain_ids.items(): + indices: list[int] = [] + for chain_id in chain_ids: + indices.append(index_by_chain_id[chain_id]) + maybe_periodic_yield() + source_indices[source] = tuple(indices) + del source_chain_ids + + predecessor_tuples: list[tuple[int, ...]] = [] + for chain in chains: + chain_predecessor_indices = [index_by_chain_id[id(predecessor)] for predecessor in chain.dependent_on] + predecessor_tuples.append(tuple(chain_predecessor_indices)) + maybe_periodic_yield() + predecessors = tuple(predecessor_tuples) + del predecessor_tuples + + successor_lists: list[list[int]] = [[] for _ in chains] + for successor, predecessor_indices in enumerate(predecessors): + for predecessor in predecessor_indices: + successor_lists[predecessor].append(successor) + maybe_periodic_yield() + + successor_tuples: list[tuple[int, ...]] = [] + for items in successor_lists: + successor_tuples.append(tuple(items)) + items.clear() + maybe_periodic_yield() + del successor_lists + successors = tuple(successor_tuples) + del successor_tuples + + return cls( + chains=chains, + index_by_chain_id=MappingProxyType(index_by_chain_id), + predecessors=predecessors, + successors=successors, + source_indices=MappingProxyType(source_indices), + ) + + def find_unclosed_source(self) -> str | None: + """Return a diagnostic if a source omits a predecessor.""" + for source, indices in self.source_indices.items(): + activated = set(indices) + for index in indices: + for predecessor in self.predecessors[index]: + if predecessor not in activated: + node = self.chains[index].executor.node + span = node.span + return ( + f'source {source.path!r} activates chain {index} ' + f'({type(node).__name__} at {span.source.path}:{span.start_line}:{span.start_pos}) ' + f'without predecessor chain {predecessor}' + ) + maybe_periodic_yield() + return None + + +_INACTIVE = -3 +_OUT = -1 +_DONE = -2 + + +class LateDependencyActivationError(RuntimeError): + """Report a source activation that violates the plan dependency order.""" + + +class ExecutionPlanState: + """Store the compact schedule state for one action.""" + + __slots__ = ('_plan', '_active', '_remaining', '_activation_rank', '_next_rank', '_ready') + + def __init__(self, plan: ExecutionPlan) -> None: + self._plan = plan + self._active = bytearray(len(plan.chains)) + self._remaining = [_INACTIVE] * len(plan.chains) + self._activation_rank = [_INACTIVE] * len(plan.chains) + self._next_rank = 0 + self._ready: list[int] = [] + + def activate_source(self, source: Source) -> None: + new_indices = tuple([index for index in self._plan.source_indices[source] if not self._active[index]]) + if not new_indices: + return + + new_set = set(new_indices) + for index in new_indices: + active_successors = tuple( + [successor for successor in self._plan.successors[index] if self._active[successor]] + ) + if active_successors: + node = self._plan.chains[index].executor.node + span = node.span + raise LateDependencyActivationError( + f'source {source.path!r} would activate chain {index} ' + f'({type(node).__name__} at {span.source.path}:{span.start_line}:{span.start_pos}) ' + f'after active successor chains {active_successors}' + ) + + counts = tuple( + [ + sum( + 1 + for predecessor in self._plan.predecessors[index] + if (self._active[predecessor] or predecessor in new_set) and self._remaining[predecessor] != _DONE + ) + for index in new_indices + ] + ) + + for index, count in zip(new_indices, counts): + self._active[index] = 1 + self._remaining[index] = count + self._activation_rank[index] = self._next_rank + self._next_rank += 1 + if count == 0: + self._ready.append(index) + + self._ready.sort(key=self._activation_rank.__getitem__) + + def get_ready(self) -> tuple[DependencyChain, ...]: + indices = tuple(self._ready) + self._ready.clear() + for index in indices: + self._remaining[index] = _OUT + ready_chains = [self._plan.chains[index] for index in indices] + return tuple(ready_chains) + + def done(self, chain: DependencyChain) -> None: + index = self._plan.index_by_chain_id[id(chain)] + if self._remaining[index] != _OUT: + raise ValueError(f'chain {index} was not passed out') + + self._remaining[index] = _DONE + newly_ready: list[int] = [] + for successor in self._plan.successors[index]: + if not self._active[successor] or self._remaining[successor] < 0: + continue + self._remaining[successor] -= 1 + if self._remaining[successor] == 0: + newly_ready.append(successor) + + newly_ready.sort(key=self._activation_rank.__getitem__) + self._ready.extend(newly_ready) diff --git a/osprey_worker/src/osprey/engine/executor/executor.py b/osprey_worker/src/osprey/engine/executor/executor.py index cc9c113a..e245e13b 100644 --- a/osprey_worker/src/osprey/engine/executor/executor.py +++ b/osprey_worker/src/osprey/engine/executor/executor.py @@ -35,6 +35,8 @@ logger = get_logger(__name__) +_UNSET_RESULT: NodeResult = Err(None) + InProgressSingletsType = dict['gevent.Greenlet[NodeResult]', DependencyChain] """ A dictionary mapping in-progress async greenlets to the chain that they are executing. @@ -193,18 +195,30 @@ def _wrapped_execution( ) -> NodeResult: caught_exception: Exception | None = None - metric_tags = _get_metric_tags(context) + call_node: CallExecutor | None = None if isinstance(chain.executor, CallExecutor): # This half step is necessary as mypy has a difficult time linting build in class variables - call_node: CallExecutor = chain.executor - metric_tags += [f'udf:{call_node._udf.__class__.__name__}'] + call_node = chain.executor + + # Most nodes never emit a metric (sync, non-UDF nodes with no exception), so tags are built + # lazily on first use and memoized -- avoids paying for an 8-element list + 3 f-strings on + # every node execution when nothing ends up consuming them. + _metric_tags: list[str] | None = None + + def _get_tags() -> list[str]: + nonlocal _metric_tags + if _metric_tags is None: + _metric_tags = _get_metric_tags(context) + if call_node is not None: + _metric_tags += [f'udf:{call_node._udf.__class__.__name__}'] + return _metric_tags # Make mypy happy - execution_result: NodeResult = Err(None) + execution_result: NodeResult = _UNSET_RESULT try: # only track time if using an async function if chain.executor.execute_async: - with metrics.timed('udf_execution_duration', tags=metric_tags, sample_rate=0.01): + with metrics.timed('udf_execution_duration', tags=_get_tags(), sample_rate=0.01): execution_result = Ok(chain.executor.execute(execution_context=context)) else: execution_result = Ok(chain.executor.execute(execution_context=context)) @@ -217,16 +231,16 @@ def _wrapped_execution( finally: # If this is a call node which executed a UDF, push the results of the execution to the datadog metrics. - if isinstance(chain.executor, CallExecutor): - if execution_result.is_ok() and chain.executor and chain.executor.execute_async: - metrics.increment('udf_execution', tags=metric_tags + ['exc_name:none', 'result:success']) + if call_node is not None: + if execution_result.is_ok() and call_node.execute_async: + metrics.increment('udf_execution', tags=_get_tags() + ['exc_name:none', 'result:success']) # Ignore some well-known "unexpected" exceptions that are spammy. elif not _is_spammy_exception(caught_exception): exc_name = caught_exception.__class__.__name__ metrics.increment( 'udf_execution', - tags=metric_tags + [f'exc_name:{exc_name}', 'result:unexpected_failure'], + tags=_get_tags() + [f'exc_name:{exc_name}', 'result:unexpected_failure'], ) return execution_result diff --git a/osprey_worker/src/osprey/engine/executor/node_executor/assign_executor.py b/osprey_worker/src/osprey/engine/executor/node_executor/assign_executor.py index 06015570..612b78e1 100644 --- a/osprey_worker/src/osprey/engine/executor/node_executor/assign_executor.py +++ b/osprey_worker/src/osprey/engine/executor/node_executor/assign_executor.py @@ -2,6 +2,7 @@ from osprey.engine.ast.grammar import Assign, ASTNode +from ..execution_context import NodeFailurePropagationException from ..node_executor_registry import NodeExecutorRegistry from ._base_node_executor import BaseNodeExecutor @@ -14,15 +15,14 @@ class AssignExecutor(BaseNodeExecutor[Assign, Any]): node_type = Assign def execute(self, execution_context: 'ExecutionContext') -> Any: - # We want to store a value in the output even if the dependency node failed. - resolved_maybe_error = execution_context.resolved(self._node.value, return_none_for_failed_values=True) + node_result = execution_context.resolved_result(self._node.value) if self._node.should_extract: - execution_context.set_output_value(self._node.target.identifier, resolved_maybe_error) + execution_context.set_output_value(self._node.target.identifier, node_result.ok()) - # Re-fetch to throw an exception in case our dependency node failed. - resolved_not_error = execution_context.resolved(self._node.value, return_none_for_failed_values=False) - return resolved_not_error + if node_result.is_err(): + raise NodeFailurePropagationException() + return node_result.unwrap() def get_dependent_nodes(self) -> list[ASTNode]: return [self._node.value] diff --git a/osprey_worker/src/osprey/engine/executor/tests/conftest.py b/osprey_worker/src/osprey/engine/executor/tests/conftest.py index e69de29b..75015c63 100644 --- a/osprey_worker/src/osprey/engine/executor/tests/conftest.py +++ b/osprey_worker/src/osprey/engine/executor/tests/conftest.py @@ -0,0 +1,14 @@ +import pytest +from osprey.engine.conftest import RunValidationFunction +from osprey.engine.executor.execution_graph import ExecutionGraph, compile_execution_graph + + +@pytest.fixture +def compiled_execution_graph(run_validation: RunValidationFunction) -> ExecutionGraph: + validated = run_validation( + { + 'main.sml': 'First = 1 + 2', + 'secondary.sml': 'Second = First + 3', + } + ) + return compile_execution_graph(validated) diff --git a/osprey_worker/src/osprey/engine/executor/tests/test_assign_executor.py b/osprey_worker/src/osprey/engine/executor/tests/test_assign_executor.py new file mode 100644 index 00000000..8a2189e4 --- /dev/null +++ b/osprey_worker/src/osprey/engine/executor/tests/test_assign_executor.py @@ -0,0 +1,35 @@ +from unittest.mock import Mock + +import pytest +from osprey.engine.ast.ast_utils import filter_nodes +from osprey.engine.ast.grammar import Assign, Source +from osprey.engine.executor.execution_context import ExecutionContext, NodeFailurePropagationException +from osprey.engine.executor.node_executor.assign_executor import AssignExecutor +from result import Err, Ok + + +def _parse_assign(contents: str) -> Assign: + source = Source(path='test.sml', contents=contents) + return next(iter(filter_nodes(source.ast_root, Assign))) + + +def test_assign_resolves_its_value_once() -> None: + assign = _parse_assign('Result = 1\n') + execution_context = Mock(spec=ExecutionContext) + execution_context.resolved_result.return_value = Ok(1) + executor = AssignExecutor(assign, Mock()) + + assert executor.execute(execution_context) == 1 + execution_context.resolved_result.assert_called_once_with(assign.value) + + +def test_assign_extracts_none_and_propagates_a_failed_value() -> None: + assign = _parse_assign('Result = Function()\n') + execution_context = Mock(spec=ExecutionContext) + execution_context.resolved_result.return_value = Err(None) + executor = AssignExecutor(assign, Mock()) + + with pytest.raises(NodeFailurePropagationException): + executor.execute(execution_context) + + execution_context.set_output_value.assert_called_once_with('Result', None) diff --git a/osprey_worker/src/osprey/engine/executor/tests/test_dependency_chain.py b/osprey_worker/src/osprey/engine/executor/tests/test_dependency_chain.py new file mode 100644 index 00000000..81359a8d --- /dev/null +++ b/osprey_worker/src/osprey/engine/executor/tests/test_dependency_chain.py @@ -0,0 +1,90 @@ +import gc + +import pytest +from osprey.engine.executor.dependency_chain import DependencyChain +from osprey.engine.executor.execution_graph import ExecutionGraph + + +class _DummyExecutor: + """Stand-in for a BaseNodeExecutor: a plain object with default identity eq/hash, exactly + like the real thing (no node executor subclass overrides __eq__/__hash__).""" + + def __init__(self, dependent_nodes: tuple[object, ...] = ()) -> None: + self._dependent_nodes = dependent_nodes + + def get_dependent_nodes(self) -> tuple[object, ...]: + return self._dependent_nodes + + +def test_chain_equals_itself() -> None: + executor = _DummyExecutor() + chain = DependencyChain(executor=executor, dependent_on=()) + same_chain = chain + + assert chain == same_chain + assert chain in {chain} + assert {chain: 'value'}[chain] == 'value' + + +def test_chains_with_same_executor_and_deps_are_not_equal() -> None: + """Two independently-constructed chains that happen to wrap the same executor and the same + dependent_on tuple are NOT equal under identity semantics, even though they would have been + dataclass-equal (and hash-colliding) under the old structural eq=True default. + + This is intentional: the engine never builds two such "equal-but-distinct" chains for the + same logical node -- `ExecutionGraph._build_dependency_chain` either mints a brand-new + executor per chain (so a genuinely different chain never collides with an existing one) or + returns the literal cached chain object for a `Load`-context Name (so a shared node is the + same object, not merely an equal one). Treating two distinct chain objects as interchangeable + just because their fields happen to match is not a case this engine relies on. + """ + executor = _DummyExecutor() + leaf = DependencyChain(executor=executor, dependent_on=()) + + chain_a = DependencyChain(executor=executor, dependent_on=(leaf,)) + chain_b = DependencyChain(executor=executor, dependent_on=(leaf,)) + + assert chain_a is not chain_b + assert chain_a != chain_b + assert chain_b not in {chain_a} + + +def test_distinct_executors_are_not_equal() -> None: + chain_a = DependencyChain(executor=_DummyExecutor(), dependent_on=()) + chain_b = DependencyChain(executor=_DummyExecutor(), dependent_on=()) + + assert chain_a != chain_b + + +def test_build_dependency_chain_does_not_resize_an_observable_tuple(monkeypatch: pytest.MonkeyPatch) -> None: + root, left, right = object(), object(), object() + executors = { + root: _DummyExecutor((left, right)), + left: _DummyExecutor(), + right: _DummyExecutor(), + } + monkeypatch.setattr(ExecutionGraph, '_get_executor_for', lambda _graph, node: executors[node]) + + original_build = ExecutionGraph._build_dependency_chain + previous_chain: DependencyChain | None = None + held_tuples: list[tuple[object, ...]] = [] + + def observing_build(graph: ExecutionGraph, node: object) -> DependencyChain: + nonlocal previous_chain + if previous_chain is not None: + held_tuples.extend( + referrer + for referrer in gc.get_referrers(previous_chain) + if isinstance(referrer, tuple) and len(referrer) > len(executors[root].get_dependent_nodes()) + ) + chain = original_build(graph, node) # type: ignore[arg-type] + previous_chain = chain + return chain + + monkeypatch.setattr(ExecutionGraph, '_build_dependency_chain', observing_build) + graph = object.__new__(ExecutionGraph) + + chain = graph._build_dependency_chain(root) # type: ignore[arg-type] + + assert len(chain.dependent_on) == 2 + assert held_tuples == [] diff --git a/osprey_worker/src/osprey/engine/executor/tests/test_execution_context.py b/osprey_worker/src/osprey/engine/executor/tests/test_execution_context.py new file mode 100644 index 00000000..fb828ce3 --- /dev/null +++ b/osprey_worker/src/osprey/engine/executor/tests/test_execution_context.py @@ -0,0 +1,78 @@ +from datetime import datetime, timezone +from unittest.mock import Mock, call, patch + +import pytest +from osprey.engine.ast.grammar import Source +from osprey.engine.executor.execution_context import Action, ExecutionContext +from osprey.engine.executor.execution_graph import ExecutionGraph +from osprey.engine.executor.execution_plan import ExecutionPlan +from osprey.engine.executor.udf_execution_helpers import UDFHelpers + + +def _action() -> Action: + return Action( + action_id=1, + action_name='test_action', + data={}, + timestamp=datetime(2026, 8, 4, tzinfo=timezone.utc), + ) + + +def test_context_uses_independent_plan_state(compiled_execution_graph: ExecutionGraph) -> None: + first = ExecutionContext(compiled_execution_graph, _action(), Mock(spec=UDFHelpers)) + second = ExecutionContext(compiled_execution_graph, _action(), Mock(spec=UDFHelpers)) + + assert first._execution_plan_state is not None + assert second._execution_plan_state is not None + assert first._execution_plan_state is not second._execution_plan_state + assert first._execution_plan_state._plan is second._execution_plan_state._plan + + +def test_enqueue_source_retries_after_enqueue_failure() -> None: + entry_source = Source(path='main.sml', contents='') + dynamic_source = Source(path='dynamic.sml', contents='') + execution_graph = Mock(spec=ExecutionGraph) + execution_graph.get_entry_point.return_value = entry_source + execution_graph.get_execution_plan.return_value = None + execution_graph.get_sorted_dependency_chain.return_value = () + action = Action( + action_id=1, + action_name='test_action', + data={}, + timestamp=datetime(2026, 8, 3, tzinfo=timezone.utc), + ) + context = ExecutionContext(execution_graph, action, Mock(spec=UDFHelpers)) + execution_graph.get_sorted_dependency_chain.reset_mock() + execution_graph.get_sorted_dependency_chain.side_effect = [RuntimeError('enqueue failed'), ()] + + with pytest.raises(RuntimeError, match='enqueue failed'): + context.enqueue_source(dynamic_source) + context.enqueue_source(dynamic_source) + + assert execution_graph.get_sorted_dependency_chain.call_args_list == [ + call(dynamic_source), + call(dynamic_source), + ] + + +def test_planned_enqueue_source_retries_after_activation_failure() -> None: + entry_source = Source(path='main.sml', contents='') + dynamic_source = Source(path='dynamic.sml', contents='') + execution_graph = Mock(spec=ExecutionGraph) + execution_graph.get_entry_point.return_value = entry_source + execution_graph.get_execution_plan.return_value = Mock(spec=ExecutionPlan) + + with patch('osprey.engine.executor.execution_context.ExecutionPlanState') as plan_state_type: + plan_state = plan_state_type.return_value + context = ExecutionContext(execution_graph, _action(), Mock(spec=UDFHelpers)) + plan_state.activate_source.reset_mock() + plan_state.activate_source.side_effect = [RuntimeError('activation failed'), None] + + with pytest.raises(RuntimeError, match='activation failed'): + context.enqueue_source(dynamic_source) + assert dynamic_source not in context._enqueued_sources + + context.enqueue_source(dynamic_source) + + assert plan_state.activate_source.call_args_list == [call(dynamic_source), call(dynamic_source)] + assert dynamic_source in context._enqueued_sources diff --git a/osprey_worker/src/osprey/engine/executor/tests/test_execution_plan.py b/osprey_worker/src/osprey/engine/executor/tests/test_execution_plan.py new file mode 100644 index 00000000..61acd33a --- /dev/null +++ b/osprey_worker/src/osprey/engine/executor/tests/test_execution_plan.py @@ -0,0 +1,343 @@ +import gc +import random +from collections.abc import Iterator +from dataclasses import dataclass +from types import MappingProxyType +from unittest.mock import patch + +import pytest +from osprey.engine.ast.grammar import Source, Span +from osprey.engine.conftest import RunValidationFunction +from osprey.engine.executor.dependency_chain import DependencyChain +from osprey.engine.executor.execution_graph import ExecutionGraph, compile_execution_graph +from osprey.engine.executor.execution_plan import ExecutionPlan, ExecutionPlanState, LateDependencyActivationError +from osprey.engine.executor.topological_sorter import TopologicalSorter + + +@dataclass(frozen=True) +class _PlanNode: + span: Span + + +@dataclass(frozen=True) +class _PlanExecutor: + node: _PlanNode + + +def test_full_graph_compiles_immutable_execution_plan(compiled_execution_graph: ExecutionGraph) -> None: + plan = compiled_execution_graph.get_execution_plan() + + assert isinstance(plan, ExecutionPlan) + assert isinstance(plan.index_by_chain_id, MappingProxyType) + assert isinstance(plan.source_indices, MappingProxyType) + assert len(plan.chains) == len(plan.index_by_chain_id) + assert set(plan.source_indices) == set(compiled_execution_graph.validated_sources.sources) + for source, indices in plan.source_indices.items(): + expected = compiled_execution_graph.get_sorted_dependency_chain(source) + assert tuple([plan.chains[index] for index in indices]) == tuple(expected) + + +def test_every_planned_predecessor_has_a_stable_index(compiled_execution_graph: ExecutionGraph) -> None: + plan = compiled_execution_graph.get_execution_plan() + + assert isinstance(plan, ExecutionPlan) + for chain_index, chain in enumerate(plan.chains): + expected = tuple([plan.index_by_chain_id[id(predecessor)] for predecessor in chain.dependent_on]) + assert plan.predecessors[chain_index] == expected + + +def test_plan_compilation_yields_through_large_phases(run_validation: RunValidationFunction) -> None: + validated = run_validation( + { + 'main.sml': 'First = 1 + 2', + 'secondary.sml': 'Second = First + 3', + } + ) + + with patch('osprey.engine.executor.execution_plan.maybe_periodic_yield') as periodic_yield: + graph = compile_execution_graph(validated) + + plan = graph.get_execution_plan() + assert isinstance(plan, ExecutionPlan) + chain_references = sum(len(indices) for indices in plan.source_indices.values()) + assert periodic_yield.call_count >= 4 * len(plan.chains) + 2 * chain_references + + +def _manual_plan( + predecessors: tuple[tuple[int, ...], ...], source_indices: tuple[tuple[int, ...], ...] +) -> tuple[ExecutionPlan, tuple[Source, ...]]: + chains: list[DependencyChain] = [] + for chain_predecessors in predecessors: + index = len(chains) + node_source = Source(path=f'chain-{index}.sml', contents='') + dependent_on = tuple([chains[predecessor] for predecessor in chain_predecessors]) + chains.append( + DependencyChain( + # The plan only reads executor.node, which this test stub provides. + executor=_PlanExecutor( # type: ignore[arg-type] + node=_PlanNode(span=Span(source=node_source, start_line=index + 1, start_pos=index)) + ), + dependent_on=dependent_on, + ) + ) + + successor_lists: list[list[int]] = [[] for _ in chains] + for successor, chain_predecessors in enumerate(predecessors): + for predecessor in chain_predecessors: + successor_lists[predecessor].append(successor) + + sources = tuple([Source(path=f'source-{index}.sml', contents='') for index in range(len(source_indices))]) + return ( + ExecutionPlan( + chains=tuple(chains), + index_by_chain_id={id(chain): index for index, chain in enumerate(chains)}, + predecessors=predecessors, + successors=tuple([tuple(items) for items in successor_lists]), + source_indices={source: indices for source, indices in zip(sources, source_indices)}, + ), + sources, + ) + + +def _ready_indices(state: ExecutionPlanState, plan: ExecutionPlan) -> tuple[int, ...]: + return tuple([plan.index_by_chain_id[id(chain)] for chain in state.get_ready()]) + + +class _TupleResizeObserver: + def __init__(self, chains: tuple[DependencyChain, ...]) -> None: + self._chains = chains + self._previous: DependencyChain | None = None + self._held_tuples: list[tuple[object, ...]] = [] + + def __getitem__(self, index: int) -> DependencyChain: + if self._previous is not None: + self._held_tuples.extend( + referrer + for referrer in gc.get_referrers(self._previous) + if isinstance(referrer, tuple) and len(referrer) > len(self._chains) + ) + chain = self._chains[index] + self._previous = chain + return chain + + +class _ObservedIndex(int): + pass + + +class _TupleResizeIndexSequence: + def __init__(self, indices: tuple[int, ...]) -> None: + self._indices = indices + self._held_tuples: list[tuple[object, ...]] = [] + + def __iter__(self) -> Iterator[int]: + previous: _ObservedIndex | None = None + for index in self._indices: + if previous is not None: + self._held_tuples.extend( + referrer + for referrer in gc.get_referrers(previous) + if isinstance(referrer, tuple) and len(referrer) > len(self._indices) + ) + previous = _ObservedIndex(index) + yield previous + + +def test_get_ready_does_not_resize_an_observable_tuple() -> None: + plan, (source,) = _manual_plan(predecessors=((), ()), source_indices=((0, 1),)) + state = ExecutionPlanState(plan) + expected_ready = plan.chains + observed_chains = _TupleResizeObserver(expected_ready) + object.__setattr__(plan, 'chains', observed_chains) + state.activate_source(source) + + ready = state.get_ready() + + assert ready == expected_ready + assert observed_chains._held_tuples == [] + + +def test_activate_source_reports_all_active_successors_without_tuple_resize() -> None: + plan, (source,) = _manual_plan(predecessors=((), (), ()), source_indices=((0,),)) + observed_successors = _TupleResizeIndexSequence((1, 2)) + object.__setattr__(plan, 'successors', (observed_successors, (), ())) + state = ExecutionPlanState(plan) + state._active[1] = 1 + state._active[2] = 1 + + with pytest.raises(LateDependencyActivationError) as exc_info: + state.activate_source(source) + + assert 'active successor chains (1, 2)' in str(exc_info.value) + assert observed_successors._held_tuples == [] + + +def test_invalid_plan_reports_source_and_chain() -> None: + plan, _ = _manual_plan(predecessors=((), (0,)), source_indices=((1,),)) + + assert plan.find_unclosed_source() == ( + "source 'source-0.sml' activates chain 1 (_PlanNode at chain-1.sml:2:1) without predecessor chain 0" + ) + + +def test_graph_falls_back_when_plan_is_invalid( + run_validation: RunValidationFunction, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + validated = run_validation('Value = 1') + invalid_plan, _ = _manual_plan(predecessors=((), (0,)), source_indices=((1,),)) + monkeypatch.setattr(ExecutionPlan, 'from_graph', lambda _graph: invalid_plan) + + graph = compile_execution_graph(validated) + + assert graph.get_execution_plan() is None + assert "Execution plan is invalid. The graph scheduler will run: source 'source-0.sml'" in caplog.text + + +def test_activation_failure_does_not_mutate_state() -> None: + plan, (dependent_source, predecessor_source) = _manual_plan( + predecessors=((), (0,)), + source_indices=((1,), (0,)), + ) + state = ExecutionPlanState(plan) + state.activate_source(dependent_source) + before = (bytes(state._active), tuple(state._remaining), tuple(state._ready)) + + with pytest.raises(LateDependencyActivationError) as exc_info: + state.activate_source(predecessor_source) + + assert "source 'source-1.sml'" in str(exc_info.value) + assert (bytes(state._active), tuple(state._remaining), tuple(state._ready)) == before + + +def test_plan_states_are_independent(compiled_execution_graph: ExecutionGraph) -> None: + plan = compiled_execution_graph.get_execution_plan() + assert isinstance(plan, ExecutionPlan) + first = ExecutionPlanState(plan) + second = ExecutionPlanState(plan) + source = compiled_execution_graph.get_entry_point() + + first.activate_source(source) + + assert first.get_ready() + assert second.get_ready() == () + second.activate_source(source) + assert second.get_ready() + + +def test_duplicate_source_activation_is_a_noop(compiled_execution_graph: ExecutionGraph) -> None: + plan = compiled_execution_graph.get_execution_plan() + assert isinstance(plan, ExecutionPlan) + state = ExecutionPlanState(plan) + source = compiled_execution_graph.get_entry_point() + + state.activate_source(source) + first_ready = state.get_ready() + state.activate_source(source) + + assert first_ready + assert state.get_ready() == () + + +def test_unactivated_source_stays_unscheduled() -> None: + plan, (entry_source, dynamic_source) = _manual_plan(predecessors=((), ()), source_indices=((0,), (1,))) + state = ExecutionPlanState(plan) + + state.activate_source(entry_source) + assert _ready_indices(state, plan) == (0,) + state.done(plan.chains[0]) + + assert state.get_ready() == () + assert dynamic_source in plan.source_indices + + +def test_successors_become_ready_in_dynamic_activation_order() -> None: + plan, (second_successor_source, first_successor_source) = _manual_plan( + predecessors=((), (0,), (0,)), + source_indices=((0, 2), (0, 1)), + ) + state = ExecutionPlanState(plan) + state.activate_source(second_successor_source) + state.activate_source(first_successor_source) + assert _ready_indices(state, plan) == (0,) + + state.done(plan.chains[0]) + + assert _ready_indices(state, plan) == (2, 1) + + +def test_source_activation_reorders_all_ready_nodes_like_legacy_prepare() -> None: + plan, (dependent_source, existing_ready_source, new_ready_source) = _manual_plan( + predecessors=((), (0,), (), ()), + source_indices=((0, 1), (2,), (3,)), + ) + state = ExecutionPlanState(plan) + state.activate_source(dependent_source) + assert _ready_indices(state, plan) == (0,) + state.activate_source(existing_ready_source) + state.done(plan.chains[0]) + assert tuple(state._ready) == (2, 1) + + state.activate_source(new_ready_source) + + assert _ready_indices(state, plan) == (1, 2, 3) + + +def test_scheduler_matches_legacy_sorter_for_randomized_valid_activations() -> None: + rng = random.Random(0) + transitions = 0 + + while transitions < 10_000: + node_count = rng.randint(2, 9) + predecessors = tuple( + [tuple([index for index in range(node) if rng.random() < 0.25]) for node in range(node_count)] + ) + + def closure_for(target: int) -> tuple[int, ...]: + closure: set[int] = set() + + def add_with_predecessors(index: int) -> None: + for predecessor in predecessors[index]: + add_with_predecessors(predecessor) + closure.add(index) + + add_with_predecessors(target) + return tuple(sorted(closure)) + + source_indices = tuple([closure_for(rng.randrange(node_count)) for _ in range(rng.randint(2, 6))]) + source_indices += (tuple(range(node_count)),) + plan, sources = _manual_plan(predecessors, source_indices) + state = ExecutionPlanState(plan) + legacy = TopologicalSorter() + legacy_added: set[int] = set() + sources_left = list(sources) + outstanding: list[int] = [] + + while True: + if sources_left and (not outstanding or rng.random() < 0.45): + source_position = rng.randrange(len(sources_left)) + source = sources_left.pop(source_position) + indices = plan.source_indices[source] + known = legacy_added | set(indices) + for index in indices: + if legacy.already_added(index): + continue + live_predecessors = [predecessor for predecessor in predecessors[index] if predecessor in known] + legacy.add(index, *live_predecessors) + legacy_added.add(index) + legacy.prepare() + state.activate_source(source) + transitions += 1 + elif outstanding and rng.random() < 0.65: + position = rng.randrange(len(outstanding)) + index = outstanding.pop(position) + legacy.done(index) + state.done(plan.chains[index]) + transitions += 1 + else: + expected = legacy.get_ready() + actual = _ready_indices(state, plan) + assert actual == expected + outstanding.extend(expected) + transitions += 1 + if not expected and not sources_left and not outstanding: + break diff --git a/osprey_worker/src/osprey/engine/executor/tests/test_executor.py b/osprey_worker/src/osprey/engine/executor/tests/test_executor.py index dcd8ca10..9056d6c5 100644 --- a/osprey_worker/src/osprey/engine/executor/tests/test_executor.py +++ b/osprey_worker/src/osprey/engine/executor/tests/test_executor.py @@ -2,7 +2,8 @@ import json from collections.abc import Sequence from dataclasses import dataclass -from typing import Any, Type +from typing import Any, Protocol, Type, cast +from unittest.mock import MagicMock import gevent import gevent.event @@ -11,11 +12,16 @@ from osprey.engine.ast_validator.validation_context import ValidationContext from osprey.engine.conftest import ExecuteFunction, ExecuteWithResultFunction from osprey.engine.executor.execution_context import ExecutionContext, ExpectedUdfException +from osprey.engine.executor.execution_graph import ExecutionGraph from osprey.engine.language_types.post_execution_convertible import PostExecutionConvertible from osprey.engine.stdlib.udfs.import_ import Import +from osprey.engine.stdlib.udfs.json_data import JsonData +from osprey.engine.stdlib.udfs.require import Require from osprey.engine.udf.arguments import ArgumentsBase from osprey.engine.udf.base import BatchableUDFBase, UDFBase from osprey.engine.udf.registry import UDFRegistry +from osprey.worker.lib.instruments import metrics +from pytest_mock import MockFixture from result import Err, Ok, Result @@ -467,6 +473,25 @@ def test_errors_propagate_to_dependencies_only(udf_registry: UDFRegistry, execut assert data == {'A': None, 'B': 5, 'C': None, 'D': 11} +def test_plan_matches_legacy_failure_propagation( + udf_registry: UDFRegistry, execute: ExecuteFunction, monkeypatch: pytest.MonkeyPatch +) -> None: + udf_registry.register(FailingUdf) + sources = """ + Failed = FailingUdf() + Dependent = Failed + 1 + Independent = 2 + 3 + """ + + with monkeypatch.context() as legacy: + legacy.setattr(ExecutionGraph, 'get_execution_plan', lambda _graph: None) + legacy_data = execute(sources, allow_errors=True) + + planned_data = execute(sources, allow_errors=True) + + assert planned_data == legacy_data == {'Failed': None, 'Dependent': None, 'Independent': 5} + + def test_errors_propagate_to_dependencies_only_with_batching( udf_registry: UDFRegistry, batch_failing_udf: Type[BatchFailingUdf], execute: ExecuteFunction ) -> None: @@ -490,6 +515,26 @@ def test_errors_propagate_to_dependencies_only_with_batching( assert batch_failing_udf.order_called() == [[1, 2]] +def test_plan_matches_legacy_batch_groups( + batch_recording_udf: Type[BatchRecordingUdf], execute: ExecuteFunction, monkeypatch: pytest.MonkeyPatch +) -> None: + sources = """ + A = BatchRecordingUdf(id="a", routing_key="shared") + B = BatchRecordingUdf(id="b", routing_key="shared") + """ + + with monkeypatch.context() as legacy: + legacy.setattr(ExecutionGraph, 'get_execution_plan', lambda _graph: None) + legacy_data = execute(sources, async_pool=gevent.pool.Pool(2)) + legacy_calls = list(batch_recording_udf.order_called()) + batch_recording_udf.order_called().clear() + + planned_data = execute(sources, async_pool=gevent.pool.Pool(2)) + + assert planned_data == legacy_data == {'A': 'a', 'B': 'b'} + assert batch_recording_udf.order_called() == legacy_calls == [['a', 'b']] + + def test_errors_propagate_through_features( udf_registry: UDFRegistry, execute_with_result: ExecuteWithResultFunction ) -> None: @@ -792,3 +837,92 @@ def test_dependent_node_imported_after_parent_node_executed( } ) assert data == {'Child': 8} + + +def test_plan_matches_legacy_source_loading_and_dynamic_selection( + udf_registry: UDFRegistry, execute: ExecuteFunction, monkeypatch: pytest.MonkeyPatch +) -> None: + for udf_type in (Import, JsonData, Require): + udf_registry.register(udf_type) + sources = { + 'main.sml': """ + ActionName: str = JsonData(path="$.action_name", coerce_type=True) + Require(rule=f"actions/{ActionName}.sml") + """, + 'actions/a.sml': """ + Import(rules=["shared.sml"]) + A = 40 + Shared + """, + 'actions/b.sml': 'B = 99', + 'shared.sml': 'Shared = 2', + } + + with monkeypatch.context() as legacy: + legacy.setattr(ExecutionGraph, 'get_execution_plan', lambda _graph: None) + legacy_data = execute(sources, data={'action_name': 'a'}) + + planned_data = execute(sources, data={'action_name': 'a'}) + + assert planned_data == legacy_data + assert planned_data['A'] == 42 + assert 'B' not in planned_data + + +class _MetricCall(Protocol): + args: tuple[object, ...] + kwargs: dict[str, object] + + +def _sole_call_for_metric(mock: MagicMock, metric_name: str) -> _MetricCall: + """`execute()` also emits an unrelated `osprey.action_health` increment; isolate the call + for the metric under test, asserting there's exactly one.""" + matching = [call for call in mock.call_args_list if call.args[0] == metric_name] + assert len(matching) == 1, f'expected exactly one {metric_name!r} call, got {matching}' + return cast(_MetricCall, matching[0]) + + +def test_metric_tags_on_sync_call_exception( + udf_registry: UDFRegistry, execute_with_result: ExecuteWithResultFunction, mocker: MockFixture +) -> None: + """A sync (execute_async=False) Call node that raises never enters the `metrics.timed` + branch, but must still emit `udf_execution` with the exact same tags the old eager-build + code would have produced.""" + udf_registry.register(FailingUdf) + base_tags = ('scope:test',) + mocker.patch('osprey.engine.executor.executor._get_metric_tags', side_effect=lambda *_args: list(base_tags)) + mocker.patch('osprey.worker.lib.instruments.metrics.increment') + mocker.patch('osprey.worker.lib.instruments.metrics.timed') + + result = execute_with_result('A = FailingUdf()') + + assert len(result.error_infos) == 1 + metrics.timed.assert_not_called() + call = _sole_call_for_metric(metrics.increment, 'udf_execution') + assert cast(list[str], call.kwargs['tags']) == [ + *base_tags, + 'udf:FailingUdf', + 'exc_name:ValueError', + 'result:unexpected_failure', + ] + + +def test_metric_tags_on_async_udf_success( + blocking_udf: Type[BlockingUdf], execute_with_result: ExecuteWithResultFunction, mocker: MockFixture +) -> None: + """An async (execute_async=True) UDF that succeeds consumes tags in both the `metrics.timed` + call and the success `metrics.increment` -- both must be built from the identical base list.""" + base_tags = ('scope:test',) + mocker.patch('osprey.engine.executor.executor._get_metric_tags', side_effect=lambda *_args: list(base_tags)) + mocker.patch('osprey.worker.lib.instruments.metrics.increment') + mocker.patch('osprey.worker.lib.instruments.metrics.timed') + + result = execute_with_result('A = BlockingUdf(block=False, id="a")', async_pool=None) + + assert result.error_infos == [] + timed_call = _sole_call_for_metric(metrics.timed, 'udf_execution_duration') + assert timed_call.kwargs['sample_rate'] == 0.01 + timed_tags = cast(list[str], timed_call.kwargs['tags']) + assert timed_tags == [*base_tags, 'udf:BlockingUdf'] + + inc_call = _sole_call_for_metric(metrics.increment, 'udf_execution') + assert cast(list[str], inc_call.kwargs['tags']) == [*timed_tags, 'exc_name:none', 'result:success'] diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/json_utils.py b/osprey_worker/src/osprey/engine/stdlib/udfs/json_utils.py index d5f2ee3a..2ebbebac 100644 --- a/osprey_worker/src/osprey/engine/stdlib/udfs/json_utils.py +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/json_utils.py @@ -1,16 +1,79 @@ from typing import Any, TypeVar -from jsonpath_rw import JSONPath, parse +from jsonpath_rw import Child, Fields, JSONPath, Root, parse +from jsonpath_rw import jsonpath as jsonpath_module from osprey.engine.executor.execution_context import ExpectedUdfException from osprey.engine.udf.arguments import ConstExpr from osprey.engine.udf.rvalue_type_checker import RValueTypeChecker from osprey.engine.udf.type_helpers import to_display_str +_MISSING: Any = object() +"""Sentinel for "the path matched nothing", which a matched value of `None` must not collide with.""" + + +class FieldChainPath(Child): # type: ignore[misc] # jsonpath_rw does not provide type information. + """A plain field chain with a direct lookup path.""" + + __slots__ = ('_keys',) + + def __init__(self, keys: tuple[str, ...], expr: Child): + super().__init__(expr.left, expr.right) + self._keys = keys + + def get_first(self, data: object) -> Any: + """The value at this path, or `_MISSING` when jsonpath_rw would report no match. + + `Fields.get_field_datum` subscripts (never `.get()`, so a present-but-null field is a + match holding `None`) and treats exactly these exceptions as "no match". + """ + value: Any = data + for key in self._keys: + try: + value = value[key] + except (TypeError, KeyError, AttributeError): + return _MISSING + return value + + def __repr__(self) -> str: + return f'{type(self).__name__}({self._keys!r})' + + +def _compile_field_chain(expr: JSONPath) -> tuple[str, ...] | None: + """The key tuple for a `$.a.b.c` expression, or `None` when the shape isn't provably a plain chain. + + Only a `Root` base and `Child(, Fields())` links qualify — + wildcards, comma selectors, indices, slices, descendants, filters and `@`/`parent` bases all + keep jsonpath_rw. Exact type checks (not `isinstance`) so a future jsonpath_rw node subclassing + one of these can't silently inherit the fast path. + """ + # The auto-id feature rewrites what a *missing* field resolves to, so only compile while it is + # off. Osprey never enables it; a monkeypatch after rule compilation would not be picked up. + if jsonpath_module.auto_id_field is not None: + return None + + keys: list[str] = [] + node = expr + while type(node) is Child: + right = node.right + if type(right) is not Fields or len(right.fields) != 1: + return None + field = right.fields[0] + if not isinstance(field, str) or field == '*': + return None + keys.append(field) + node = node.left + + if type(node) is not Root or not keys: + return None + + keys.reverse() + return tuple(keys) + def parse_path(path: ConstExpr[str]) -> JSONPath: with path.attribute_errors(): try: - return parse(path.value) + expr = parse(path.value) except Exception as e: # There is a bug in jsonpath_rw that is throwing an error while trying to generate an error. # that's cool, but we can catch it and transmute it to something a little more relevant for now. @@ -18,6 +81,12 @@ def parse_path(path: ConstExpr[str]) -> JSONPath: raise Exception('invalid json-path supplied') raise + keys = _compile_field_chain(expr) + if keys is None: + return expr + assert isinstance(expr, Child) + return FieldChainPath(keys, expr) + class MissingJsonPath(Exception): def __init__(self, path: str) -> None: @@ -50,11 +119,14 @@ def get_from_data( coerce_type: bool, rvalue_type_checker: RValueTypeChecker, ) -> Any: - matches = expr.find(data) - - if matches: - value = matches[0].value + value: Any + if type(expr) is FieldChainPath: + value = expr.get_first(data) else: + matches = expr.find(data) + value = matches[0].value if matches else _MISSING + + if value is _MISSING: # If we can return None, do that if rvalue_type_checker.check(None): return None diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_json_data.py b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_json_data.py index c8f3529e..6e7ef2fe 100644 --- a/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_json_data.py +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_json_data.py @@ -59,6 +59,30 @@ def test_execute_value_not_present(execute: ExecuteFunction, execute_with_result assert result.extracted_features['Foo'] is None +def test_execute_nested_path(execute: ExecuteFunction, execute_with_result: ExecuteWithResultFunction) -> None: + source = "Foo: str = JsonData(path='$.a.b.c', required=False)" + + assert execute(source, data={'a': {'b': {'c': 'hello'}}}) == {'Foo': 'hello'} + # Missing leaf, missing intermediate, and a non-dict intermediate all read as absent. + for data in ({'a': {'b': {}}}, {'a': {}}, {'a': None}, {'a': ['b']}, {}): + result = execute_with_result(source, data=data) + assert not result.error_infos, (data, result.error_infos) + assert result.extracted_features['Foo'] is None + + +@pytest.mark.parametrize( + ('path', 'data'), + [ + ('$.a.*', {'a': {'k': 'x'}}), + ('$.a[0]', {'a': ['x']}), + ('$..c', {'b': {'c': 'x'}}), + ], +) +def test_execute_non_simple_path_uses_jsonpath_rw(path: str, data: dict[str, Any], execute: ExecuteFunction) -> None: + """Wildcard/index/descendant paths can't be specialized, so they keep the jsonpath_rw walk.""" + assert execute(f"Foo: str = JsonData(path='{path}')", data=data) == {'Foo': 'x'} + + def test_execute_value_present_but_null(execute: ExecuteFunction) -> None: data = execute("Foo: str = JsonData(path='$.foo', required=False)", data={'foo': None}) diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_json_utils.py b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_json_utils.py new file mode 100644 index 00000000..4524f109 --- /dev/null +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_json_utils.py @@ -0,0 +1,120 @@ +from typing import Any, Dict, List, Tuple + +import pytest +from jsonpath_rw import JSONPath, parse +from osprey.engine.stdlib.udfs.json_utils import ( + _MISSING, + FieldChainPath, + _compile_field_chain, + parse_path, +) +from osprey.engine.udf.arguments import ConstExpr + + +def _parse_path(path: str) -> Any: + return parse_path(ConstExpr.for_default('path', path)) + + +@pytest.mark.parametrize( + ('path', 'expected_keys'), + [ + ('$.foo', ('foo',)), + ('$.foo.bar', ('foo', 'bar')), + ('$.http.headers.Cookie', ('http', 'headers', 'Cookie')), + ], +) +def test_compiles_simple_field_chains(path: str, expected_keys: Tuple[str, ...]) -> None: + compiled = _parse_path(path) + assert isinstance(compiled, FieldChainPath) + assert _compile_field_chain(parse(path)) == expected_keys + # str() is what MissingJsonPath/InvalidJsonType report, so it must survive specialization. + assert str(compiled) == path + + +@pytest.mark.parametrize( + 'path', + [ + '$.a.*', # wildcard field + '$.a[*]', # slice + '$.a[0]', # index + '$.a[1:2]', # bounded slice + '$.a[0].b', # index in the middle of a chain + '$..b', # descendants + '$.a,b', # comma selector + '$', # bare root + 'foo.bar', # no root + '@.a', # `this` base + ], +) +def test_falls_back_to_jsonpath_rw(path: str) -> None: + assert not isinstance(_parse_path(path), FieldChainPath) + assert _compile_field_chain(parse(path)) is None + + +_EQUIVALENCE_CASES: List[Tuple[str, Dict[str, object]]] = [ + # Simple chains against documents that hit, miss, or hold None. + ('$.a', {'a': 1}), + ('$.a', {'a': None}), + ('$.a', {}), + ('$.a', {'b': 1}), + ('$.a.b', {'a': {'b': 'x'}}), + ('$.a.b', {'a': {'b': None}}), + ('$.a.b', {'a': {}}), + ('$.a.b', {}), + ('$.a.b.c', {'a': {'b': {'c': 0}}}), + ('$.a.b.c', {'a': {'b': {}}}), + # Missing intermediate node. + ('$.a.b.c', {'a': None}), + ('$.a.b.c', {'a': {'b': None}}), + # Non-dict intermediates: list, str, int, and an object with no __getitem__. + ('$.a.b', {'a': [1, 2]}), + ('$.a.b', {'a': 'string'}), + ('$.a.b', {'a': 7}), + ('$.a.b', {'a': object()}), + # A key whose value shadows a falsy match. + ('$.a.b', {'a': {'b': False}}), + ('$.a.b', {'a': {'b': []}}), + # Non-simple shapes still route through jsonpath_rw. + ('$.a.*', {'a': {'b': 'x'}}), + ('$.a[0]', {'a': ['x', 'y']}), + ('$.a[0].b', {'a': [{'b': 'x'}]}), + ('$..b', {'a': {'b': 'x'}}), +] + + +@pytest.mark.parametrize(('path', 'document'), _EQUIVALENCE_CASES) +def test_fast_path_matches_jsonpath_rw(path: str, document: Dict[str, object]) -> None: + """The specialized accessor and jsonpath_rw must agree on both "did it match" and the value.""" + matches = parse(path).find(document) + expected: Any = matches[0].value if matches else _MISSING + + compiled = _parse_path(path) + if isinstance(compiled, FieldChainPath): + actual = compiled.get_first(document) + else: + found = compiled.find(document) + actual = found[0].value if found else _MISSING + + assert actual is expected or actual == expected + + +@pytest.mark.parametrize(('path', 'document'), _EQUIVALENCE_CASES) +def test_find_delegates_to_jsonpath_rw(path: str, document: Dict[str, object]) -> None: + """`find()` stays available and identical for any caller holding a parsed path.""" + assert _parse_path(path).find(document) == parse(path).find(document) + + +@pytest.mark.parametrize('path', ['$.a', '$.missing']) +def test_update_delegates_to_jsonpath_rw(path: str) -> None: + assert _parse_path(path).update({'a': 1}, 2) == parse(path).update({'a': 1}, 2) + + +def test_repr_and_equality() -> None: + compiled = _parse_path('$.a.b') + original = parse('$.a.b') + assert isinstance(compiled, JSONPath) + assert repr(compiled) == "FieldChainPath(('a', 'b'))" + assert compiled == original + assert original == compiled + assert compiled == _parse_path('$.a.b') + assert compiled != _parse_path('$.a.c') diff --git a/osprey_worker/src/osprey/engine/udf/arguments.py b/osprey_worker/src/osprey/engine/udf/arguments.py index 0499f4fd..8d6bc500 100644 --- a/osprey_worker/src/osprey/engine/udf/arguments.py +++ b/osprey_worker/src/osprey/engine/udf/arguments.py @@ -20,6 +20,7 @@ # will be collected in it. The arguments will be typechecked against the value type, e.g. extra_args: dict[str, int] # would require all extra arguments to be ints EXTRA_ARGS_ATTR = 'extra_arguments' +ARGUMENT_METADATA_CACHE_SIZE = 1024 class ConstExpr(Generic[T]): @@ -227,7 +228,7 @@ def __eq__(self, o: object) -> bool: def __hash__(self) -> int: # raises TypeError('unhashable type: yada yada yada') if an argument is not hashable assert self._resolved, 'arguments are not comparable until resolved' - return hash(tuple(v for _, v in sorted(self._arguments.items()))) + return hash(tuple([value for _, value in sorted(self._arguments.items())])) def get_call_node(self) -> grammar.Call: return self._call_node @@ -265,7 +266,7 @@ def traverse_mro(klass: Any) -> None: return list(ordered_mro) @classmethod - @lru_cache(1) + @lru_cache(maxsize=ARGUMENT_METADATA_CACHE_SIZE) def items(cls) -> dict[str, type]: fields: dict[str, type] = {} @@ -315,12 +316,12 @@ def is_generic(cls) -> bool: return cls.get_generic_param() is not None @classmethod - @lru_cache(1) + @lru_cache(maxsize=ARGUMENT_METADATA_CACHE_SIZE) def is_extra_arguments_allowed(cls) -> bool: return EXTRA_ARGS_ATTR in cls.items() @classmethod - @lru_cache(1) + @lru_cache(maxsize=ARGUMENT_METADATA_CACHE_SIZE) def get_extra_arguments_values_type(cls) -> type: """returns the type allowed by unexpected kwargs""" assert cls.is_extra_arguments_allowed(), 'check if is_extra_arguments_allowed() first' @@ -332,6 +333,7 @@ def get_extra_arguments_values_type(cls) -> type: return val_type @classmethod + @lru_cache(maxsize=ARGUMENT_METADATA_CACHE_SIZE) def kwarg_can_be_none(cls, name: str) -> bool: """Whether or not the kwarg can accept None as an input. diff --git a/osprey_worker/src/osprey/engine/udf/tests/test_arguments.py b/osprey_worker/src/osprey/engine/udf/tests/test_arguments.py index 877da299..9b607c55 100644 --- a/osprey_worker/src/osprey/engine/udf/tests/test_arguments.py +++ b/osprey_worker/src/osprey/engine/udf/tests/test_arguments.py @@ -1,3 +1,8 @@ +import builtins +from types import GeneratorType +from typing import cast + +import osprey.engine.udf.arguments as arguments_module from osprey.engine.udf.arguments import ArgumentsBase, ConstExpr StrConstExpr = ConstExpr[str] # This being inside the below function is causing mypy to crash @@ -15,6 +20,20 @@ class Arguments(ArgumentsBase): assert items['bar'] is StrConstExpr +def test_arguments_items_cache_is_scoped_to_each_subclass() -> None: + class FirstArguments(ArgumentsBase): + first: str + + class SecondArguments(ArgumentsBase): + second: int + + first_items = FirstArguments.items() + SecondArguments.items() + + assert FirstArguments.items() is first_items + assert FirstArguments.items.cache_info().maxsize is not None + + def test_arguments_can_be_none() -> None: class Arguments(ArgumentsBase): optional: str | None @@ -30,3 +49,66 @@ class Arguments(ArgumentsBase): assert Arguments.kwarg_can_be_none('obj') assert not Arguments.kwarg_can_be_none('string') assert not Arguments.kwarg_can_be_none('integer') + + +def test_arguments_can_be_none_is_cached() -> None: + class Arguments(ArgumentsBase): + optional: str | None + string: str + + hits_before = Arguments.kwarg_can_be_none.cache_info().hits + + assert Arguments.kwarg_can_be_none('optional') is True + assert Arguments.kwarg_can_be_none('string') is False + + # Repeat the same calls; both should now be served from the cache. + assert Arguments.kwarg_can_be_none('optional') is True + assert Arguments.kwarg_can_be_none('string') is False + + assert Arguments.kwarg_can_be_none.cache_info().hits == hits_before + 2 + assert Arguments.kwarg_can_be_none.cache_info().maxsize is not None + + +def test_extra_argument_metadata_is_cached() -> None: + class Arguments(ArgumentsBase): + extra_arguments: dict[str, int] + + allowed_hits_before = Arguments.is_extra_arguments_allowed.cache_info().hits + value_type_hits_before = Arguments.get_extra_arguments_values_type.cache_info().hits + + assert Arguments.is_extra_arguments_allowed() is True + assert Arguments.get_extra_arguments_values_type() is int + assert Arguments.is_extra_arguments_allowed() is True + assert Arguments.get_extra_arguments_values_type() is int + + assert Arguments.is_extra_arguments_allowed.cache_info().hits >= allowed_hits_before + 2 + assert Arguments.get_extra_arguments_values_type.cache_info().hits == value_type_hits_before + 1 + assert Arguments.is_extra_arguments_allowed.cache_info().maxsize is not None + assert Arguments.get_extra_arguments_values_type.cache_info().maxsize is not None + + +def test_arguments_hash_does_not_pass_a_generator_to_tuple(monkeypatch) -> None: + tuple_inputs: list[object] = [] + + def recording_tuple(values): + tuple_inputs.append(values) + return builtins.tuple(values) + + class Arguments(ArgumentsBase): + value: object + + class CallNode: + def argument_dict(self) -> dict[str, object]: + return {} + + monkeypatch.setattr(arguments_module, 'tuple', recording_tuple, raising=False) + + arguments = Arguments( + call_node=cast(arguments_module.grammar.Call, CallNode()), + arguments={'value': object()}, + resolved=True, + ) + + hash(arguments) + + assert all(not isinstance(values, GeneratorType) for values in tuple_inputs) diff --git a/osprey_worker/src/osprey/engine/utils/tests/test_types.py b/osprey_worker/src/osprey/engine/utils/tests/test_types.py new file mode 100644 index 00000000..612c8ebe --- /dev/null +++ b/osprey_worker/src/osprey/engine/utils/tests/test_types.py @@ -0,0 +1,91 @@ +import builtins +from collections.abc import Sequence +from types import GeneratorType +from typing import Protocol, cast + +import osprey.engine.utils.types as types_module + + +class _StatefulValue(Protocol): + value: object + + def __getstate__(self) -> Sequence[object]: + pass + + def __setstate__(self, state: Sequence[object]) -> None: + pass + + +def _record_tuple_inputs(monkeypatch) -> list[object]: + tuple_inputs: list[object] = [] + + def recording_tuple(values): + tuple_inputs.append(values) + return builtins.tuple(values) + + monkeypatch.setattr(types_module, 'tuple', recording_tuple, raising=False) + return tuple_inputs + + +def test_add_state_functions_does_not_pass_a_generator_to_tuple(monkeypatch) -> None: + tuple_inputs = _record_tuple_inputs(monkeypatch) + + cls_dict: dict[str, object] = {} + + types_module._add_state_functions(cls_dict, ('value',)) + + assert tuple_inputs + assert all(not isinstance(values, GeneratorType) for values in tuple_inputs) + + +def test_slots_getstate_does_not_pass_a_generator_to_tuple(monkeypatch) -> None: + tuple_inputs = _record_tuple_inputs(monkeypatch) + cls_dict: dict[str, object] = {} + types_module._add_state_functions(cls_dict, ('value',)) + getstate = cls_dict['__getstate__'] + + class Instance: + value = object() + + assert callable(getstate) + instance = Instance() + state = getstate(instance) + assert state == (instance.value,) + + restored = Instance() + restored.value = object() + setstate = cls_dict['__setstate__'] + assert callable(setstate) + setstate(restored, state) + assert restored.value is instance.value + assert tuple_inputs + assert all(not isinstance(values, GeneratorType) for values in tuple_inputs) + + +def test_add_slots_does_not_pass_a_generator_to_tuple(monkeypatch) -> None: + tuple_inputs = _record_tuple_inputs(monkeypatch) + + class Field: + name = 'value' + + class Fields: + def __iter__(self): + yield Field() + + @types_module.dataclasses.dataclass + class Value: + value: object + + monkeypatch.setattr(types_module.dataclasses, 'fields', lambda _: Fields()) + + slotted_value_cls = types_module.add_slots(Value) + instance = cast(_StatefulValue, slotted_value_cls(value=object())) + state = instance.__getstate__() + restored = cast(_StatefulValue, slotted_value_cls(value=object())) + restored.__setstate__(state) + + assert 'value' in getattr(slotted_value_cls, '__slots__') + assert state == (instance.value,) + assert restored.value is instance.value + assert tuple_inputs + assert all(not isinstance(values, GeneratorType) for values in tuple_inputs) diff --git a/osprey_worker/src/osprey/engine/utils/types.py b/osprey_worker/src/osprey/engine/utils/types.py index 6a6944d2..65324d73 100644 --- a/osprey_worker/src/osprey/engine/utils/types.py +++ b/osprey_worker/src/osprey/engine/utils/types.py @@ -26,11 +26,11 @@ def _add_state_functions(cls_dict: dict[str, object], field_names: Sequence[str] # https://github.com/python-attrs/attrs/blob/33b61316f8fd97d78374818e5ecc21068cf69ae3/src/attr/_make.py#L660-L689 # __weakref__ is not writable. - state_attr_names = tuple(name for name in field_names if name != '__weakref__') + state_attr_names = tuple([name for name in field_names if name != '__weakref__']) def slots_getstate(self: object) -> Sequence[object]: """Automatically created by slotted_dataclass.""" - return tuple(getattr(self, name) for name in state_attr_names) + return tuple([getattr(self, name) for name in state_attr_names]) def slots_setstate(self: object, state: Sequence[object]) -> None: """Automatically created by slotted_dataclass.""" @@ -59,7 +59,7 @@ def add_slots(cls: TypeT) -> TypeT: # Create a new dict for our new class. cls_dict = dict(cls.__dict__) - field_names = tuple(f.name for f in dataclasses.fields(cls)) + field_names = tuple([field.name for field in dataclasses.fields(cls)]) # Some extra things to make sure deepcopy and weakref work. _add_state_functions(cls_dict, field_names)