Optimize deterministic set serialization without JSON parse round trips

CHANGELOG.md

@@ -14,6 +14,15 @@ by one or more task hubs.
 
 FIXED
 
+- Optimized set and frozenset serialization, particularly for deeply nested
+frozensets and large singleton elements. JSON object names within sets are now
+ordered consistently after conversion to strings, including numeric names.
+Equivalent surrogate-pair and non-BMP object names retain the last value
+rather than emitting duplicate JSON properties.
+Set reconstruction remains opt-in.
+- With set reconstruction enabled, annotated sets of typed tuples are now
+recognized during input and activity-result discovery, and bare
+`collections.abc.Set` annotations reconstruct JSON arrays as concrete sets.
 - Implemented `__hash__` and `__repr__` for `EntityInstanceId`, making identity attributes immutable to guarantee hash stability when used in sets or dict keys.
 - Large-payload externalization now checks every eligible payload's maximum
 size before uploading or replacing any protobuf values, including carryover

durabletask/serialization.py

@@ -38,6 +38,7 @@
 import inspect
 import json
 import logging
+import re
 import sys
 import types
 import typing
@@ -56,6 +57,8 @@
 # replay. Internal detail; re-exported from ``durabletask.internal.shared`` only
 # for backwards compatibility.
 _AUTO_SERIALIZED = "__durabletask_autoobject__"
+_SURROGATE_PAIR = re.compile("[\ud800-\udbff][\udc00-\udfff]")
+_JSON_SCALAR_TYPES = (int, str, float, type(None))
 
 
 class DataConverter(ABC):
@@ -137,6 +140,8 @@ class JsonDataConverter(DataConverter):
     own hooks. Sets and frozensets serialize to JSON arrays ordered
     deterministically by their canonical JSON representation (using lexical
     string order, so numeric sets like ``{2, 10}`` serialize as ``[10, 2]``).
+    Within sets, object keys are normalized to strings and sorted, and custom
+    hooks run once per occurrence. Singleton sets need no ordering key.
 
     Deserialization (and value-level :meth:`coerce`) is **best-effort**: when a
     ``target_type`` is supplied and the value cannot be coerced to it, the raw
@@ -247,6 +252,11 @@ def _is_supported_set_element(converter: DataConverter, elem_type: Any) -> bool:
         if elem_origin is frozenset:
             args = typing.get_args(elem_type)
             return not args or _is_supported_set_element(converter, args[0])
+        if elem_origin is tuple:
+            args = typing.get_args(elem_type)
+            if len(args) == 2 and args[1] is Ellipsis:
+                args = args[:1]
+            return all(_is_supported_set_element(converter, arg) for arg in args)
         return False
     if not isinstance(elem_type, type):
         return False
@@ -368,25 +378,9 @@ def _encode_custom_object(o: Any) -> Any:
     if callable(to_json_hook):
         return to_json_hook(o)
     if isinstance(o, (set, frozenset, AbstractSet)):
-        items: AbstractSet[Any] = cast("AbstractSet[Any]", o)
-        # Sets and frozensets are unordered collections. To ensure deterministic
-        # serialization across processes, worker replays, and PYTHONHASHSEED
-        # variations, elements are individually encoded to canonical JSON
-        # (with sort_keys=True), sorted lexically by their encoded JSON string,
-        # and emitted as normalized parsed values.
-        #
-        # Emitting the normalized parsed values rather than original objects
-        # guarantees identical order and eliminates leaks from dict insertion
-        # order variations, NaN/non-finite comparison anomalies, or partial
-        # ordering in frozensets. A consequence of sorting canonical JSON strings
-        # is lexical ordering for numeric sets (e.g., {2, 10} serializes as
-        # [10, 2]).
+        items = cast(AbstractSet[Any], o)
         try:
