💥 Payload limit configuration and validation (#1288)

Author: jmaeagle99

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

@@ -9178,7 +9178,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

@@ -930,7 +930,10 @@ def to_failure(
         else:
             # Convert to failure error
             failure_error = temporalio.exceptions.ApplicationError(
-                str(exception), type=exception.__class__.__name__
+                str(exception),
+                type="PayloadSizeError"
+                if isinstance(exception, _PayloadSizeError)
+                else exception.__class__.__name__,
             )
             failure_error.__traceback__ = exception.__traceback__
             failure_error.__cause__ = exception.__cause__
@@ -1220,6 +1223,46 @@ 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."""
+
+
+class _PayloadSizeError(temporalio.exceptions.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
+
+
+@dataclass(frozen=True)
+class _ServerPayloadErrorLimits:
+    """Error limits for payloads as described by the Temporal server."""
+
+    memo_size_error: int
+    """The limit (in bytes) at which a memo size error is raised."""
+
+    payload_size_error: int
+    """The limit (in bytes) at which a payload size error is raised."""
+
+
 @dataclass(frozen=True)
 class DataConverter(WithSerializationContext):
     """Data converter for converting and encoding payloads to/from Python values.
@@ -1243,9 +1286,15 @@ 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."""
 
+    _payload_error_limits: _ServerPayloadErrorLimits | None = None
+    """Server-reported limits for payloads."""
+
     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())
@@ -1347,6 +1396,11 @@ def with_context(self, context: SerializationContext) -> Self:
         object.__setattr__(cloned, "failure_converter", failure_converter)
         return cloned
 
+    def _with_payload_error_limits(
+        self, limits: _ServerPayloadErrorLimits | None
+    ) -> DataConverter:
+        return dataclasses.replace(self, _payload_error_limits=limits)
+
     async def _decode_memo(
         self,
         source: temporalio.api.common.v1.Memo,
@@ -1385,30 +1439,38 @@ async def _encode_memo_existing(
             if not isinstance(v, temporalio.api.common.v1.Payload):
                 payload = (await self.encode([v]))[0]
             memo.fields[k].CopyFrom(payload)
+        # Memos have their field payloads validated all together in one unit
+        DataConverter._validate_limits(
+            list(memo.fields.values()),
+            self._payload_error_limits.memo_size_error
+            if self._payload_error_limits
+            else None,
+            "[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
@@ -1465,6 +1527,42 @@ 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_error_limits.payload_size_error
+            if self._payload_error_limits
+            else None,
+            "[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 | None,
+        error_message: str,
+        warning_limit: int,
+        warning_message: str,
+    ):
+        total_size = sum(payload.ByteSize() for payload in payloads)
+
+        if error_limit and error_limit > 0 and total_size > error_limit:
+            raise _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(),