💥 Payload limit configuration and validation

temporalio/bridge/worker.py

@@ -315,7 +315,6 @@ async def encode_completion(
     encode_headers: bool,
 ) -> None:
     """Encode all payloads in the completion."""
-    if data_converter._encode_payload_has_effect:
-        await CommandAwarePayloadVisitor(
-            skip_search_attributes=True, skip_headers=not encode_headers
-        ).visit(_Visitor(data_converter._encode_payload_sequence), completion)
+    await CommandAwarePayloadVisitor(
+        skip_search_attributes=True, skip_headers=not encode_headers
+    ).visit(_Visitor(data_converter._encode_payload_sequence), completion)

temporalio/client.py

@@ -6837,7 +6837,7 @@ async def _apply_headers(
 ) -> None:
     if source is None:
         return
-    if encode_headers and data_converter._encode_payload_has_effect:
+    if encode_headers:
         for payload in source.values():
             payload.CopyFrom(await data_converter._encode_payload(payload))
     temporalio.common._apply_headers(source, dest)

temporalio/converter.py

@@ -1207,6 +1207,27 @@ def __init__(self) -> None:
         super().__init__(encode_common_attributes=True)
 
 
+@dataclass(frozen=True)
+class PayloadLimitsConfig:
+    """Configuration for when payload sizes exceed limits."""
+
+    memo_size_warning: int = 2 * 1024
+    """The limit (in bytes) at which a memo size warning is logged."""
+
+    payload_size_warning: int = 512 * 1024
+    """The limit (in bytes) at which a payload size warning is logged."""
+
+
+class PayloadSizeWarning(RuntimeWarning):
+    """The size of payloads is above the warning limit."""
+
+
+@dataclass
+class _PayloadErrorLimits:
+    memo_size_error: int
+    payload_size_error: int
+
+
 @dataclass(frozen=True)
 class DataConverter(WithSerializationContext):
     """Data converter for converting and encoding payloads to/from Python values.
@@ -1230,9 +1251,16 @@ class DataConverter(WithSerializationContext):
     failure_converter: FailureConverter = dataclasses.field(init=False)
     """Failure converter created from the :py:attr:`failure_converter_class`."""
 
+    payload_limits: PayloadLimitsConfig = PayloadLimitsConfig()
+    """Settings for payload size limits."""
+
     default: ClassVar[DataConverter]
     """Singleton default data converter."""
 
+    _memo_size_error_limit: int = 0
+
+    _payload_size_error_limit: int = 0
+
     def __post_init__(self) -> None:  # noqa: D105
         object.__setattr__(self, "payload_converter", self.payload_converter_class())
         object.__setattr__(self, "failure_converter", self.failure_converter_class())
@@ -1334,6 +1362,13 @@ def with_context(self, context: SerializationContext) -> Self:
         object.__setattr__(cloned, "failure_converter", failure_converter)
         return cloned
 
+    def _with_payload_error_limits(self, options: _PayloadErrorLimits) -> DataConverter:
+        return dataclasses.replace(
+            self,
+            _memo_size_error_limit=options.memo_size_error,
+            _payload_size_error_limit=options.payload_size_error,
+        )
+
     async def _decode_memo(
         self,
         source: temporalio.api.common.v1.Memo,
@@ -1367,35 +1402,43 @@ async def _encode_memo(
     async def _encode_memo_existing(
         self, source: Mapping[str, Any], memo: temporalio.api.common.v1.Memo
     ):
+        payloads = []
         for k, v in source.items():
             payload = v
             if not isinstance(v, temporalio.api.common.v1.Payload):
                 payload = (await self.encode([v]))[0]
             memo.fields[k].CopyFrom(payload)
+            payloads.append(payload)
+        # Memos have their field payloads validated all together in one unit
+        DataConverter._validate_limits(
+            payloads,
+            self._memo_size_error_limit,
+            "[TMPRL1103] Attempted to upload memo with size that exceeded the error limit.",
+            self.payload_limits.memo_size_warning,
+            "[TMPRL1103] Attempted to upload memo with size that exceeded the warning limit.",
+        )
 
     async def _encode_payload(
         self, payload: temporalio.api.common.v1.Payload
     ) -> temporalio.api.common.v1.Payload:
         if self.payload_codec:
             payload = (await self.payload_codec.encode([payload]))[0]
+        self._validate_payload_limits([payload])
         return payload
 
     async def _encode_payloads(self, payloads: temporalio.api.common.v1.Payloads):
         if self.payload_codec:
             await self.payload_codec.encode_wrapper(payloads)
+        self._validate_payload_limits(payloads.payloads)
 
     async def _encode_payload_sequence(
         self, payloads: Sequence[temporalio.api.common.v1.Payload]
     ) -> list[temporalio.api.common.v1.Payload]:
-        if not self.payload_codec:
-            return list(payloads)
-        return await self.payload_codec.encode(payloads)
-
-    # Temporary shortcircuit detection while the _encode_* methods may no-op if
-    # a payload codec is not configured. Remove once those paths have more to them.
-    @property
-    def _encode_payload_has_effect(self) -> bool:
-        return self.payload_codec is not None
+        encoded_payloads = list(payloads)
+        if self.payload_codec:
+            encoded_payloads = await self.payload_codec.encode(encoded_payloads)
+        self._validate_payload_limits(encoded_payloads)
+        return encoded_payloads
 
     async def _decode_payload(
         self, payload: temporalio.api.common.v1.Payload
@@ -1452,6 +1495,40 @@ async def _apply_to_failure_payloads(
         if failure.HasField("cause"):
             await DataConverter._apply_to_failure_payloads(failure.cause, cb)
 
+    def _validate_payload_limits(
+        self,
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+    ):
+        DataConverter._validate_limits(
+            payloads,
+            self._payload_size_error_limit,
+            "[TMPRL1103] Attempted to upload payloads with size that exceeded the error limit.",
+            self.payload_limits.payload_size_warning,
+            "[TMPRL1103] Attempted to upload payloads with size that exceeded the warning limit.",
+        )
+
+    @staticmethod
+    def _validate_limits(
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+        error_limit: int,
+        error_message: str,
+        warning_limit: int,
+        warning_message: str,
+    ):
+        total_size = sum(payload.ByteSize() for payload in payloads)
+
+        if error_limit > 0 and total_size > error_limit:
+            raise temporalio.exceptions.PayloadSizeError(
+                f"{error_message} Size: {total_size} bytes, Limit: {error_limit} bytes"
+            )
+
+        if warning_limit > 0 and total_size > warning_limit:
+            # TODO: Use a context aware logger to log extra information about workflow/activity/etc
+            warnings.warn(
+                f"{warning_message} Size: {total_size} bytes, Limit: {warning_limit} bytes",
+                PayloadSizeWarning,
+            )
+
 
 DefaultPayloadConverter.default_encoding_payload_converters = (
     BinaryNullPayloadConverter(),

temporalio/exceptions.py

@@ -446,3 +446,17 @@ def is_cancelled_exception(exception: BaseException) -> bool:
             and isinstance(exception.cause, CancelledError)
         )
     )
+
+
+class PayloadSizeError(TemporalError):
+    """Error raised when payloads size exceeds payload size limits."""
+
+    def __init__(self, message: str):
+        """Initialize a payloads size error."""
+        super().__init__(message)
+        self._message = message
+
+    @property
+    def message(self) -> str:
+        """Message."""
+        return self._message