-            encoded_items = sorted(
-                json.dumps(item, default=_encode_custom_object, sort_keys=True)
-                for item in items
-            )
-            return [json.loads(encoded) for encoded in encoded_items]
+            return _normalize_set_value(items, set())
         except Exception as e:
             raise TypeError(
                 f"Cannot deterministically serialize set of type '{type(items).__name__}' "
@@ -415,6 +409,66 @@ def _encode_custom_object(o: Any) -> Any:
     raise TypeError(f"Object of type '{type(o).__name__}' is not JSON serializable")
 
 
+def _normalize_set_value(value: Any, active: set[int]) -> Any:
+    """Normalize only set subtrees; leave ordinary payload encoding unchanged."""
+    identity = id(value)
+    is_mapping = isinstance(value, dict)
+    if not is_mapping and isinstance(value, _JSON_SCALAR_TYPES):
+        return value
+    if identity in active:
+        raise ValueError("Circular reference detected")
+    if is_mapping:
+        # Match sort_keys=True, including rejection of incomparable keys.
+        # JSON object names are strings; duplicate names keep the last value,
+        # just as the previous dumps/loads normalization did.
+        result: dict[str, Any] = {}
+        string_names = True
+        for key, item in sorted(cast(dict[Any, Any], value).items()):
+            if isinstance(key, str):
+                name = key
+                if not str.isascii(name) and _SURROGATE_PAIR.search(name):
+                    # Ordinary Unicode names and lone surrogates need no codec
+                    # round trip; only a high/low pair changes after decoding.
+                    name = str.encode(name, "utf-16-le", "surrogatepass").decode(
+                        "utf-16-le", "surrogatepass")
+                    string_names = False
+            else:
+                if key is not None and not isinstance(key, (int, float)):
+                    raise TypeError(f"Unsupported JSON object key: {type(key).__name__}")
+                name = json.dumps(key)
+                string_names = False
+            if isinstance(item, _JSON_SCALAR_TYPES):
+                result[name] = item
+            else:
+                # Scalar-only records cannot cycle. Track the parent only
+                # while descending into a child that could refer back to it.
+                active.add(identity)
+                try:
+                    result[name] = _normalize_set_value(item, active)
+                finally:
+                    active.remove(identity)
+        return result if string_names else dict(sorted(result.items()))
+    active.add(identity)
+    try:
+        if isinstance(value, (list, tuple)):
+            return [
+                item if isinstance(item, _JSON_SCALAR_TYPES)
+                else _normalize_set_value(item, active)
+                for item in cast(Sequence[Any], value)
+            ]
+        to_json_hook = getattr(cast(Any, type(value)), "to_json", None)
+        if isinstance(value, AbstractSet) and not callable(to_json_hook):
+            items = [_normalize_set_value(item, active) for item in cast(AbstractSet[Any], value)]
+            if len(items) > 1:
+                # list.sort evaluates each key once and discards the temporary
+                # documents afterwards. Emit these same normalized values.
+                items.sort(key=json.dumps)
+            return items
+        return _normalize_set_value(_encode_custom_object(value), active)
+    finally:
+        active.remove(identity)
+
+
 def _legacy_object_hook(d: dict[str, Any]) -> Any:
     # If the object carries the legacy marker, deserialize it as a SimpleNamespace.
     if d.pop(_AUTO_SERIALIZED, False):
@@ -458,6 +512,10 @@ def _coerce_to_type(value: Any, expected_type: Any, converter: DataConverter | N
     if isinstance(value, expected_type):
         return value
 
+    if (expected_type is AbstractSet and isinstance(value, list)
+            and getattr(converter, "enable_generic_sets", False)):
+        return _coerce_generic(value, expected_type, AbstractSet, converter)
+
     from_json_hook = getattr(expected_type, "from_json", None)
     if callable(from_json_hook):
         return _invoke_from_json(from_json_hook, value, converter)
@@ -471,7 +529,7 @@ def _coerce_to_type(value: Any, expected_type: Any, converter: DataConverter | N
     except Exception as e:
         type_name = getattr(type_ctor, "__name__", None) or str(type_ctor)
         raise TypeError(
-            f"Could not coerce value of type '{type(value).__name__}' to "
+            f"Could not coerce value of type '{type(cast(object, value)).__name__}' to "
             f"'{type_name}'"
         ) from e
 

tests/durabletask/test_data_converter.py

@@ -6,20 +6,24 @@
 import inspect
 import json
 import logging
+from collections.abc import Set as AbstractSet
 from dataclasses import dataclass, field
 from functools import wraps
-from typing import Any
+from typing import AbstractSet as TypingAbstractSet, Any
 from unittest.mock import patch
 
 import pytest
+from durabletask import task
+from durabletask.internal import helpers
+from durabletask.internal import orchestrator_service_pb2 as pb
 from durabletask.internal.entity_state_shim import StateShim
 from durabletask.serialization import (
     DEFAULT_DATA_CONVERTER,
     DataConverter,
     JsonDataConverter,
 )
 from durabletask.internal import type_discovery
-from durabletask.worker import _ActivityExecutor, _Registry
+from durabletask.worker import _ActivityExecutor, _OrchestrationExecutor, _Registry
 
 
 @dataclass
@@ -423,6 +427,60 @@ def test_can_reconstruct_set_and_frozenset():
     assert opt_in_converter.can_reconstruct(set[dict[str, int]]) is False
 
 
+@pytest.mark.parametrize("option", ["enable_generic_sets", "enable_set_coercion"])
+@pytest.mark.parametrize("target_type,payload,expected", [
+    (set[tuple[int, int]], "[[1, 2], [3, 4]]", {(1, 2), (3, 4)}),
+    (frozenset[tuple[int, ...]], "[[1], [2, 3]]", frozenset({(1,), (2, 3)})),
+    (set[tuple[int, frozenset[str]]], '[[1, ["b", "a"]]]', {(1, frozenset({"a", "b"}))}),
+    (frozenset[frozenset[tuple[int, int]]], "[[[1, 2], [3, 4]]]",
+     frozenset({frozenset({(1, 2), (3, 4)})})),
+    (set[tuple[()]], "[[]]", {()}),
+    (set[tuple[int, str | None]], '[[1, null], [2, "x"]]', {(1, None), (2, "x")}),
+    (AbstractSet, "[2, 1, 2]", {1, 2}),
+    (TypingAbstractSet, "[2, 1, 2]", {1, 2}),
+])
+def test_opt_in_set_discovery_matches_reconstruction(
+    option: str, target_type: Any, payload: str, expected: Any,
+):
+    converter = JsonDataConverter(**{option: True})
+
+    def activity(ctx: Any, values: Any) -> Any:
+        assert isinstance(values, type(expected))
+        assert values == expected
+        return values
+
+    activity.__annotations__["values"] = target_type
+    activity.__annotations__["return"] = target_type
+    assert converter.can_reconstruct(target_type)
+    assert converter.deserialize(payload, target_type) == expected
+    assert converter.coerce(json.loads(payload), target_type) == expected
+    assert type_discovery.activity_input_type(activity, converter) == target_type
+    assert type_discovery.orchestrator_input_type(activity, converter) == target_type
+    assert type_discovery.activity_output_type(activity, converter) == target_type
+    registry = _Registry()
+    name = registry.add_activity(activity)
+    executor = _ActivityExecutor(registry, logging.getLogger("test_set_discovery"), converter)
+    output = executor.execute("instance", name, 1, payload)
+    assert converter.deserialize(output, target_type) == expected
+
+    default = JsonDataConverter()
+    assert not default.can_reconstruct(target_type)
+    assert type_discovery.activity_input_type(activity, default) is None
+    assert type_discovery.orchestrator_input_type(activity, default) is None
+    assert type_discovery.activity_output_type(activity, default) is None
+    assert default.deserialize(payload, target_type) == json.loads(payload)
+
+
+@pytest.mark.parametrize("target_type", [
+    set[tuple[list[int], int]],
+    set[tuple[dict[str, int], ...]],
+    set[tuple[frozenset[list[int]], int]],
+    frozenset[tuple[set[int], ...]],
+])
+def test_set_tuple_discovery_rejects_unhashable_elements(target_type: Any):
+    assert not JsonDataConverter(enable_generic_sets=True).can_reconstruct(target_type)
+
+
 def test_existing_instance_replay_compatibility_preserves_list():
     """Preserve existing-instance replay behavior for recorded arrays under set annotations.
 
@@ -453,6 +511,29 @@ def test_existing_instance_replay_compatibility_preserves_list():
     assert len(opt_in_restored.flags) == 2
 
 
+def test_default_converter_replays_recorded_arrays_under_set_annotations():
+    def orchestrator(ctx: task.OrchestrationContext, data: TaggedOrder):
+        yield ctx.call_activity(f"count-{len(data.tags)}")
+        return {"tags": data.tags, "flags": data.flags}
+
+    registry = _Registry()
+    name = registry.add_orchestrator(orchestrator)
+    history = [
+        helpers.new_orchestrator_started_event(),
+        helpers.new_execution_started_event(
+            name, "legacy-sets",
+            encoded_input='{"item": "laptop", "tags": ["b", "a", "b"], "flags": [1, 2, 1]}',
+        ),
+        helpers.new_task_scheduled_event(1, "count-3"),
+    ]
+    executor = _OrchestrationExecutor(registry, logging.getLogger("test_set_replay"), JsonDataConverter())
+    result = executor.execute("legacy-sets", history, [helpers.new_task_completed_event(1, "null")])
+    assert len(result.actions) == 1
+    completion = result.actions[0].completeOrchestration
+    assert completion.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert completion.result.value == '{"tags": ["b", "a", "b"], "flags": [1, 2, 1]}'
+
+
 def test_type_discovery_and_activity_execution_with_sets():
     """Verify that type discovery recognizes set annotations when opted in,
     enabling activities to receive sets rather than raw lists."""