Preflight all externalized payload sizes before uploads (#305)

Fixes #298

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: c95a41ba-3f18-427d-8a6c-a7438e7ec782

Author: andystaples

CHANGELOG.md

@@ -9,6 +9,11 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
+- Large-payload externalization now checks every eligible payload's maximum
+size before uploading or replacing any protobuf values, including carryover
+events and entity payloads, for both synchronous and asynchronous calls.
+Oversized messages remain unchanged and cause no uploads; their size error
+now takes precedence over hypothetical earlier storage failures.
 - Orchestrator generators are now closed at the end of each dispatch, including
 while waiting for durable work, releasing resources held by async bridges.
 `finally` blocks and context-manager exits can run before orchestration completion

azure-functions-durable/CHANGELOG.md

@@ -21,6 +21,9 @@ CHANGED
 
 FIXED
 
+- With the corresponding core `durabletask` SDK fix, oversized messages fail
+large-payload validation before any uploads or payload replacements. Size
+errors now take precedence over hypothetical earlier storage failures.
 - With the corresponding core `durabletask` SDK fix, orchestrator generators
 are closed at the end of each dispatch, including while waiting for durable work.
 See [cleanup semantics](../docs/supported-patterns.md#cleanup-in-orchestrator-finally-blocks)

durabletask-azuremanaged/CHANGELOG.md

@@ -9,6 +9,9 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
+- With the corresponding core `durabletask` SDK fix, oversized messages fail
+large-payload validation before any uploads or payload replacements. Size
+errors now take precedence over hypothetical earlier storage failures.
 - With the corresponding core `durabletask` SDK fix, orchestrator generators
 are closed at the end of each dispatch, including while waiting for durable work.
 See [cleanup semantics](../docs/supported-patterns.md#cleanup-in-orchestrator-finally-blocks)

durabletask/payload/helpers.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any
 
 from google.protobuf import message as proto_message
@@ -38,11 +39,23 @@ def externalize_payloads(
 
     Any ``StringValue`` whose UTF-8 byte length exceeds
     ``store.options.threshold_bytes`` is uploaded and its value replaced
-    with the token returned by the store.
+    with the token returned by the store. Empty values and known tokens
+    are skipped.
+
+    All eligible payloads are snapshotted and checked against
+    ``store.options.max_stored_payload_bytes`` before any upload or mutation.
+    A size violation therefore raises ``ValueError`` before any hypothetical
+    earlier upload failure. Upload failures can still leave earlier uploads
+    and token replacements in place; this is not a storage transaction.
     """
-    threshold = store.options.threshold_bytes
-    max_bytes = store.options.max_stored_payload_bytes
-    _walk_and_externalize(msg, store, threshold, max_bytes, instance_id)
+    payloads = _collect_payloads_to_externalize(msg, store)
+    for field_name, sv, payload_bytes in payloads:
+        token = store.upload(payload_bytes, instance_id=instance_id)
+        sv.value = token
+        logger.debug(
+            "Externalized %d-byte payload in field '%s' -> %s",
+            len(payload_bytes), field_name, token,
+        )
 
 
 def deexternalize_payloads(
@@ -66,10 +79,19 @@ async def externalize_payloads_async(
     *,
     instance_id: str | None = None,
 ) -> None:
-    """Async version of :func:`externalize_payloads`."""
-    threshold = store.options.threshold_bytes
-    max_bytes = store.options.max_stored_payload_bytes
-    await _walk_and_externalize_async(msg, store, threshold, max_bytes, instance_id)
+    """Async version of :func:`externalize_payloads`, with the same preflight.
+
+    Uploads remain sequential; upload errors and cancellation propagate without
+    rolling back earlier uploads or token replacements.
+    """
+    payloads = _collect_payloads_to_externalize(msg, store)
+    for field_name, sv, payload_bytes in payloads:
+        token = await store.upload_async(payload_bytes, instance_id=instance_id)
+        sv.value = token
+        logger.debug(
+            "Externalized %d-byte payload in field '%s' -> %s",
+            len(payload_bytes), field_name, token,
+        )
 
 
 async def deexternalize_payloads_async(
@@ -81,7 +103,7 @@ async def deexternalize_payloads_async(
 
 
 # ------------------------------------------------------------------
-# Internal recursive walkers – sync
+# Internal preflight and recursive walkers – sync
 # ------------------------------------------------------------------
 
 def _is_map_field(fd: Any) -> bool:
@@ -90,88 +112,50 @@ def _is_map_field(fd: Any) -> bool:
     return mt is not None and fd.is_repeated and mt.GetOptions().map_entry
 
 
-def _walk_and_externalize(
+def _iter_payload_fields(
     msg: proto_message.Message,
-    store: PayloadStore,
-    threshold: int,
-    max_bytes: int,
-    instance_id: str | None,
-) -> None:
+) -> Iterator[tuple[str, wrappers_pb2.StringValue]]:
     for fd in msg.DESCRIPTOR.fields:
         if fd.message_type is None:
             continue
 
         if _is_map_field(fd):
-            # Map fields: iterate values. ScalarMap values are not
-            # messages and will be skipped by the isinstance check.
-            for map_value in getattr(msg, fd.name).values():
-                if isinstance(map_value, proto_message.Message):
-                    if isinstance(map_value, wrappers_pb2.StringValue):
-                        _try_externalize_field(
-                            fd.name, map_value, store,
-                            threshold, max_bytes, instance_id,
-                        )
-                    else:
-                        _walk_and_externalize(
-                            map_value, store, threshold, max_bytes, instance_id
-                        )
+            values = getattr(msg, fd.name).values()
         elif fd.is_repeated:
-            value = getattr(msg, fd.name)
-            for item in value:
-                if isinstance(item, proto_message.Message):
-                    if isinstance(item, wrappers_pb2.StringValue):
-                        _try_externalize_field(
-                            fd.name, item, store,
-                            threshold, max_bytes, instance_id,
-                        )
-                    else:
-                        _walk_and_externalize(
-                            item, store, threshold, max_bytes, instance_id
-                        )
+            values = getattr(msg, fd.name)
         else:
-            # Singular message field — only recurse if actually set
             if not msg.HasField(fd.name):
                 continue
-            value = getattr(msg, fd.name)
+            values = (getattr(msg, fd.name),)
+
+        for value in values:
             if isinstance(value, wrappers_pb2.StringValue):
-                _try_externalize_field(
-                    fd.name, value, store,
-                    threshold, max_bytes, instance_id,
-                )
-            else:
-                _walk_and_externalize(
-                    value, store, threshold, max_bytes, instance_id
-                )
+                yield fd.name, value
+            elif isinstance(value, proto_message.Message):
+                yield from _iter_payload_fields(value)
 
 
-def _try_externalize_field(
-    field_name: str,
-    sv: wrappers_pb2.StringValue,
+def _collect_payloads_to_externalize(
+    msg: proto_message.Message,
     store: PayloadStore,
-    threshold: int,
-    max_bytes: int,
-    instance_id: str | None,
-) -> None:
-    val = sv.value
-    if not val:
-        return
-    # Already a token – skip
-    if store.is_known_token(val):
-        return
-    payload_bytes = val.encode("utf-8")
-    if len(payload_bytes) <= threshold:
-        return
-    if len(payload_bytes) > max_bytes:
-        raise ValueError(
-            f"Payload size {len(payload_bytes)} bytes exceeds the maximum "
-            f"allowed size of {max_bytes} bytes."
-        )
-    token = store.upload(payload_bytes, instance_id=instance_id)
-    sv.value = token
-    logger.debug(
-        "Externalized %d-byte payload in field '%s' -> %s",
-        len(payload_bytes), field_name, token,
-    )
+) -> list[tuple[str, wrappers_pb2.StringValue, bytes]]:
+    threshold = store.options.threshold_bytes
+    max_bytes = store.options.max_stored_payload_bytes
+    payloads: list[tuple[str, wrappers_pb2.StringValue, bytes]] = []
+    for field_name, sv in _iter_payload_fields(msg):
+        val = sv.value
+        if not val or store.is_known_token(val):
+            continue
+        payload_bytes = val.encode("utf-8")
+        if len(payload_bytes) <= threshold:
+            continue
+        if len(payload_bytes) > max_bytes:
+            raise ValueError(
+                f"Payload size {len(payload_bytes)} bytes exceeds the maximum "
+                f"allowed size of {max_bytes} bytes."
+            )
+        payloads.append((field_name, sv, payload_bytes))
+    return payloads
 
 
 def _walk_and_deexternalize(
@@ -223,87 +207,6 @@ def _try_deexternalize_field(
 # Internal recursive walkers – async
 # ------------------------------------------------------------------
 
-async def _walk_and_externalize_async(
-    msg: proto_message.Message,
-    store: PayloadStore,
-    threshold: int,
-    max_bytes: int,
-    instance_id: str | None,
-) -> None:
-    for fd in msg.DESCRIPTOR.fields:
-        if fd.message_type is None:
-            continue
-
-        if _is_map_field(fd):
-            for map_value in getattr(msg, fd.name).values():
-                if isinstance(map_value, proto_message.Message):
-                    if isinstance(map_value, wrappers_pb2.StringValue):
-                        await _try_externalize_field_async(
-                            fd.name, map_value, store,
-                            threshold, max_bytes, instance_id,
-                        )
-                    else:
-                        await _walk_and_externalize_async(
-                            map_value, store, threshold, max_bytes, instance_id,
-                        )
-        elif fd.is_repeated:
-            value = getattr(msg, fd.name)
-            for item in value:
-                if isinstance(item, proto_message.Message):
-                    if isinstance(item, wrappers_pb2.StringValue):
-                        await _try_externalize_field_async(
-                            fd.name, item, store,
-                            threshold, max_bytes, instance_id,
-                        )
-                    else:
-                        await _walk_and_externalize_async(
-                            item, store, threshold, max_bytes, instance_id,
-                        )
-        else:
-            if not msg.HasField(fd.name):
-                continue
-            value = getattr(msg, fd.name)
-            if isinstance(value, wrappers_pb2.StringValue):
-                await _try_externalize_field_async(
-                    fd.name, value, store,
-                    threshold, max_bytes, instance_id,
-                )
-            else:
-                await _walk_and_externalize_async(
-                    value, store, threshold, max_bytes, instance_id,
-                )
-
-
-async def _try_externalize_field_async(
-    field_name: str,
-    sv: wrappers_pb2.StringValue,
-    store: PayloadStore,
-    threshold: int,
-    max_bytes: int,
-    instance_id: str | None,
-) -> None:
-    val = sv.value
-    if not val:
-        return
-    # Already a token – skip
-    if store.is_known_token(val):
-        return
-    payload_bytes = val.encode("utf-8")
-    if len(payload_bytes) <= threshold:
-        return
-    if len(payload_bytes) > max_bytes:
-        raise ValueError(
-            f"Payload size {len(payload_bytes)} bytes exceeds the maximum "
-            f"allowed size of {max_bytes} bytes."
-        )
-    token = await store.upload_async(payload_bytes, instance_id=instance_id)
-    sv.value = token
-    logger.debug(
-        "Externalized %d-byte payload in field '%s' -> %s",
-        len(payload_bytes), field_name, token,
-    )
-
-
 async def _walk_and_deexternalize_async(
     msg: proto_message.Message,
     store: PayloadStore,