feat!(extstore): add new concurrency controls for external storage.

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.

changelog/breaking-changes/napping-narwhal.md

@@ -0,0 +1 @@
+**Experimental**: Removed the `max_workflow_task_external_storage_concurrency` `Worker` argument. Use `ExternalStorage.concurrency`.

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

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

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:

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,

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,

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)

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,
         )