feat(endode-extstore-refs): external storage references are now encoded and decoded.

changelog/added/discerning-murlock.md

@@ -0,0 +1 @@
+Workers and clients can now read external storage reference payloads that were encoded with the configured payload codec and marked with `__temporal_external_storage_reference` metadata. This release does not write encoded references yet; a later release will, so upgrade all workers before then.

temporalio/bridge/worker.py

@@ -315,6 +315,11 @@ async def decode_activation(
     Returns:
         Metrics from any external storage retrieval operations that occurred.
     """
+    await CommandAwarePayloadVisitor(
+        skip_search_attributes=True,
+        skip_headers=not decode_headers,
+    ).visit(_Visitor(data_converter._decode_reference_payload_sequence), activation)
+
     metrics = temporalio.converter._extstore.StorageOperationMetrics()
     with metrics.track():
         await CommandAwarePayloadVisitor(
@@ -363,4 +368,12 @@ async def encode_completion(
             completion,
         )
 
+    await CommandAwarePayloadVisitor(
+        skip_search_attributes=True,
+        skip_headers=not encode_headers,
+    ).visit(
+        _Visitor(data_converter._encode_reference_payload_sequence),
+        completion,
+    )
+
     return metrics

temporalio/converter/_data_converter.py

@@ -13,8 +13,10 @@
 import temporalio.api.common.v1
 import temporalio.api.failure.v1
 import temporalio.common
+import temporalio.converter._extstore
 from temporalio.converter._extstore import (
     _REFERENCE_ENCODING,
+    _REFERENCE_MARKER_KEY,
     _REFERENCE_MESSAGE_TYPE,
     ExternalStorage,
     StorageDriverStoreContext,
@@ -44,6 +46,11 @@ def _is_reference_payload(p: temporalio.api.common.v1.Payload) -> bool:
     )
 
 
+def _has_reference_marker(p: temporalio.api.common.v1.Payload) -> bool:
+    """Return True if *p* carries the external-storage reference marker."""
+    return p.metadata.get(_REFERENCE_MARKER_KEY) == b"true"
+
+
 # Import defaults from public API to avoid pydoctor cross-reference issues
 if TYPE_CHECKING:
     from temporalio.converter import DefaultFailureConverter, DefaultPayloadConverter
@@ -118,6 +125,7 @@ async def encode(
         payloads = self.payload_converter.to_payloads(values)
         payloads = await self._encode_payload_sequence(payloads)
         payloads = await self._external_store_payload_sequence(payloads)
+        payloads = await self._encode_reference_payload_sequence(payloads)
         return payloads
 
     async def decode(
@@ -135,6 +143,7 @@ async def decode(
         Returns:
             Decoded and converted values.
         """
+        payloads = await self._decode_reference_payload_sequence(payloads)
         payloads = await self._external_retrieve_payload_sequence(payloads)
         payloads = await self._decode_payload_sequence(payloads)
         return self.payload_converter.from_payloads(payloads, type_hints)
@@ -269,7 +278,7 @@ async def _transform_outbound_payload(
             payload = (await self.payload_codec.encode([payload]))[0]
         if self.external_storage:
             payload = await self.external_storage._store_payload(payload)
-        return payload
+        return (await self._encode_reference_payload_sequence([payload]))[0]
 
     async def _transform_outbound_payloads(
         self, payloads: temporalio.api.common.v1.Payloads
@@ -278,10 +287,16 @@ async def _transform_outbound_payloads(
             await self.payload_codec.encode_wrapper(payloads)
         if self.external_storage:
             await self.external_storage._store_payloads(payloads)
+        encoded_payloads = await self._encode_reference_payload_sequence(
+            payloads.payloads
+        )
+        for i, payload in enumerate(encoded_payloads):
+            payloads.payloads[i].CopyFrom(payload)
 
     async def _transform_inbound_payload(
         self, payload: temporalio.api.common.v1.Payload
     ) -> temporalio.api.common.v1.Payload:
+        payload = (await self._decode_reference_payload_sequence([payload]))[0]
         if self.external_storage:
             payload = await self.external_storage._retrieve_payload(payload)
         if self.payload_codec:
@@ -291,6 +306,11 @@ async def _transform_inbound_payload(
     async def _transform_inbound_payloads(
         self, payloads: temporalio.api.common.v1.Payloads
     ):
+        decoded_payloads = await self._decode_reference_payload_sequence(
+            payloads.payloads
+        )
+        for i, payload in enumerate(decoded_payloads):
+            payloads.payloads[i].CopyFrom(payload)
         if self.external_storage:
             await self.external_storage._retrieve_payloads(payloads)
         else:
@@ -321,6 +341,55 @@ async def _external_store_payload_sequence(
             )
         return stored_payloads
 
+    async def _encode_reference_payload_sequence(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        """Codec encode external storage reference payloads only."""
+        encoded_payloads = list(payloads)
+        if (
+            not temporalio.converter._extstore._ENCODE_STORAGE_REFERENCES
+            or not self.payload_codec
+        ):
+            return encoded_payloads
+        for i, payload in enumerate(encoded_payloads):
+            if not _is_reference_payload(payload):
+                continue
+            [encoded_reference] = await self.payload_codec.encode([payload])
+            if not encoded_reference.external_payloads:
+                encoded_reference.external_payloads.extend(payload.external_payloads)
+            encoded_reference.metadata[_REFERENCE_MARKER_KEY] = b"true"
+            encoded_payloads[i] = encoded_reference
+        return encoded_payloads
+
+    async def _decode_reference_payload_sequence(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        """Codec decode marked external storage reference payloads only."""
+        decoded_payloads = list(payloads)
+        for i, payload in enumerate(decoded_payloads):
+            if not _has_reference_marker(payload):
+                continue
+            if not self.external_storage:
+                raise RuntimeError(
+                    "[TMPRL1105] Detected externally stored payload(s) but external storage is not configured."
+                )
+            reference = payload
+            if self.payload_codec:
+                encoded_reference = temporalio.api.common.v1.Payload()
+                encoded_reference.CopyFrom(payload)
+                del encoded_reference.metadata[_REFERENCE_MARKER_KEY]
+                del encoded_reference.external_payloads[:]
+                [reference] = await self.payload_codec.decode([encoded_reference])
+            if not _is_reference_payload(reference):
+                raise RuntimeError(
+                    "[TMPRL1106] Encountered an encoded external storage reference that could not be decoded. "
+                    "Check that the payload codec used to write it is configured."
+                )
+            if not reference.external_payloads:
+                reference.external_payloads.extend(payload.external_payloads)
+            decoded_payloads[i] = reference
+        return decoded_payloads
+
     async def _external_retrieve_payload_sequence(
         self, payloads: Sequence[temporalio.api.common.v1.Payload]
     ) -> list[temporalio.api.common.v1.Payload]:

temporalio/converter/_extstore.py

@@ -28,6 +28,15 @@
 
 _REFERENCE_ENCODING = b"json/external-storage-reference"
 _REFERENCE_MESSAGE_TYPE = ExternalStorageReference.DESCRIPTOR.full_name.encode()
+# Indicates this payload represents an externally stored payload. If encoded, the
+# _REFERENCE_ENCODING could be any value. The presence of the _REFERENCE_MARKER_KEY
+# with a value of b"true" is used to reliably identify such payloads.
+_REFERENCE_MARKER_KEY = "__temporal_external_storage_reference"
+# Writing marked, codec-encoded references is off until a release that can read
+# them has been out long enough for workers to upgrade: workers on earlier
+# versions cannot read encoded references. Reading is always on. Set to True in a
+# later release to start writing them.
+_ENCODE_STORAGE_REFERENCES = False
 
 
 @dataclass
@@ -430,6 +439,8 @@ async def _store_payload(self, payload: Payload) -> Payload:
                 f"Failed to serialize storage reference for driver '{driver.name()}'"
             )
         reference_payload.external_payloads.add().size_bytes = external_size
+        if _ENCODE_STORAGE_REFERENCES:
+            reference_payload.metadata[_REFERENCE_MARKER_KEY] = b"true"
 
         ExternalStorage._record_metrics(1, external_size, start_time, {driver.name()})
 
@@ -494,6 +505,8 @@ async def _store_payload_sequence(
                         f"Failed to serialize storage reference for driver '{driver.name()}'"
                     )
                 reference_payload.external_payloads.add().size_bytes = sizes[i]
+                if _ENCODE_STORAGE_REFERENCES:
+                    reference_payload.metadata[_REFERENCE_MARKER_KEY] = b"true"
                 results[indices[i]] = reference_payload
                 external_size += sizes[i]
 

temporalio/worker/_nexus.py

@@ -234,7 +234,9 @@ async def _encode_completion(
         completion: temporalio.bridge.proto.nexus.NexusTaskCompletion,
         data_converter: temporalio.converter.DataConverter,
     ) -> None:
-        """Apply the payload codec then external storage to the completion's payloads."""
+        """Apply the payload codec, external storage, then the payload codec to
+        any storage references in the completion's payloads.
+        """
         await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit(
             _PayloadTransformVisitor(data_converter._encode_payload_sequence),
             completion,
@@ -243,6 +245,10 @@ async def _encode_completion(
             _PayloadTransformVisitor(data_converter._external_store_payload_sequence),
             completion,
         )
+        await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit(
+            _PayloadTransformVisitor(data_converter._encode_reference_payload_sequence),
+            completion,
+        )
 
     def _data_converter_for_nexus_task(
         self, endpoint: str, service: str, operation: str
@@ -566,6 +572,17 @@ async def deserialize(
         # payload untouched.
         payload = temporalio.api.common.v1.Payload()
         payload.CopyFrom(self.payload)
+        try:
+            await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit(
+                _PayloadTransformVisitor(dc._decode_reference_payload_sequence),
+                payload,
+            )
+        except Exception as err:
+            raise nexusrpc.HandlerError(
+                "Payload codec failed to decode Nexus operation input storage reference",
+                type=nexusrpc.HandlerErrorType.INTERNAL,
+            ) from err
+
         try:
             await PayloadVisitor(skip_search_attributes=True).visit(
                 _PayloadTransformVisitor(dc._external_retrieve_payload_sequence),

temporalio/worker/_workflow.py

@@ -956,6 +956,16 @@ async def _external_store_payload_sequence(
         dc = self._get_current_dc()._with_store_context(store_ctx)
         return await dc._external_store_payload_sequence(payloads)
 
+    async def _encode_reference_payload_sequence(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        return await self._get_current_dc()._encode_reference_payload_sequence(payloads)
+
+    async def _decode_reference_payload_sequence(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        return await self._get_current_dc()._decode_reference_payload_sequence(payloads)
+
     async def _external_retrieve_payload_sequence(
         self, payloads: Sequence[temporalio.api.common.v1.Payload]
     ) -> list[temporalio.api.common.v1.Payload]:

tests/conftest.py

@@ -14,6 +14,7 @@
 from opentelemetry.metrics import NoOpMeterProvider
 from opentelemetry.util._once import Once
 
+import temporalio.converter._extstore
 from temporalio.client import Client
 from temporalio.envconfig import ClientConfigProfile
 from temporalio.testing import WorkflowEnvironment
@@ -351,3 +352,12 @@ def reset_otel_logger_provider():
     yield
     opentelemetry._logs._internal._LOGGER_PROVIDER_SET_ONCE = Once()
     opentelemetry._logs._internal._LOGGER_PROVIDER = None
+
+
+@pytest.fixture
+def encode_storage_references(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Turn on writing marked, codec-encoded external storage references, which
+    is off by default until a later release."""
+    monkeypatch.setattr(
+        temporalio.converter._extstore, "_ENCODE_STORAGE_REFERENCES", True
+    )

tests/nexus/test_temporal_extstore.py

@@ -30,6 +30,7 @@
 from temporalio.client import Client, WorkflowFailureError
 from temporalio.converter import (
     ExternalStorage,
+    PayloadCodec,
     StorageDriverClaim,
     StorageDriverRetrieveContext,
     StorageDriverStoreContext,
@@ -39,7 +40,7 @@
 from temporalio.types import MethodAsyncSingleParam
 from temporalio.worker import UnsandboxedWorkflowRunner, Worker
 from tests.helpers.nexus import make_nexus_endpoint_name
-from tests.test_extstore import InMemoryTestDriver
+from tests.test_extstore import InMemoryTestDriver, RecordingPayloadCodec
 
 # Cloud CI's namespace credentials cannot manage Nexus endpoints.
 # See https://github.com/temporalio/sdk-python/issues/1704.
@@ -137,11 +138,14 @@ async def retrieve(
 
 
 def _client_with_extstore(
-    env: WorkflowEnvironment, driver: InMemoryTestDriver
+    env: WorkflowEnvironment,
+    driver: InMemoryTestDriver,
+    payload_codec: PayloadCodec | None = None,
 ) -> Client:
     config = env.client.config()
     config["data_converter"] = dataclasses.replace(
         temporalio.converter.default(),
+        payload_codec=payload_codec,
         external_storage=ExternalStorage(
             drivers=[driver],
             payload_size_threshold=PAYLOAD_SIZE_THRESHOLD,
@@ -154,8 +158,11 @@ async def _run_caller(
     env: WorkflowEnvironment,
     driver: InMemoryTestDriver,
     workflow_run: MethodAsyncSingleParam[Any, str, int],
+    *,
+    payload_codec: PayloadCodec | None = None,
+    workflow_id: str | None = None,
 ) -> int:
-    client = _client_with_extstore(env, driver)
+    client = _client_with_extstore(env, driver, payload_codec)
     task_queue = str(uuid.uuid4())
     async with Worker(
         client,
@@ -170,7 +177,7 @@ async def _run_caller(
         return await client.execute_workflow(
             workflow_run,
             task_queue,
-            id=str(uuid.uuid4()),
+            id=workflow_id or str(uuid.uuid4()),
             task_queue=task_queue,
             execution_timeout=timedelta(seconds=30),
         )
@@ -213,6 +220,69 @@ async def test_nexus_operation_sync_result_offloaded_and_retrieved(
     assert driver._retrieve_calls >= 1
 
 
+def _assert_encoded_reference(payload: Payload) -> None:
+    assert payload.metadata["encoding"] == b"binary/test-encoded"
+    assert payload.metadata["__temporal_external_storage_reference"] == b"true"
+
+
+@pytest.mark.usefixtures("encode_storage_references")
+async def test_nexus_operation_input_reference_encoded_with_codec(
+    env: WorkflowEnvironment,
+):
+    """With a codec, the operation input reference is encoded, and the Nexus
+    worker decodes it before retrieving the input."""
+    if env.supports_time_skipping:
+        pytest.skip("Nexus tests don't work with the Java test server")
+
+    driver = InMemoryTestDriver()
+    workflow_id = str(uuid.uuid4())
+    result = await _run_caller(
+        env,
+        driver,
+        SizeOpCallerWorkflow.run,
+        payload_codec=RecordingPayloadCodec("binary/test-encoded"),
+        workflow_id=workflow_id,
+    )
+
+    assert result == PAYLOAD_SIZE
+    history = await env.client.get_workflow_handle(workflow_id).fetch_history()
+    [scheduled] = [
+        e.nexus_operation_scheduled_event_attributes
+        for e in history.events
+        if e.HasField("nexus_operation_scheduled_event_attributes")
+    ]
+    _assert_encoded_reference(scheduled.input)
+
+
+@pytest.mark.usefixtures("encode_storage_references")
+async def test_nexus_operation_result_reference_encoded_with_codec(
+    env: WorkflowEnvironment,
+):
+    """With a codec, the Nexus worker encodes the reference to an offloaded
+    operation result, and the caller decodes it."""
+    if env.supports_time_skipping:
+        pytest.skip("Nexus tests don't work with the Java test server")
+
+    driver = InMemoryTestDriver()
+    workflow_id = str(uuid.uuid4())
+    result = await _run_caller(
+        env,
+        driver,
+        BigResultOpCallerWorkflow.run,
+        payload_codec=RecordingPayloadCodec("binary/test-encoded"),
+        workflow_id=workflow_id,
+    )
+
+    assert result == PAYLOAD_SIZE
+    history = await env.client.get_workflow_handle(workflow_id).fetch_history()
+    [completed] = [
+        e.nexus_operation_completed_event_attributes
+        for e in history.events
+        if e.HasField("nexus_operation_completed_event_attributes")
+    ]
+    _assert_encoded_reference(completed.result)
+
+
 async def test_nexus_operation_transient_retrieve_failure_recovers(
     env: WorkflowEnvironment,
 ):