diff --git a/changelog/added/juggling-jackalope.md b/changelog/added/juggling-jackalope.md new file mode 100644 index 0000000000..2d692c5a09 --- /dev/null +++ b/changelog/added/juggling-jackalope.md @@ -0,0 +1 @@ +**Experimental**: New external storage concurrency controls. `ExternalStorage.concurrency` sets `max_driver_operations` (for one client or worker) and `max_operations_per_message` (for one message). Each client and each worker gets its own `max_driver_operations` limit. Drivers must use `context.limiter` for store and retrieve operations. diff --git a/changelog/breaking-changes/napping-narwhal.md b/changelog/breaking-changes/napping-narwhal.md new file mode 100644 index 0000000000..0e0272472a --- /dev/null +++ b/changelog/breaking-changes/napping-narwhal.md @@ -0,0 +1 @@ +**Experimental**: Removed the `max_workflow_task_external_storage_concurrency` `Worker` argument. Use `ExternalStorage.concurrency`. diff --git a/scripts/gen_payload_visitor.py b/scripts/gen_payload_visitor.py index 7fab85dcda..664e3432a6 100644 --- a/scripts/gen_payload_visitor.py +++ b/scripts/gen_payload_visitor.py @@ -164,7 +164,7 @@ def __init__( *, skip_search_attributes: bool = False, skip_headers: bool = False, - concurrency_limit: int = 1, + concurrency_limit: int | None = 1, ): \"\"\"Creates a new payload visitor. @@ -173,9 +173,10 @@ def __init__( skip_headers: If True, headers are not visited. concurrency_limit: Maximum number of payload visits that may run concurrently during a single call to visit(). Defaults to 1 - (sequential). + (sequential). None runs every visit at once, for callers + that do their own limiting. \"\"\" - if concurrency_limit < 1: + if concurrency_limit is not None and concurrency_limit < 1: raise ValueError("concurrency_limit must be positive") self.skip_search_attributes = skip_search_attributes self.skip_headers = skip_headers diff --git a/temporalio/bridge/_visitor.py b/temporalio/bridge/_visitor.py index d8fa05edf3..3f93072daf 100644 --- a/temporalio/bridge/_visitor.py +++ b/temporalio/bridge/_visitor.py @@ -22,7 +22,7 @@ def __init__( *, skip_search_attributes: bool = False, skip_headers: bool = False, - concurrency_limit: int = 1, + concurrency_limit: int | None = 1, ): """Creates a new payload visitor. @@ -31,9 +31,10 @@ def __init__( skip_headers: If True, headers are not visited. concurrency_limit: Maximum number of payload visits that may run concurrently during a single call to visit(). Defaults to 1 - (sequential). + (sequential). None runs every visit at once, for callers + that do their own limiting. """ - if concurrency_limit < 1: + if concurrency_limit is not None and concurrency_limit < 1: raise ValueError("concurrency_limit must be positive") self.skip_search_attributes = skip_search_attributes self.skip_headers = skip_headers diff --git a/temporalio/bridge/_visitor_functions.py b/temporalio/bridge/_visitor_functions.py index da8c67ae2b..8382d246b2 100644 --- a/temporalio/bridge/_visitor_functions.py +++ b/temporalio/bridge/_visitor_functions.py @@ -36,16 +36,33 @@ async def drain_since(self, _checkpoint: int) -> None: return None +class _UnboundedSemaphore: + """A semaphore that never blocks. Used when visits are not limited.""" + + async def acquire(self) -> None: + """Returns immediately.""" + + def release(self) -> None: + """Does nothing.""" + + class BoundedVisitorFunctions(VisitorFunctions): """Wraps VisitorFunctions to cap concurrent payload visits via a semaphore. After the full traversal, call drain() to await all in-flight tasks. """ - def __init__(self, inner: VisitorFunctions, concurrency_limit: int) -> None: - """Create a bounded wrapper around the given visitor functions.""" + def __init__(self, inner: VisitorFunctions, concurrency_limit: int | None) -> None: + """Create a bounded wrapper around the given visitor functions. + + ``None`` runs visits without a limit, still tracking tasks to drain. + """ self._inner = inner - self._sem = asyncio.Semaphore(concurrency_limit) + self._sem: asyncio.Semaphore | _UnboundedSemaphore = ( + _UnboundedSemaphore() + if concurrency_limit is None + else asyncio.Semaphore(concurrency_limit) + ) self._tasks: list[asyncio.Task[None]] = [] async def visit_payload(self, payload: Payload) -> None: diff --git a/temporalio/bridge/worker.py b/temporalio/bridge/worker.py index 4b7f55d092..3a4f2e6250 100644 --- a/temporalio/bridge/worker.py +++ b/temporalio/bridge/worker.py @@ -5,7 +5,8 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable, Sequence +import contextlib +from collections.abc import Awaitable, Callable, Iterator, Sequence from dataclasses import dataclass from typing import ( TypeAlias, @@ -304,23 +305,40 @@ async def visit_payloads(self, payloads: PayloadSequence) -> None: payloads.extend(new_payloads) +@contextlib.contextmanager +def _external_storage_message( + data_converter: temporalio.converter.DataConverter, +) -> Iterator[tuple[temporalio.converter._extstore.StorageOperationMetrics, int]]: + """Yields one message's external storage metrics and payload visit limit.""" + metrics = temporalio.converter._extstore.StorageOperationMetrics() + storage = data_converter.external_storage + if storage is None: + with metrics.track(): + yield metrics, 1 + return + visit_limit = storage.concurrency.max_operations_per_message + with ( + metrics.track(), + temporalio.converter._extstore.message_scope(storage), + ): + yield metrics, visit_limit + + async def decode_activation( activation: temporalio.bridge.proto.workflow_activation.WorkflowActivation, data_converter: temporalio.converter.DataConverter, decode_headers: bool, - storage_concurrency_limit: int, ) -> temporalio.converter._extstore.StorageOperationMetrics: """Decode all payloads in the activation. Returns: Metrics from any external storage retrieval operations that occurred. """ - metrics = temporalio.converter._extstore.StorageOperationMetrics() - with metrics.track(): + with _external_storage_message(data_converter) as (metrics, visit_limit): await CommandAwarePayloadVisitor( skip_search_attributes=True, skip_headers=not decode_headers, - concurrency_limit=storage_concurrency_limit, + concurrency_limit=visit_limit, ).visit( _Visitor(data_converter._external_retrieve_payload_sequence), activation ) @@ -337,7 +355,6 @@ async def encode_completion( completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion, data_converter: temporalio.converter.DataConverter, encode_headers: bool, - storage_concurrency_limit: int, ) -> temporalio.converter._extstore.StorageOperationMetrics: """Encode all payloads in the completion. @@ -352,12 +369,11 @@ async def encode_completion( completion, ) - metrics = temporalio.converter._extstore.StorageOperationMetrics() - with metrics.track(): + with _external_storage_message(data_converter) as (metrics, visit_limit): await CommandAwarePayloadVisitor( skip_search_attributes=True, skip_headers=not encode_headers, - concurrency_limit=storage_concurrency_limit, + concurrency_limit=visit_limit, ).visit( _Visitor(data_converter._external_store_payload_sequence), completion, diff --git a/temporalio/client/_activity.py b/temporalio/client/_activity.py index df7cd5c283..18573c28ad 100644 --- a/temporalio/client/_activity.py +++ b/temporalio/client/_activity.py @@ -943,8 +943,8 @@ def with_context(self, context: SerializationContext) -> Self: converter that makes use of this context then you can use this method to supply matching context data to the data converter used to serialize and encode the outbound payloads. """ - data_converter = self._client.data_converter.with_context(context) - if data_converter is self._client.data_converter: + data_converter = self._client._data_converter.with_context(context) + if data_converter is self._client._data_converter: return self cls = type(self) if cls.__init__ is not AsyncActivityHandle.__init__: @@ -981,7 +981,7 @@ def __init__( @functools.cached_property def _data_converter(self) -> temporalio.converter.DataConverter: - return self._client.data_converter.with_context( + return self._client._data_converter.with_context( ActivitySerializationContext( namespace=self._client.namespace, activity_id=self._id, diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index 5a99c0b0fe..28d4b01a29 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import dataclasses from collections.abc import ( Awaitable, Callable, @@ -25,6 +26,7 @@ import temporalio.api.workflowservice.v1 import temporalio.common import temporalio.converter +import temporalio.converter._extstore import temporalio.runtime import temporalio.service import temporalio.workflow @@ -293,6 +295,20 @@ def __init__( def _init_from_config(self, config: ClientConfig): self._config = config + # The driver operation limit is state of this client, so it rides on a + # private copy of the converter rather than the config or the public + # data_converter. Clients built from this config get their own. + converter = config["data_converter"] + self._data_converter = ( + dataclasses.replace( + converter, + _driver_operation_semaphore=temporalio.converter._extstore.new_driver_operation_semaphore( + converter.external_storage + ), + ) + if converter.external_storage is not None + else converter + ) # Iterate over interceptors in reverse building the impl self._impl: OutboundInterceptor = _ClientImpl(self) diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index 0c1e67b85e..3da01eb159 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -235,7 +235,7 @@ async def _build_signal_with_start_workflow_execution_request( self, input: StartWorkflowInput ) -> temporalio.api.workflowservice.v1.SignalWithStartWorkflowExecutionRequest: assert input.start_signal - data_converter = self._client.data_converter._with_contexts( + data_converter = self._client._data_converter._with_contexts( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=input.id, @@ -274,7 +274,7 @@ async def _populate_start_workflow_execution_request( ), input: StartWorkflowInput | UpdateWithStartStartWorkflowInput, ) -> None: - data_converter = self._client.data_converter._with_contexts( + data_converter = self._client._data_converter._with_contexts( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=input.id, @@ -366,7 +366,7 @@ async def describe_workflow( timeout=input.rpc_timeout, ), namespace=self._client.namespace, - converter=self._client.data_converter.with_context( + converter=self._client._data_converter.with_context( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=input.id, @@ -400,7 +400,7 @@ async def count_workflows( ) async def query_workflow(self, input: QueryWorkflowInput) -> Any: - data_converter = self._client.data_converter._with_contexts( + data_converter = self._client._data_converter._with_contexts( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=input.id, @@ -466,7 +466,7 @@ async def query_workflow(self, input: QueryWorkflowInput) -> Any: return results[0] async def signal_workflow(self, input: SignalWorkflowInput) -> None: - data_converter = self._client.data_converter._with_contexts( + data_converter = self._client._data_converter._with_contexts( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=input.id, @@ -504,7 +504,7 @@ async def signal_workflow(self, input: SignalWorkflowInput) -> None: ) async def terminate_workflow(self, input: TerminateWorkflowInput) -> None: - data_converter = self._client.data_converter._with_contexts( + data_converter = self._client._data_converter._with_contexts( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=input.id, @@ -579,7 +579,7 @@ async def _build_start_activity_execution_request( self, input: StartActivityInput ) -> temporalio.api.workflowservice.v1.StartActivityExecutionRequest: """Build StartActivityExecutionRequest from input.""" - data_converter = self._client.data_converter._with_contexts( + data_converter = self._client._data_converter._with_contexts( ActivitySerializationContext( namespace=self._client.namespace, activity_id=input.id, @@ -805,7 +805,7 @@ async def describe_activity( return ActivityExecutionDescription._from_resp( resp=resp, namespace=self._client.namespace, - data_converter=self._client.data_converter.with_context( + data_converter=self._client._data_converter.with_context( ActivitySerializationContext( namespace=self._client.namespace, activity_id=resp.info.activity_id, @@ -894,7 +894,7 @@ async def _build_update_workflow_execution_request( input: StartWorkflowUpdateInput | UpdateWithStartUpdateWorkflowInput, workflow_id: str, ) -> temporalio.api.workflowservice.v1.UpdateWorkflowExecutionRequest: - data_converter = self._client.data_converter._with_contexts( + data_converter = self._client._data_converter._with_contexts( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=workflow_id, @@ -1126,7 +1126,7 @@ async def heartbeat_async_activity( self, input: HeartbeatAsyncActivityInput ) -> None: data_converter = ( - input.data_converter_override or self._client.data_converter + input.data_converter_override or self._client._data_converter )._with_store_context(self._get_async_activity_store_context(input.id_or_token)) details = ( None @@ -1183,7 +1183,7 @@ async def heartbeat_async_activity( async def complete_async_activity(self, input: CompleteAsyncActivityInput) -> None: data_converter = ( - input.data_converter_override or self._client.data_converter + input.data_converter_override or self._client._data_converter )._with_store_context(self._get_async_activity_store_context(input.id_or_token)) result = ( None @@ -1219,7 +1219,7 @@ async def complete_async_activity(self, input: CompleteAsyncActivityInput) -> No async def fail_async_activity(self, input: FailAsyncActivityInput) -> None: data_converter = ( - input.data_converter_override or self._client.data_converter + input.data_converter_override or self._client._data_converter )._with_store_context(self._get_async_activity_store_context(input.id_or_token)) failure = temporalio.api.failure.v1.Failure() @@ -1262,7 +1262,7 @@ async def report_cancellation_async_activity( self, input: ReportCancellationAsyncActivityInput ) -> None: data_converter = ( - input.data_converter_override or self._client.data_converter + input.data_converter_override or self._client._data_converter )._with_store_context(self._get_async_activity_store_context(input.id_or_token)) details = ( None @@ -1338,7 +1338,7 @@ async def create_schedule(self, input: CreateScheduleInput) -> ScheduleHandle: initial_patch=initial_patch, identity=self._client.identity, request_id=str(uuid.uuid4()), - memo=await self._client.data_converter._encode_memo(input.memo) + memo=await self._client._data_converter._encode_memo(input.memo) if input.memo else None, ) @@ -1410,7 +1410,7 @@ async def describe_schedule( metadata=input.rpc_metadata, timeout=input.rpc_timeout, ), - self._client.data_converter, + self._client._data_converter, ) async def pause_schedule(self, input: PauseScheduleInput) -> None: @@ -1485,7 +1485,7 @@ async def update_schedule(self, input: UpdateScheduleInput) -> None: metadata=input.rpc_metadata, timeout=input.rpc_timeout, ), - self._client.data_converter, + self._client._data_converter, ) ) ) @@ -1563,7 +1563,7 @@ async def start_nexus_operation( service=input.service, operation=input.operation, ) - data_converter = self._client.data_converter.with_context(nexus_context) + data_converter = self._client._data_converter.with_context(nexus_context) req = temporalio.api.workflowservice.v1.StartNexusOperationExecutionRequest( namespace=self._client.namespace, identity=self._client.identity, @@ -1651,7 +1651,7 @@ async def describe_nexus_operation( metadata=input.rpc_metadata, timeout=input.rpc_timeout, ) - data_converter = self._client.data_converter.with_context( + data_converter = self._client._data_converter.with_context( temporalio.converter.NexusSerializationContext( endpoint=resp.info.endpoint, service=resp.info.service, @@ -1667,7 +1667,7 @@ async def get_nexus_operation_result( self, input: GetNexusOperationResultInput ) -> Any: """Poll for nexus operation result until it's available.""" - data_converter = self._client.data_converter + data_converter = self._client._data_converter # These three are set together or not at all: a handle that started the operation has all # of them, and a handle obtained by operation ID alone has none and defaults them to "". # An empty endpoint therefore means "no operation to build a context from", not "an @@ -1783,5 +1783,5 @@ async def _apply_headers( dest, self._client.config(active_config=True)["header_codec_behavior"] == HeaderCodecBehavior.CODEC, - self._client.data_converter, + self._client._data_converter, ) diff --git a/temporalio/client/_schedule.py b/temporalio/client/_schedule.py index e0a4bcc478..d155c609bf 100644 --- a/temporalio/client/_schedule.py +++ b/temporalio/client/_schedule.py @@ -794,7 +794,7 @@ async def _to_proto( priority: temporalio.api.common.v1.Priority | None = None if self.priority: priority = self.priority._to_proto() - data_converter = client.data_converter._with_contexts( + data_converter = client._data_converter._with_contexts( WorkflowSerializationContext( namespace=client.namespace, workflow_id=self.id, @@ -858,7 +858,7 @@ async def _to_proto( client.config(active_config=True)["header_codec_behavior"] == HeaderCodecBehavior.CODEC and not self._from_raw, - client.data_converter, + client._data_converter, ) return action @@ -1572,7 +1572,7 @@ async def fetch_next_page(self, *, page_size: int | None = None) -> None: timeout=self._input.rpc_timeout, ) self._current_page = [ - ScheduleListDescription._from_proto(v, self._client.data_converter) + ScheduleListDescription._from_proto(v, self._client._data_converter) for v in resp.schedules ] self._current_page_index = 0 diff --git a/temporalio/client/_workflow.py b/temporalio/client/_workflow.py index 94bc2bb23d..ad5c9b1e8c 100644 --- a/temporalio/client/_workflow.py +++ b/temporalio/client/_workflow.py @@ -130,7 +130,7 @@ def __init__( @functools.cached_property def _data_converter(self) -> temporalio.converter.DataConverter: - return self._client.data_converter.with_context( + return self._client._data_converter.with_context( temporalio.converter.WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=self._id ) @@ -1615,7 +1615,7 @@ async def fetch_next_page(self, *, page_size: int | None = None) -> None: self._current_page = [ WorkflowExecution._from_raw_info( - v, self._client.namespace, self._client.data_converter + v, self._client.namespace, self._client._data_converter ) for v in resp.executions ] @@ -1867,7 +1867,7 @@ def __init__( @functools.cached_property def _data_converter(self) -> temporalio.converter.DataConverter: - return self._client.data_converter.with_context( + return self._client._data_converter.with_context( WorkflowSerializationContext( namespace=self._client.namespace, workflow_id=self.workflow_id, diff --git a/temporalio/contrib/aws/s3driver/_driver.py b/temporalio/contrib/aws/s3driver/_driver.py index 4bcf9de254..e86b46f0ac 100644 --- a/temporalio/contrib/aws/s3driver/_driver.py +++ b/temporalio/contrib/aws/s3driver/_driver.py @@ -199,11 +199,15 @@ async def _upload(payload: Payload) -> StorageDriverClaim: }, ) - return await _gather_with_cancellation([_upload(p) for p in payloads]) + async def _permitted_upload(payload: Payload) -> StorageDriverClaim: + async with context.limiter.permit(payload): + return await _upload(payload) + + return await _gather_with_cancellation([_permitted_upload(p) for p in payloads]) async def retrieve( self, - context: StorageDriverRetrieveContext, # noqa: ARG002 + context: StorageDriverRetrieveContext, claims: Sequence[StorageDriverClaim], ) -> list[Payload]: """Retrieves payloads from S3 for the given ``temporalio.extstore.DriverClaim`` list.""" @@ -247,4 +251,8 @@ async def _download(claim: StorageDriverClaim) -> Payload: payload.ParseFromString(payload_bytes) return payload - return await _gather_with_cancellation([_download(c) for c in claims]) + async def _permitted_download(claim: StorageDriverClaim) -> Payload: + async with context.limiter.permit(claim): + return await _download(claim) + + return await _gather_with_cancellation([_permitted_download(c) for c in claims]) diff --git a/temporalio/converter/__init__.py b/temporalio/converter/__init__.py index 324b477f25..3840e61734 100644 --- a/temporalio/converter/__init__.py +++ b/temporalio/converter/__init__.py @@ -6,9 +6,11 @@ ) from temporalio.converter._extstore import ( ExternalStorage, + ExternalStorageConcurrency, StorageDriver, StorageDriverActivityInfo, StorageDriverClaim, + StorageDriverLimiter, StorageDriverRetrieveContext, StorageDriverSelectContext, StorageDriverStoreContext, @@ -59,10 +61,12 @@ __all__ = [ "ActivitySerializationContext", "ExternalStorage", + "ExternalStorageConcurrency", "StorageDriver", "StorageDriverActivityInfo", "StorageDriverClaim", "StorageDriverRetrieveContext", + "StorageDriverLimiter", "StorageDriverSelectContext", "StorageDriverStoreContext", "StorageDriverWorkflowInfo", diff --git a/temporalio/converter/_data_converter.py b/temporalio/converter/_data_converter.py index f924ccf6c8..00d346655d 100644 --- a/temporalio/converter/_data_converter.py +++ b/temporalio/converter/_data_converter.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import dataclasses from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -18,6 +19,7 @@ _REFERENCE_MESSAGE_TYPE, ExternalStorage, StorageDriverStoreContext, + default_driver_operation_scope, ) from temporalio.converter._failure_converter import ( FailureConverter, @@ -87,6 +89,11 @@ class DataConverter(WithSerializationContext): This API is experimental. """ + _driver_operation_semaphore: asyncio.Semaphore | None = dataclasses.field( + default=None, repr=False, compare=False + ) + """The driver operation limit of the client that owns this converter.""" + default: ClassVar[DataConverter] """Singleton default data converter.""" @@ -316,9 +323,10 @@ async def _external_store_payload_sequence( """External storage store, then validate payload limits.""" stored_payloads = list(payloads) if self.external_storage: - stored_payloads = await self.external_storage._store_payload_sequence( - stored_payloads - ) + with default_driver_operation_scope(self._driver_operation_semaphore): + stored_payloads = await self.external_storage._store_payload_sequence( + stored_payloads + ) return stored_payloads async def _external_retrieve_payload_sequence( @@ -327,9 +335,12 @@ async def _external_retrieve_payload_sequence( """External storage retrieve only.""" retrieved_payloads = list(payloads) if self.external_storage: - retrieved_payloads = await self.external_storage._retrieve_payload_sequence( - retrieved_payloads - ) + with default_driver_operation_scope(self._driver_operation_semaphore): + retrieved_payloads = ( + await self.external_storage._retrieve_payload_sequence( + retrieved_payloads + ) + ) else: if any(_is_reference_payload(p) for p in retrieved_payloads): raise RuntimeError( diff --git a/temporalio/converter/_extstore.py b/temporalio/converter/_extstore.py index c300448b92..468c053826 100644 --- a/temporalio/converter/_extstore.py +++ b/temporalio/converter/_extstore.py @@ -10,9 +10,17 @@ import dataclasses import time from abc import ABC, abstractmethod -from collections.abc import Callable, Coroutine, Generator, Mapping, Sequence +from collections.abc import ( + AsyncIterator, + Callable, + Coroutine, + Generator, + Mapping, + Sequence, +) from dataclasses import dataclass from datetime import timedelta +from logging import getLogger from typing import Any, ClassVar, TypeVar from typing_extensions import Self @@ -24,6 +32,8 @@ JSONProtoPayloadConverter, ) +logger = getLogger(__name__) + _T = TypeVar("_T") _REFERENCE_ENCODING = b"json/external-storage-reference" @@ -164,6 +174,41 @@ class StorageDriverActivityInfo: """The activity type name, if available.""" +class StorageDriverLimiter(ABC): + """Limits the concurrent external storage operations a driver performs. + + Drivers must wrap each operation in :meth:`permit`. Operations performed + outside a permit are not limited. Do not take a permit inside another + permit. + + .. warning:: + This API is experimental. + """ + + @abstractmethod + def permit( + self, item: Payload | StorageDriverClaim + ) -> contextlib.AbstractAsyncContextManager[None]: + """Holds a permit for the duration of the ``async with`` block. + + ``item`` is the payload being stored or the claim being retrieved. + """ + ... + + +class _PassthroughLimiter(StorageDriverLimiter): + """A limiter that never blocks.""" + + @contextlib.asynccontextmanager + async def permit( # type: ignore[override] + self, item: Payload | StorageDriverClaim + ) -> AsyncIterator[None]: + yield + + +_PASSTHROUGH_LIMITER = _PassthroughLimiter() + + @dataclass(frozen=True) class StorageDriverStoreContext: """Context passed to :meth:`StorageDriver.store` calls. @@ -181,6 +226,9 @@ class StorageDriverStoreContext: exists the current execution context (workflow or activity) is used as the target instead.""" + limiter: StorageDriverLimiter = _PASSTHROUGH_LIMITER + """Wrap each store operation in :meth:`StorageDriverLimiter.permit`.""" + @dataclass(frozen=True) class StorageDriverSelectContext: @@ -208,6 +256,9 @@ class StorageDriverRetrieveContext: This API is experimental. """ + limiter: StorageDriverLimiter = _PASSTHROUGH_LIMITER + """Wrap each retrieve operation in :meth:`StorageDriverLimiter.permit`.""" + class StorageDriver(ABC): """Base driver for storing and retrieve payloads from external storage systems. @@ -280,6 +331,113 @@ class _StorageReference: driver_claim: StorageDriverClaim +@dataclass(frozen=True) +class ExternalStorageConcurrency: + """Configures concurrency limits for external storage operations. + + .. warning:: + This API is experimental. + """ + + max_driver_operations: int = 64 + """The maximum number of concurrent external storage operations that drivers + can execute for a single client or worker. Each client and each worker gets + its own limit of this size. Defaults to 64. + """ + + max_operations_per_message: int = 8 + """The maximum number of concurrent external storage operations that drivers + can execute for a single message (a message is input or output such as a + workflow activation, completion, or client request). This prevents one + message from monopolizing resources. Defaults to 8. + """ + + +class _SemaphoreLimiter(StorageDriverLimiter): + """Applies the storage limits to a single driver call.""" + + def __init__( + self, + per_message_semaphore: asyncio.Semaphore, + per_driver_operation_semaphore: asyncio.Semaphore, + ) -> None: + self._per_message_semaphore = per_message_semaphore + self._per_driver_operation_semaphore = per_driver_operation_semaphore + self.permit_taken = False + + @contextlib.asynccontextmanager + async def permit( # type: ignore[override] + self, item: Payload | StorageDriverClaim + ) -> AsyncIterator[None]: + self.permit_taken = True + async with self._per_message_semaphore: + async with self._per_driver_operation_semaphore: + yield + + +_current_per_message_semaphore: contextvars.ContextVar[asyncio.Semaphore | None] = ( + contextvars.ContextVar("_temporal_extstore_message_semaphore", default=None) +) +"""The limit for the message being processed.""" + + +_current_driver_operation_semaphore: contextvars.ContextVar[ + asyncio.Semaphore | None +] = contextvars.ContextVar( + "_temporal_extstore_driver_operation_semaphore", default=None +) +"""The driver operation limit of the client or worker running the operation.""" + + +@contextlib.contextmanager +def message_scope( + storage: ExternalStorage, + driver_operation_semaphore: asyncio.Semaphore | None = None, +) -> Generator[None, Any, None]: + """Gives everything in this block one shared per-message limit. + + ``driver_operation_semaphore`` is the limit of the client or worker running + the message. Pass None to leave the current one in place. + """ + message_token = _current_per_message_semaphore.set( + asyncio.Semaphore(storage.concurrency.max_operations_per_message) + ) + instance_token = ( + None + if driver_operation_semaphore is None + else _current_driver_operation_semaphore.set(driver_operation_semaphore) + ) + try: + yield + finally: + if instance_token is not None: + _current_driver_operation_semaphore.reset(instance_token) + _current_per_message_semaphore.reset(message_token) + + +def new_driver_operation_semaphore(storage: ExternalStorage) -> asyncio.Semaphore: + """Builds a driver operation limit for one client or worker.""" + return asyncio.Semaphore(storage.concurrency.max_driver_operations) + + +@contextlib.contextmanager +def default_driver_operation_scope( + driver_operation_semaphore: asyncio.Semaphore | None, +) -> Generator[None, Any, None]: + """Uses driver_operation_semaphore only when none is already in scope.""" + if ( + driver_operation_semaphore is None + or _current_driver_operation_semaphore.get() is not None + ): + yield + return + token = _current_driver_operation_semaphore.set(driver_operation_semaphore) + try: + yield + finally: + _current_driver_operation_semaphore.reset(token) + + @dataclass(frozen=True) class ExternalStorage: """Configuration for external storage behavior. @@ -314,6 +472,9 @@ class ExternalStorage: Defaults to 256 KiB. Must be greater than or equal to zero. """ + concurrency: ExternalStorageConcurrency = ExternalStorageConcurrency() + """Concurrency limits for external storage operations.""" + _driver_map: dict[str, StorageDriver] = dataclasses.field( init=False, repr=False, compare=False ) @@ -371,8 +532,47 @@ def __post_init__(self) -> None: "Each driver must have a unique name." ) driver_map[name] = driver + for name, value in ( + ("max_driver_operations", self.concurrency.max_driver_operations), + ( + "max_operations_per_message", + self.concurrency.max_operations_per_message, + ), + ): + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ValueError( + f"ExternalStorage.concurrency.{name} must be a positive integer, got {value!r}." + ) object.__setattr__(self, "_driver_map", driver_map) + def _make_limiter(self) -> _SemaphoreLimiter: + """Builds the limiter for one driver call.""" + per_driver_operation_semaphore = _current_driver_operation_semaphore.get() + if per_driver_operation_semaphore is None: + per_driver_operation_semaphore = asyncio.Semaphore( + self.concurrency.max_driver_operations + ) + per_message_semaphore = _current_per_message_semaphore.get() + if per_message_semaphore is None: + per_message_semaphore = asyncio.Semaphore( + self.concurrency.max_operations_per_message + ) + return _SemaphoreLimiter(per_message_semaphore, per_driver_operation_semaphore) + + @staticmethod + def _warn_if_limiter_unused( + limiter: _SemaphoreLimiter, driver: StorageDriver, operation: str + ) -> None: + """Warns when a driver did not use the limiter.""" + if limiter.permit_taken: + return + logger.warning( + "Storage driver '%s' performed a %s without using context.limiter. " + "Its operations are not limited.", + driver.name(), + operation, + ) + def _select_driver( self, context: StorageDriverSelectContext, payload: Payload ) -> StorageDriver | None: @@ -415,7 +615,11 @@ async def _store_payload(self, payload: Payload) -> Payload: if driver is None: return payload - claims = await driver.store(self._store_context, [payload]) + limiter = self._make_limiter() + claims = await driver.store( + dataclasses.replace(self._store_context, limiter=limiter), [payload] + ) + self._warn_if_limiter_unused(limiter, driver, "store") self._validate_claim_length(claims, expected=1, driver=driver) @@ -467,12 +671,20 @@ async def _store_payload_sequence( driver_group_list = list(driver_groups.items()) + store_limiters = [self._make_limiter() for _ in driver_group_list] all_claims = await _gather_cancel_on_error( [ - driver.store(self._store_context, [p for _, p in indexed_payloads]) - for driver, indexed_payloads in driver_group_list + driver.store( + dataclasses.replace(self._store_context, limiter=limiter), + [p for _, p in indexed_payloads], + ) + for (driver, indexed_payloads), limiter in zip( + driver_group_list, store_limiters + ) ] ) + for (driver, _), limiter in zip(driver_group_list, store_limiters): + self._warn_if_limiter_unused(limiter, driver, "store") external_count = 0 external_size = 0 @@ -534,10 +746,12 @@ async def _retrieve_payload(self, payload: Payload) -> Payload: start_time = time.monotonic() driver = self._get_driver_by_name(ref.driver_name) - context = StorageDriverRetrieveContext() + limiter = self._make_limiter() + context = StorageDriverRetrieveContext(limiter=limiter) claim = StorageDriverClaim(claim_data=dict(ref.claim_data)) stored_payloads = await driver.retrieve(context, [claim]) + self._warn_if_limiter_unused(limiter, driver, "retrieve") self._validate_payload_length(stored_payloads, expected=1, driver=driver) @@ -577,17 +791,24 @@ async def _retrieve_payload_sequence( if not driver_claims: return results - context = StorageDriverRetrieveContext() stored_by_index: dict[int, Payload] = {} driver_claim_list = list(driver_claims.items()) + retrieve_limiters = [self._make_limiter() for _ in driver_claim_list] all_stored = await _gather_cancel_on_error( [ - driver.retrieve(context, [claim for _, claim in indexed_claims]) - for driver, indexed_claims in driver_claim_list + driver.retrieve( + StorageDriverRetrieveContext(limiter=limiter), + [claim for _, claim in indexed_claims], + ) + for (driver, indexed_claims), limiter in zip( + driver_claim_list, retrieve_limiters + ) ] ) + for (driver, _), limiter in zip(driver_claim_list, retrieve_limiters): + self._warn_if_limiter_unused(limiter, driver, "retrieve") external_count = 0 external_size = 0 diff --git a/temporalio/worker/_command_aware_visitor.py b/temporalio/worker/_command_aware_visitor.py index 7c03c2cd43..ac1e60cab8 100644 --- a/temporalio/worker/_command_aware_visitor.py +++ b/temporalio/worker/_command_aware_visitor.py @@ -53,7 +53,7 @@ def __init__( *, skip_search_attributes: bool = False, skip_headers: bool = False, - concurrency_limit: int = 1, + concurrency_limit: int | None = 1, ) -> None: """Creates a new command-aware payload visitor. @@ -62,6 +62,7 @@ def __init__( skip_headers: If True, headers are not visited. concurrency_limit: Maximum number of payload visits that may run concurrently during a single call to visit(). Defaults to 1. + None runs every visit concurrently. """ super().__init__( skip_search_attributes=skip_search_attributes, diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index b3eb1a4d12..b42db8f7cf 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -283,7 +283,6 @@ def on_eviction_hook( "header_codec_behavior", HeaderCodecBehavior.NO_CODEC ) != HeaderCodecBehavior.NO_CODEC, - max_workflow_task_external_storage_concurrency=1, default_workflow_logic_flags=frozenset( self._default_workflow_logic_flags ), diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 39faf615ae..8a975ccd05 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -4,6 +4,7 @@ import asyncio import concurrent.futures +import dataclasses import hashlib import logging import sys @@ -22,6 +23,7 @@ import temporalio.bridge.worker import temporalio.client import temporalio.common +import temporalio.converter._extstore import temporalio.runtime import temporalio.service from temporalio.common import ( @@ -35,10 +37,7 @@ from ._nexus import _NexusWorker from ._plugin import Plugin from ._tuning import WorkerTuner -from ._workflow import ( - _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, - _WorkflowWorker, -) +from ._workflow import _WorkflowWorker from ._workflow_instance import ( PatchActivationInput, UnsandboxedWorkflowRunner, @@ -151,7 +150,6 @@ def __init__( maximum=5 ), disable_payload_error_limit: bool = False, - max_workflow_task_external_storage_concurrency: int = _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, ) -> None: """Create a worker to process workflows and/or activities. @@ -338,13 +336,6 @@ def __init__( and cause a task failure if the size limit is exceeded. The default is False. See https://docs.temporal.io/troubleshooting/blob-size-limit-error for more details. - max_workflow_task_external_storage_concurrency: Maximum number of - external storage payload operations (store/retrieve) that may run - concurrently within a single workflow task activation. - Defaults to 3. Adjust this value based on your workload's needs. - Please report any issues you encounter with this setting or if you - feel the default should be changed. - WARNING: This setting is experimental. """ config = WorkerConfig( @@ -392,7 +383,6 @@ def __init__( activity_task_poller_behavior=activity_task_poller_behavior, nexus_task_poller_behavior=nexus_task_poller_behavior, disable_payload_error_limit=disable_payload_error_limit, - max_workflow_task_external_storage_concurrency=max_workflow_task_external_storage_concurrency, ) plugins_from_client = cast( @@ -453,14 +443,6 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf raise ValueError( "default_versioning_behavior must be UNSPECIFIED when use_worker_versioning is False" ) - max_workflow_task_external_storage_concurrency = config.get( - "max_workflow_task_external_storage_concurrency", - _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, - ) - if max_workflow_task_external_storage_concurrency < 1: - raise ValueError( - "max_workflow_task_external_storage_concurrency must be positive" - ) if config.get("max_eager_activity_reservations_per_workflow_task", 3) < 1: raise ValueError( "max_eager_activity_reservations_per_workflow_task must be positive; " @@ -478,6 +460,21 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf # Extract storage drivers from the client's data converter _ext_storage = client_config["data_converter"].external_storage self._storage_drivers = list(_ext_storage.drivers) if _ext_storage else [] + self._driver_operation_semaphore = ( + temporalio.converter._extstore.new_driver_operation_semaphore(_ext_storage) + if _ext_storage + else None + ) + # Every worker operation draws on the worker's limit, not the client's, + # so the worker carries its own copy of the converter. + worker_data_converter = ( + dataclasses.replace( + client_config["data_converter"], + _driver_operation_semaphore=self._driver_operation_semaphore, + ) + if self._driver_operation_semaphore is not None + else client_config["data_converter"] + ) # Extract the bridge service client bridge_client = _extract_bridge_client_for_worker(config["client"]) # type: ignore[reportTypedDictNotRequiredAccess] @@ -502,7 +499,7 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf activities=activities, activity_executor=config.get("activity_executor"), shared_state_manager=config.get("shared_state_manager"), - data_converter=client_config["data_converter"], + data_converter=worker_data_converter, interceptors=interceptors, metric_meter=self._runtime.metric_meter, client=client, @@ -520,7 +517,7 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf namespace=client_config["namespace"], task_queue=config["task_queue"], # type: ignore[reportTypedDictNotRequiredAccess] service_handlers=nexus_service_handlers, - data_converter=client_config["data_converter"], + data_converter=worker_data_converter, interceptors=interceptors, metric_meter=self._runtime.metric_meter, executor=config.get("nexus_task_executor"), @@ -555,7 +552,7 @@ def check_activity(activity: str): ), workflow_runner=config["workflow_runner"], # type: ignore[reportTypedDictNotRequiredAccess] unsandboxed_workflow_runner=config["unsandboxed_workflow_runner"], # type: ignore[reportTypedDictNotRequiredAccess] - data_converter=client_config["data_converter"], + data_converter=worker_data_converter, interceptors=interceptors, workflow_failure_exception_types=config[ "workflow_failure_exception_types" @@ -572,7 +569,6 @@ def check_activity(activity: str): assert_local_activity_valid=check_activity, encode_headers=client_config["header_codec_behavior"] != HeaderCodecBehavior.NO_CODEC, - max_workflow_task_external_storage_concurrency=max_workflow_task_external_storage_concurrency, ) tuner = config.get("tuner") @@ -1027,7 +1023,6 @@ class WorkerConfig(TypedDict, total=False): activity_task_poller_behavior: PollerBehavior nexus_task_poller_behavior: PollerBehavior disable_payload_error_limit: bool - max_workflow_task_external_storage_concurrency: int def _warn_if_activity_executor_max_workers_is_inconsistent( diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index f50955c4cf..3b7832fa63 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -55,15 +55,6 @@ LOG_PROTOS = False -# Value was chosen abitrarily as a small number that allows some concurrency and prevents -# large numbers of concurrent external storage operations causing resource contention. -# This default limit is per workflow task activation and does not limit the total number -# of concurrent external storage operations across all workflow task activations. -# Advise customers to adjust based on their workload needs and to report issues with the -# value if problems are encountered. This setting is experimental. -_DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY: int = 3 - - def _set_external_storage_metrics( target: temporalio.bridge.proto.common.ExternalStorageMetrics, metrics: temporalio.converter._extstore.StorageOperationMetrics, @@ -102,7 +93,6 @@ def __init__( should_enforce_versioning_behavior: bool, assert_local_activity_valid: Callable[[str], None], encode_headers: bool, - max_workflow_task_external_storage_concurrency: int, default_workflow_logic_flags: frozenset[_WorkflowLogicFlag] | None = None, ) -> None: # Debug mode is enabled if specified or if the TEMPORAL_DEBUG env var is truthy @@ -170,9 +160,6 @@ def __init__( self._on_eviction_hook = on_eviction_hook self._disable_safe_eviction = disable_safe_eviction self._encode_headers = encode_headers - self._max_workflow_task_external_storage_concurrency = ( - max_workflow_task_external_storage_concurrency - ) self._throw_after_activation: Exception | None = None # If debug mode is enabled, disable deadlock detection @@ -386,7 +373,6 @@ async def _handle_activation( act, data_converter, decode_headers=self._encode_headers, - storage_concurrency_limit=self._max_workflow_task_external_storage_concurrency, ) if not workflow: assert init_job @@ -501,7 +487,6 @@ async def _handle_activation( completion, data_converter, encode_headers=self._encode_headers, - storage_concurrency_limit=self._max_workflow_task_external_storage_concurrency, ) except Exception as err: logger.exception( diff --git a/tests/nexus/test_link_propagation.py b/tests/nexus/test_link_propagation.py index f656f57c10..efe353a509 100644 --- a/tests/nexus/test_link_propagation.py +++ b/tests/nexus/test_link_propagation.py @@ -168,7 +168,7 @@ def _make_client_impl(workflow_service: Any) -> _ClientImpl: client.namespace = NAMESPACE client.identity = "test-identity" client.workflow_service = workflow_service - client.data_converter = temporalio.converter.DataConverter.default + client._data_converter = temporalio.converter.DataConverter.default return _ClientImpl(client) diff --git a/tests/test_extstore.py b/tests/test_extstore.py index 943fc7be2a..2e9599785d 100644 --- a/tests/test_extstore.py +++ b/tests/test_extstore.py @@ -11,6 +11,7 @@ from temporalio.converter import ( DataConverter, ExternalStorage, + ExternalStorageConcurrency, JSONPlainPayloadConverter, PayloadCodec, StorageDriver, @@ -24,6 +25,8 @@ _REFERENCE_ENCODING, StorageOperationMetrics, _StorageReference, + message_scope, + new_driver_operation_semaphore, ) from temporalio.converter._payload_converter import JSONProtoPayloadConverter from temporalio.exceptions import ApplicationError @@ -870,3 +873,285 @@ def test_storage_metrics_duration_merges_adjacent_and_nested_batches() -> None: if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +async def _settle() -> None: + """Let queued semaphore waiters reach their final state.""" + for _ in range(5): + await asyncio.sleep(0) + + +def _payload() -> Payload: + return Payload(metadata={"encoding": b"binary/plain"}, data=b"x" * 8) + + +class _Gate: + """Blocks every operation until released, recording peak concurrency.""" + + def __init__(self) -> None: + self._opened = asyncio.Event() + self._in_flight = 0 + self.peak = 0 + + def release(self) -> None: + self._opened.set() + + async def hold(self) -> None: + self._in_flight += 1 + self.peak = max(self.peak, self._in_flight) + await self._opened.wait() + self._in_flight -= 1 + + +class _PermittingDriver(StorageDriver): + """Takes a permit around every request, so peak concurrency is observable.""" + + def __init__(self, gate: _Gate) -> None: + self._gate = gate + + def name(self) -> str: + return "permitting" + + async def store( + self, context: StorageDriverStoreContext, payloads: Sequence[Payload] + ) -> list[StorageDriverClaim]: + async def one(payload: Payload) -> StorageDriverClaim: + async with context.limiter.permit(payload): + await self._gate.hold() + return StorageDriverClaim(claim_data={"id": "x"}) + + return list(await asyncio.gather(*[one(p) for p in payloads])) + + async def retrieve( + self, + context: StorageDriverRetrieveContext, + claims: Sequence[StorageDriverClaim], + ) -> list[Payload]: + async def one(claim: StorageDriverClaim) -> Payload: + async with context.limiter.permit(claim): + await self._gate.hold() + return _payload() + + return list(await asyncio.gather(*[one(c) for c in claims])) + + +class TestExternalStorageConcurrency: + """The two concurrency limits and the limiter handed to drivers.""" + + def test_defaults(self): + storage = ExternalStorage(drivers=[InMemoryTestDriver()]) + assert storage.concurrency == ExternalStorageConcurrency( + max_driver_operations=64, max_operations_per_message=8 + ) + + def test_explicit_values_kept(self): + storage = ExternalStorage( + drivers=[InMemoryTestDriver()], + concurrency=ExternalStorageConcurrency( + max_driver_operations=5, max_operations_per_message=2 + ), + ) + assert storage.concurrency.max_driver_operations == 5 + assert storage.concurrency.max_operations_per_message == 2 + + @pytest.mark.parametrize( + "concurrency", + [ + ExternalStorageConcurrency(max_driver_operations=0), + ExternalStorageConcurrency(max_operations_per_message=0), + ExternalStorageConcurrency(max_driver_operations=-1), + ], + ) + def test_rejects_limits_below_one(self, concurrency: ExternalStorageConcurrency): + with pytest.raises(ValueError, match="must be a positive integer"): + ExternalStorage(drivers=[InMemoryTestDriver()], concurrency=concurrency) + + async def test_max_operations_per_message_bounds_operations(self): + """One message may not exceed its own limit, however many sites it has.""" + gate = _Gate() + storage = ExternalStorage( + drivers=[_PermittingDriver(gate)], + payload_size_threshold=0, + concurrency=ExternalStorageConcurrency( + max_driver_operations=100, max_operations_per_message=3 + ), + ) + + with message_scope(storage): + sites = [ + asyncio.create_task( + storage._store_payload_sequence([_payload(), _payload()]) + ) + for _ in range(4) + ] + await _settle() + assert gate.peak == 3 + gate.release() + await asyncio.gather(*sites) + assert gate.peak == 3 + + async def test_each_message_gets_its_own_budget(self): + gate = _Gate() + storage = ExternalStorage( + drivers=[_PermittingDriver(gate)], + payload_size_threshold=0, + concurrency=ExternalStorageConcurrency( + max_driver_operations=100, max_operations_per_message=2 + ), + ) + + async def one_message() -> None: + with message_scope(storage): + await storage._store_payload_sequence( + [_payload(), _payload(), _payload()] + ) + + first = asyncio.create_task(one_message()) + second = asyncio.create_task(one_message()) + await _settle() + assert gate.peak == 4 # two messages at 2 each, not 2 shared + + gate.release() + await asyncio.gather(first, second) + + async def test_separate_owners_do_not_share_max_driver_operations(self): + """Two owners of one ExternalStorage each get their own limit.""" + gate = _Gate() + storage = ExternalStorage( + drivers=[_PermittingDriver(gate)], + payload_size_threshold=0, + concurrency=ExternalStorageConcurrency( + max_driver_operations=2, max_operations_per_message=10 + ), + ) + + async def one_owner() -> None: + with message_scope(storage, new_driver_operation_semaphore(storage)): + await storage._store_payload_sequence( + [_payload(), _payload(), _payload()] + ) + + first = asyncio.create_task(one_owner()) + second = asyncio.create_task(one_owner()) + await _settle() + assert gate.peak == 4, "each owner allows 2, so two owners allow 4" + + gate.release() + await asyncio.gather(first, second) + + async def test_max_driver_operations_shared_across_messages(self): + """One owner's driver operation limit spans every message it runs.""" + gate = _Gate() + storage = ExternalStorage( + drivers=[_PermittingDriver(gate)], + payload_size_threshold=0, + concurrency=ExternalStorageConcurrency( + max_driver_operations=3, max_operations_per_message=10 + ), + ) + + driver_operations = new_driver_operation_semaphore(storage) + + async def one_message() -> None: + with message_scope(storage, driver_operations): + await storage._store_payload_sequence( + [_payload(), _payload(), _payload()] + ) + + first = asyncio.create_task(one_message()) + second = asyncio.create_task(one_message()) + await _settle() + assert gate.peak == 3 + + gate.release() + await asyncio.gather(first, second) + + async def test_store_context_copies_share_max_driver_operations(self): + """Workers bind a store context per message by copying the instance. + Those copies must still draw from the original's limit.""" + gate = _Gate() + storage = ExternalStorage( + drivers=[_PermittingDriver(gate)], + payload_size_threshold=0, + concurrency=ExternalStorageConcurrency( + max_driver_operations=3, max_operations_per_message=10 + ), + ) + + driver_operations = new_driver_operation_semaphore(storage) + + async def one_message() -> None: + copy = storage._with_store_context(StorageDriverStoreContext(target=None)) + with message_scope(copy, driver_operations): + await copy._store_payload_sequence([_payload(), _payload(), _payload()]) + + first = asyncio.create_task(one_message()) + second = asyncio.create_task(one_message()) + await _settle() + assert gate.peak == 3 + + gate.release() + await asyncio.gather(first, second) + + async def test_data_converter_store_context_copies_share_max_driver_operations( + self, + ): + gate = _Gate() + data_converter = DataConverter( + external_storage=ExternalStorage( + drivers=[_PermittingDriver(gate)], + payload_size_threshold=0, + concurrency=ExternalStorageConcurrency( + max_driver_operations=3, max_operations_per_message=10 + ), + ) + ) + + assert data_converter.external_storage is not None + driver_operations = new_driver_operation_semaphore( + data_converter.external_storage + ) + + async def one_message() -> None: + copy = data_converter._with_store_context( + StorageDriverStoreContext(target=None) + ) + assert copy.external_storage is not None + with message_scope(copy.external_storage, driver_operations): + await copy._external_store_payload_sequence( + [_payload(), _payload(), _payload()] + ) + + first = asyncio.create_task(one_message()) + second = asyncio.create_task(one_message()) + await _settle() + assert gate.peak == 3 + + gate.release() + await asyncio.gather(first, second) + + async def test_warns_when_driver_takes_no_permit( + self, caplog: pytest.LogCaptureFixture + ): + storage = ExternalStorage( + drivers=[InMemoryTestDriver()], payload_size_threshold=0 + ) + with caplog.at_level("WARNING", logger="temporalio.converter"): + await storage._store_payload_sequence([_payload(), _payload()]) + assert any( + "without using context.limiter" in r.message for r in caplog.records + ), caplog.records + + async def test_no_warning_when_driver_takes_a_permit( + self, caplog: pytest.LogCaptureFixture + ): + gate = _Gate() + gate.release() + storage = ExternalStorage( + drivers=[_PermittingDriver(gate)], payload_size_threshold=0 + ) + with caplog.at_level("WARNING", logger="temporalio.converter"): + await storage._store_payload_sequence([_payload(), _payload()]) + assert not [ + r for r in caplog.records if "without using context.limiter" in r.message + ] diff --git a/tests/worker/test_extstore.py b/tests/worker/test_extstore.py index e8ef8edb2c..2130ada859 100644 --- a/tests/worker/test_extstore.py +++ b/tests/worker/test_extstore.py @@ -1,12 +1,15 @@ +import asyncio import contextlib import dataclasses import uuid -from collections.abc import Iterator, Sequence +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass from datetime import timedelta from unittest import mock +import nexusrpc import pytest +from nexusrpc.handler import StartOperationContext, service_handler, sync_operation import temporalio import temporalio.bridge.client @@ -20,6 +23,7 @@ from temporalio.common import RetryPolicy from temporalio.converter import ( ExternalStorage, + ExternalStorageConcurrency, StorageDriver, StorageDriverActivityInfo, StorageDriverClaim, @@ -29,9 +33,10 @@ ) from temporalio.exceptions import ActivityError, ApplicationError from temporalio.testing._workflow import WorkflowEnvironment -from temporalio.worker import Replayer +from temporalio.worker import Replayer, UnsandboxedWorkflowRunner from tests.helpers import assert_task_fail_eventually, new_worker -from tests.test_extstore import InMemoryTestDriver +from tests.helpers.nexus import make_nexus_endpoint_name +from tests.test_extstore import InMemoryTestDriver, _settle @dataclass(frozen=True) @@ -1345,3 +1350,354 @@ async def test_extstore_continue_as_new_result_stored_under_current_run( assert isinstance(result_ctx.target, StorageDriverWorkflowInfo) assert result_ctx.target.run_id is not None assert result_ctx.target.run_id == last_run_id + + +class _HeldStoreDriver(InMemoryTestDriver): + """Holds every store inside its permit until released, recording how many + stores have been requested and how many are in flight at once.""" + + def __init__(self) -> None: + super().__init__() + self.requested = 0 + self.in_flight = 0 + self.peak = 0 + self._changed = asyncio.Event() + self._released = asyncio.Event() + + def release(self) -> None: + self._released.set() + + async def wait_until(self, condition: Callable[[], bool]) -> None: + # The timeout only stops a broken limiter from hanging the test. + async def wait() -> None: + while not condition(): + self._changed.clear() + await self._changed.wait() + + await asyncio.wait_for(wait(), timeout=10) + + async def store( + self, + context: StorageDriverStoreContext, + payloads: Sequence[Payload], + ) -> list[StorageDriverClaim]: + async def one(payload: Payload) -> StorageDriverClaim: + self.requested += 1 + self._changed.set() + async with context.limiter.permit(payload): + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + self._changed.set() + try: + await self._released.wait() + claims = await InMemoryTestDriver.store(self, context, [payload]) + return claims[0] + finally: + self.in_flight -= 1 + + return list(await asyncio.gather(*[one(p) for p in payloads])) + + async def retrieve( + self, + context: StorageDriverRetrieveContext, + claims: Sequence[StorageDriverClaim], + ) -> list[Payload]: + async def one(claim: StorageDriverClaim) -> Payload: + async with context.limiter.permit(claim): + payloads = await InMemoryTestDriver.retrieve(self, context, [claim]) + return payloads[0] + + return list(await asyncio.gather(*[one(c) for c in claims])) + + +@activity.defn +async def large_input_length(data: str) -> int: + return len(data) + + +@workflow.defn +class ManyLargeActivityInputsWorkflow: + """Schedules every activity in one workflow task, so a single completion + carries all of their large inputs.""" + + @workflow.run + async def run(self, activity_count: int) -> list[int]: + return list( + await asyncio.gather( + *[ + workflow.execute_activity( + large_input_length, + "x" * 2048, + schedule_to_close_timeout=timedelta(seconds=10), + ) + for _ in range(activity_count) + ] + ) + ) + + +async def _connect_with_held_driver( + env: WorkflowEnvironment, + driver: _HeldStoreDriver, + concurrency: ExternalStorageConcurrency, +) -> Client: + return await env.connect_client( + data_converter=dataclasses.replace( + temporalio.converter.default(), + external_storage=ExternalStorage( + drivers=[driver], + payload_size_threshold=1024, + concurrency=concurrency, + ), + ), + ) + + +async def test_extstore_workflow_completion_respects_max_operations_per_message( + env: WorkflowEnvironment, +) -> None: + driver = _HeldStoreDriver() + client = await _connect_with_held_driver( + env, + driver, + ExternalStorageConcurrency( + max_driver_operations=100, max_operations_per_message=2 + ), + ) + + async with new_worker( + client, ManyLargeActivityInputsWorkflow, activities=[large_input_length] + ) as worker: + handle = await client.start_workflow( + ManyLargeActivityInputsWorkflow.run, + 5, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + # The payload visitor also caps visits per message, so the other stores + # wait before reaching the driver. + await driver.wait_until(lambda: driver.in_flight == 2) + await _settle() + assert driver.in_flight == 2 + + driver.release() + assert await handle.result() == [2048] * 5 + + assert driver.peak == 2 + + +async def test_extstore_max_driver_operations_shared_across_workflows( + env: WorkflowEnvironment, +) -> None: + """Each completion stores 2 payloads and the worker allows 3 at once, so the + second workflow's completion can only start one store until the first + finishes.""" + driver = _HeldStoreDriver() + client = await _connect_with_held_driver( + env, + driver, + ExternalStorageConcurrency( + max_driver_operations=3, max_operations_per_message=10 + ), + ) + + async with new_worker( + client, ManyLargeActivityInputsWorkflow, activities=[large_input_length] + ) as worker: + handles = [ + await client.start_workflow( + ManyLargeActivityInputsWorkflow.run, + 2, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + for _ in range(2) + ] + await driver.wait_until(lambda: driver.requested == 4) + await _settle() + assert driver.in_flight == 3 + + driver.release() + for handle in handles: + assert await handle.result() == [2048, 2048] + + assert driver.peak == 3 + + +_LARGE_DATA = "x" * 2048 + + +@activity.defn +async def large_result() -> str: + return _LARGE_DATA + + +@workflow.defn +class LargeActivityResultWorkflow: + @workflow.run + async def run(self) -> int: + result = await workflow.execute_activity( + large_result, schedule_to_close_timeout=timedelta(seconds=10) + ) + return len(result) + + +@nexusrpc.service +class LargeResultNexusService: + large_result_op: nexusrpc.Operation[None, str] + + +@service_handler(service=LargeResultNexusService) +class LargeResultNexusServiceHandler: + @sync_operation + async def large_result_op(self, _ctx: StartOperationContext, _input: None) -> str: + return _LARGE_DATA + + +@workflow.defn +class LargeNexusResultWorkflow: + @workflow.run + async def run(self, task_queue: str) -> int: + nexus_client = workflow.create_nexus_client( + service=LargeResultNexusService, + endpoint=make_nexus_endpoint_name(task_queue), + ) + result = await nexus_client.execute_operation( + LargeResultNexusService.large_result_op, None + ) + return len(result) + + +def _start_unrun_workflow_with_large_input(client: Client) -> asyncio.Task: + """Starts a workflow no worker polls for, so the only store is the client's + store of its input.""" + return asyncio.create_task( + client.start_workflow( + "UnregisteredWorkflow", + _LARGE_DATA, + id=f"workflow-{uuid.uuid4()}", + task_queue=f"tq-{uuid.uuid4()}", + ) + ) + + +@pytest.mark.parametrize("active_config", [False, True]) +async def test_extstore_client_rebuilt_from_config_gets_its_own_driver_operation_limit( + env: WorkflowEnvironment, active_config: bool +) -> None: + driver = _HeldStoreDriver() + client = await _connect_with_held_driver( + env, driver, ExternalStorageConcurrency(max_driver_operations=1) + ) + rebuilt_client = Client(**client.config(active_config=active_config)) + + starts = [ + _start_unrun_workflow_with_large_input(client), + _start_unrun_workflow_with_large_input(client), + _start_unrun_workflow_with_large_input(rebuilt_client), + ] + await driver.wait_until(lambda: driver.requested == 3) + await _settle() + # One store from each client. The original's second store waits. + assert driver.in_flight == 2 + + driver.release() + await asyncio.gather(*starts) + assert driver.peak == 2 + + +async def test_extstore_client_data_converter_is_the_one_passed_in( + env: WorkflowEnvironment, +) -> None: + """The client's driver operation limit stays internal, so the public + converter holds no state tied to the client's event loop.""" + data_converter = dataclasses.replace( + temporalio.converter.default(), + external_storage=ExternalStorage(drivers=[InMemoryTestDriver()]), + ) + client = await env.connect_client(data_converter=data_converter) + assert client.data_converter is data_converter + + +async def test_extstore_worker_driver_operation_limit_is_separate_from_client( + env: WorkflowEnvironment, +) -> None: + driver = _HeldStoreDriver() + client = await _connect_with_held_driver( + env, driver, ExternalStorageConcurrency(max_driver_operations=1) + ) + + async with new_worker( + client, LargeActivityResultWorkflow, activities=[large_result] + ) as worker: + handle = await client.start_workflow( + LargeActivityResultWorkflow.run, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + # The worker stores the activity result. + await driver.wait_until(lambda: driver.requested == 1) + client_start = _start_unrun_workflow_with_large_input(client) + await driver.wait_until(lambda: driver.requested == 2) + await _settle() + assert driver.in_flight == 2 + + driver.release() + await client_start + assert await handle.result() == len(_LARGE_DATA) + + assert driver.peak == 2 + + +@pytest.mark.requires_local_server +async def test_extstore_workflow_activity_and_nexus_workers_share_driver_operation_limit( + env: WorkflowEnvironment, +) -> None: + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + driver = _HeldStoreDriver() + client = await _connect_with_held_driver( + env, driver, ExternalStorageConcurrency(max_driver_operations=1) + ) + + async with new_worker( + client, + ManyLargeActivityInputsWorkflow, + LargeActivityResultWorkflow, + LargeNexusResultWorkflow, + activities=[large_input_length, large_result], + nexus_service_handlers=[LargeResultNexusServiceHandler()], + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + await env.create_nexus_endpoint( + make_nexus_endpoint_name(worker.task_queue), worker.task_queue + ) + # Each workflow causes exactly one store: a workflow completion, an + # activity result and a Nexus operation result. + workflow_completion = await client.start_workflow( + ManyLargeActivityInputsWorkflow.run, + 1, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + activity_result = await client.start_workflow( + LargeActivityResultWorkflow.run, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + nexus_result = await client.start_workflow( + LargeNexusResultWorkflow.run, + worker.task_queue, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await driver.wait_until(lambda: driver.requested == 3) + await _settle() + assert driver.in_flight == 1 + + driver.release() + assert await workflow_completion.result() == [len(_LARGE_DATA)] + assert await activity_result.result() == len(_LARGE_DATA) + assert await nexus_result.result() == len(_LARGE_DATA) + + assert driver.peak == 1 diff --git a/tests/worker/test_visitor.py b/tests/worker/test_visitor.py index 362288f348..f6e3488b71 100644 --- a/tests/worker/test_visitor.py +++ b/tests/worker/test_visitor.py @@ -550,9 +550,7 @@ async def test_bridge_encoding(): payload_codec=SimpleCodec(), ) - await temporalio.bridge.worker.encode_completion( - comp, data_converter, True, storage_concurrency_limit=1 - ) + await temporalio.bridge.worker.encode_completion(comp, data_converter, True) cmd = comp.successful.commands[0] sa = cmd.schedule_activity