Batch SDK payload codec visits

Issue: zorporation/durable-workflow#450

Loop-ID: build-01

Author: durable-workflow-ops

src/durable_workflow/_avro.py

@@ -65,6 +65,19 @@ def encode(value: Any) -> str:
     The generic wrapper accepts the same value shapes as ``json.dumps``; adapt
     domain objects to JSON-native data before encoding.
     """
+    return encode_many([value])[0]
+
+
+def encode_many(values: list[Any]) -> list[str]:
+    """Encode several Avro generic-wrapper payloads through one codec visit.
+
+    The Avro runtime does not provide a useful vectorized API for independent
+    payload blobs, but batching here still avoids repeated import/schema/writer
+    setup at high fan-out command boundaries.
+    """
+    if not values:
+        return []
+
     try:
         import avro.io
     except ImportError as exc:
@@ -74,17 +87,8 @@ def encode(value: Any) -> str:
         ) from exc
 
     schema = _load_avro_schema()
-    buf = io.BytesIO()
-    encoder = avro.io.BinaryEncoder(buf)
     writer = avro.io.DatumWriter(schema)
-    writer.write(
-        {
-            "json": json.dumps(value, separators=(",", ":"), ensure_ascii=False),
-            "version": WRAPPER_VERSION,
-        },
-        encoder,
-    )
-    return base64.b64encode(_PREFIX_GENERIC_WRAPPER + buf.getvalue()).decode("ascii")
+    return [_encode_with_writer(value, writer, avro.io.BinaryEncoder) for value in values]
 
 
 def decode(blob: str) -> Any:
@@ -94,6 +98,14 @@ def decode(blob: str) -> Any:
     schemas (prefix ``0x01``) raise :class:`ValueError` because the SDK
     has no schema registry.
     """
+    return decode_many([blob])[0]
+
+
+def decode_many(blobs: list[str]) -> list[Any]:
+    """Decode several Avro payload blobs through one codec visit."""
+    if not blobs:
+        return []
+
     try:
         import avro.io
     except ImportError as exc:
@@ -102,6 +114,25 @@ def decode(blob: str) -> Any:
             "codec. Reinstall durable-workflow with its runtime dependencies."
         ) from exc
 
+    schema = _load_avro_schema()
+    reader = avro.io.DatumReader(schema)
+    return [_decode_with_reader(blob, reader, avro.io.BinaryDecoder) for blob in blobs]
+
+
+def _encode_with_writer(value: Any, writer: Any, encoder_cls: Any) -> str:
+    buf = io.BytesIO()
+    encoder = encoder_cls(buf)
+    writer.write(
+        {
+            "json": json.dumps(value, separators=(",", ":"), ensure_ascii=False),
+            "version": WRAPPER_VERSION,
+        },
+        encoder,
+    )
+    return base64.b64encode(_PREFIX_GENERIC_WRAPPER + buf.getvalue()).decode("ascii")
+
+
+def _decode_with_reader(blob: str, reader: Any, decoder_cls: Any) -> Any:
     try:
         raw = base64.b64decode(blob, validate=True)
     except (ValueError, TypeError) as exc:
@@ -124,9 +155,7 @@ def decode(blob: str) -> Any:
             f"These bytes were not produced by a Durable Workflow Avro serializer."
         )
 
-    schema = _load_avro_schema()
-    reader = avro.io.DatumReader(schema)
-    decoder = avro.io.BinaryDecoder(io.BytesIO(raw[1:]))
+    decoder = decoder_cls(io.BytesIO(raw[1:]))
     try:
         record = reader.read(decoder)
     except Exception as exc:

src/durable_workflow/client.py

@@ -644,16 +644,20 @@ async def describe_workflow(self, workflow_id: str) -> WorkflowExecution:
         ``output`` envelopes when present.
         """
         data = await self._request("GET", f"/workflows/{workflow_id}", context=workflow_id)
-        input_val = None
-        output_val = None
+        input_val = data.get("input")
+        output_val = data.get("output")
+        envelope_jobs: list[tuple[str, Any]] = []
         if data.get("input_envelope"):
-            input_val = serializer.decode_envelope(data["input_envelope"])
-        elif data.get("input") is not None:
-            input_val = data["input"]
+            envelope_jobs.append(("input", data["input_envelope"]))
         if data.get("output_envelope"):
-            output_val = serializer.decode_envelope(data["output_envelope"])
-        elif data.get("output") is not None:
-            output_val = data["output"]
+            envelope_jobs.append(("output", data["output_envelope"]))
+        if envelope_jobs:
+            decoded = serializer.decode_envelopes([envelope for _, envelope in envelope_jobs])
+            for (field, _), value in zip(envelope_jobs, decoded, strict=True):
+                if field == "input":
+                    input_val = value
+                else:
+                    output_val = value
         return WorkflowExecution(
             workflow_id=data.get("workflow_id", workflow_id),
             run_id=data.get("run_id"),

src/durable_workflow/serializer.py

@@ -228,15 +228,26 @@ def encode_many(
     call sites.
     """
     contexts = _warning_contexts_for_values(values, warning_context)
-    return [
-        encode(
-            value,
+    if codec == JSON_CODEC:
+        blobs = [
+            json.dumps(value, separators=(",", ":"), ensure_ascii=False)
+            for value in values
+        ]
+    elif codec == AVRO_CODEC:
+        blobs = _avro.encode_many(list(values))
+    else:
+        raise ValueError(
+            f"Unsupported payload codec {codec!r}; this SDK supports {SUPPORTED_CODECS!r}."
+        )
+
+    for index, blob in enumerate(blobs):
+        warn_if_payload_near_limit(
+            blob,
             codec=codec,
             size_warning=size_warning,
             warning_context=contexts[index],
         )
-        for index, value in enumerate(values)
-    ]
+    return blobs
 
 
 def envelope(
@@ -377,6 +388,60 @@ def decode_envelope(value: Any, codec: str | None = None) -> Any:
     return decode(value, codec=codec)
 
 
+def decode_envelopes(values: Sequence[Any], codec: str | None = None) -> list[Any]:
+    """Decode several raw blobs or ``{codec, blob}`` envelopes in order."""
+    jobs: list[tuple[str | None, str | None]] = []
+    passthroughs: dict[int, Any] = {}
+    for index, value in enumerate(values):
+        if isinstance(value, dict) and "codec" in value and "blob" in value:
+            jobs.append((value["blob"], value["codec"]))
+        elif value is None or value == "":
+            passthroughs[index] = None
+            jobs.append((None, None))
+        else:
+            jobs.append((value, codec))
+
+    results: list[Any] = [None] * len(values)
+    grouped: dict[str | None, list[tuple[int, str | None]]] = {}
+    for index, (blob, item_codec) in enumerate(jobs):
+        if index in passthroughs:
+            continue
+        grouped.setdefault(item_codec, []).append((index, blob))
+
+    for item_codec, group in grouped.items():
+        decoded = decode_many([blob for _, blob in group], codec=item_codec)
+        for (index, _), value in zip(group, decoded, strict=True):
+            results[index] = value
+
+    return results
+
+
+def decode_many(blobs: Sequence[str | None], codec: str | None = None) -> list[Any]:
+    """Decode several payload blobs with one codec visit when possible."""
+    if codec is None or codec == JSON_CODEC:
+        return [decode(blob, codec=codec) for blob in blobs]
+
+    if codec == AVRO_CODEC:
+        decoded: list[Any] = [None] * len(blobs)
+        avro_jobs: list[tuple[int, str]] = []
+        for index, blob in enumerate(blobs):
+            if blob is None or blob == "":
+                continue
+            avro_jobs.append((index, blob))
+        if not avro_jobs:
+            return decoded
+        avro_values = _avro.decode_many([blob for _, blob in avro_jobs])
+        for (index, _), value in zip(avro_jobs, avro_values, strict=True):
+            decoded[index] = value
+        return decoded
+
+    raise ValueError(
+        f"Cannot decode payload with codec {codec!r}: this SDK supports "
+        f"{SUPPORTED_CODECS!r}. Ensure the workflow was started with a "
+        f"compatible codec or an explicit {{'codec': '<codec>', 'blob': '...'}} envelope."
+    )
+
+
 def decode(blob: str | None, codec: str | None = None) -> Any:
     """Decode a payload blob into a Python value.
 

tests/test_client.py

@@ -217,6 +217,34 @@ async def test_envelope_fields(self, client: Client) -> None:
             assert desc.output == {"greeting": "Hello, Ada!"}
             assert desc.payload_codec == "json"
 
+    @pytest.mark.asyncio
+    async def test_envelope_fields_decode_in_one_batch(self, client: Client) -> None:
+        resp = _mock_response(200, {
+            "workflow_id": "wf-1",
+            "run_id": "run-1",
+            "workflow_type": "greeter",
+            "status": "completed",
+            "input_envelope": {"codec": "json", "blob": '["Ada"]'},
+            "output_envelope": {"codec": "json", "blob": '{"greeting":"Hello, Ada!"}'},
+        })
+
+        with (
+            patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp),
+            patch.object(
+                serializer,
+                "decode_envelopes",
+                return_value=[["Ada"], {"greeting": "Hello, Ada!"}],
+            ) as decode_envelopes,
+        ):
+            desc = await client.describe_workflow("wf-1")
+
+        decode_envelopes.assert_called_once_with([
+            {"codec": "json", "blob": '["Ada"]'},
+            {"codec": "json", "blob": '{"greeting":"Hello, Ada!"}'},
+        ])
+        assert desc.input == ["Ada"]
+        assert desc.output == {"greeting": "Hello, Ada!"}
+
     @pytest.mark.asyncio
     async def test_not_found(self, client: Client) -> None:
         resp = _mock_response(404, {"reason": "workflow_not_found", "message": "not found"})