From 32933721f2abcdfc4010939714d77a6fa2d7836a Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:14:39 -0700 Subject: [PATCH 01/15] feat: add external storage runner --- crates/common-wasm/src/external_storage.rs | 178 +-- crates/common/Cargo.toml | 2 +- .../common/src/external_storage_references.rs | 62 +- crates/common/src/external_storage_runner.rs | 1081 +++++++++++++++++ crates/common/src/lib.rs | 2 + crates/common/src/payload_visitor.rs | 9 +- 6 files changed, 1200 insertions(+), 134 deletions(-) create mode 100644 crates/common/src/external_storage_runner.rs diff --git a/crates/common-wasm/src/external_storage.rs b/crates/common-wasm/src/external_storage.rs index 398dcfed6..1913921d9 100644 --- a/crates/common-wasm/src/external_storage.rs +++ b/crates/common-wasm/src/external_storage.rs @@ -146,11 +146,13 @@ impl Default for ExternalStorageConcurrency { } } +#[bon::bon] impl ExternalStorageConcurrency { - /// Create concurrency limits, setting both of them. Each must be at least 1. + /// Create concurrency limits. Each must be at least 1, and each defaults when omitted. + #[builder(state_mod(vis = "pub"))] pub fn new( - max_driver_operations: usize, - max_operations_per_message: usize, + #[builder(default = DEFAULT_MAX_DRIVER_OPERATIONS)] max_driver_operations: usize, + #[builder(default = DEFAULT_MAX_OPERATIONS_PER_MESSAGE)] max_operations_per_message: usize, ) -> Result { check_limit(max_driver_operations, "max_driver_operations")?; check_limit(max_operations_per_message, "max_operations_per_message")?; @@ -159,30 +161,11 @@ impl ExternalStorageConcurrency { max_operations_per_message, }) } - - /// Set the limit shared across the whole configuration. Must be at least 1. - pub fn with_max_driver_operations( - mut self, - max_driver_operations: usize, - ) -> Result { - check_limit(max_driver_operations, "max_driver_operations")?; - self.max_driver_operations = max_driver_operations; - Ok(self) - } - - /// Set the limit for one message. Must be at least 1. - pub fn with_max_operations_per_message( - mut self, - max_operations_per_message: usize, - ) -> Result { - check_limit(max_operations_per_message, "max_operations_per_message")?; - self.max_operations_per_message = max_operations_per_message; - Ok(self) - } } /// Context given to [`StorageDriver::store`]. -#[derive(Clone)] +#[derive(Clone, bon::Builder)] +#[builder(state_mod(vis = "pub"))] #[non_exhaustive] pub struct StorageDriverStoreContext { /// The execution the payloads are being stored on behalf of, when there is one. @@ -200,7 +183,8 @@ impl fmt::Debug for StorageDriverStoreContext { } /// Context given to [`StorageDriverSelector::select_driver`]. -#[derive(Clone, Debug, Default)] +#[derive(Clone, Debug, Default, bon::Builder)] +#[builder(state_mod(vis = "pub"))] #[non_exhaustive] pub struct StorageDriverSelectContext { /// The execution the payload is being stored on behalf of, when there is one. @@ -212,7 +196,8 @@ pub struct StorageDriverSelectContext { /// This deliberately carries no target. A payload must be retrievable from its /// [`StorageDriverClaim`] alone, since the same payload can be read from a different execution than /// the one that stored it. -#[derive(Clone)] +#[derive(Clone, bon::Builder)] +#[builder(state_mod(vis = "pub"))] #[non_exhaustive] pub struct StorageDriverRetrieveContext { /// Limiter for this driver's retrieve requests. @@ -287,6 +272,27 @@ pub enum ExternalStorageError { /// A driver failed to store or retrieve. #[error("Storage driver error: {0}")] Driver(#[source] Box), + /// Walking the message's payloads failed, outside any storage operation. + #[error("External storage payload walk failed: {0}")] + Walk(#[source] Box), + /// A reference names a driver that is not registered. + #[error("No storage driver is registered under the name '{0}'")] + UnregisteredDriver(String), + /// The selector returned a driver that is not one of the registered instances. + #[error("Selector returned a driver that is not registered: '{0}'")] + UnregisteredSelection(String), + /// A driver returned a different number of results than it was given inputs. A store is given + /// payloads and returns claims; a retrieve is given claims and returns payloads. Either way + /// there must be one result per input, in order. + #[error("Storage driver '{driver}' returned {returned} results for {requested} inputs")] + DriverContract { + /// Name of the driver. + driver: String, + /// How many inputs it was given. + requested: usize, + /// How many results it returned. + returned: usize, + }, } /// Configuration for offloading large payloads to an external storage system. @@ -424,8 +430,7 @@ impl StorageDriverSelector for SingleDriverSelector { #[cfg(test)] mod tests { use super::*; - use futures::{FutureExt, future}; - use std::sync::atomic::{AtomicBool, Ordering}; + use futures::future; struct FakeDriver(&'static str); @@ -468,122 +473,49 @@ mod tests { Arc::new(FakeDriver(name)) } - struct ImmediateLimiter; - - impl StorageDriverLimiter for ImmediateLimiter { - fn acquire<'a>( - &'a self, - _item: &'a Item, - ) -> BoxFuture<'a, Result> { - Box::pin(future::ready(Ok(StorageDriverPermit::new(Box::new(()))))) - } - } - #[test] - fn concurrency_defaults_until_overridden() { - let storage = ExternalStorage::new(Arc::new(FakeDriver("mem"))).unwrap(); + fn concurrency_defaults_to_the_spec_values_and_each_limit_overrides_alone() { + // The spec's defaults; `#[non_exhaustive]` makes the builder the only way to set one limit + // without restating the other. + let defaults = ExternalStorageConcurrency::default(); + assert_eq!(defaults.max_driver_operations, 64); + assert_eq!(defaults.max_operations_per_message, 8); assert_eq!( - storage.concurrency(), - &ExternalStorageConcurrency::default() + ExternalStorage::new(Arc::new(FakeDriver("mem"))) + .unwrap() + .concurrency(), + &defaults ); - let storage = storage.with_concurrency(ExternalStorageConcurrency::new(2, 1).unwrap()); - assert_eq!(storage.concurrency().max_driver_operations, 2); - assert_eq!(storage.concurrency().max_operations_per_message, 1); - } - - #[test] - fn one_limit_can_be_set_without_restating_the_other() { - // `#[non_exhaustive]` blocks struct literals and functional update syntax downstream, so - // without these builders a caller would have to repeat the default it did not want to - // change. - let only_shared = ExternalStorageConcurrency::default() - .with_max_driver_operations(10) + let only_shared = ExternalStorageConcurrency::builder() + .max_driver_operations(10) + .build() .unwrap(); assert_eq!(only_shared.max_driver_operations, 10); assert_eq!(only_shared.max_operations_per_message, 8); - let only_message = ExternalStorageConcurrency::default() - .with_max_operations_per_message(2) + let only_message = ExternalStorageConcurrency::builder() + .max_operations_per_message(2) + .build() .unwrap(); assert_eq!(only_message.max_driver_operations, 64); assert_eq!(only_message.max_operations_per_message, 2); - - let both = ExternalStorageConcurrency::default() - .with_max_driver_operations(10) - .unwrap() - .with_max_operations_per_message(2) - .unwrap(); - assert_eq!(both, ExternalStorageConcurrency::new(10, 2).unwrap()); - } - - #[test] - fn builders_reject_a_limit_of_zero() { - assert!(matches!( - ExternalStorageConcurrency::default().with_max_driver_operations(0), - Err(ExternalStorageError::InvalidConfiguration(_)) - )); - assert!(matches!( - ExternalStorageConcurrency::default().with_max_operations_per_message(0), - Err(ExternalStorageError::InvalidConfiguration(_)) - )); } #[test] fn a_limit_of_zero_is_rejected() { assert!(matches!( - ExternalStorageConcurrency::new(0, 8), + ExternalStorageConcurrency::builder() + .max_driver_operations(0) + .build(), Err(ExternalStorageError::InvalidConfiguration(_)) )); assert!(matches!( - ExternalStorageConcurrency::new(64, 0), + ExternalStorageConcurrency::builder() + .max_operations_per_message(0) + .build(), Err(ExternalStorageError::InvalidConfiguration(_)) )); - assert!(ExternalStorageConcurrency::new(1, 1).is_ok()); - } - - #[test] - fn dropping_a_permit_releases_the_limit() { - struct ReleaseOnDrop(Arc); - impl Drop for ReleaseOnDrop { - fn drop(&mut self) { - self.0.store(true, Ordering::SeqCst); - } - } - - let released = Arc::new(AtomicBool::new(false)); - let permit = StorageDriverPermit::new(Box::new(ReleaseOnDrop(Arc::clone(&released)))); - - assert!(!released.load(Ordering::SeqCst)); - drop(permit); - assert!(released.load(Ordering::SeqCst)); - } - - #[test] - fn contexts_carry_a_limiter() { - let store = StorageDriverStoreContext { - target: None, - limiter: Arc::new(ImmediateLimiter), - }; - let retrieve = StorageDriverRetrieveContext { - limiter: Arc::new(ImmediateLimiter), - }; - - let payload = Payload::default(); - let claim = StorageDriverClaim::new(HashMap::new()); - // The fake limiter is always ready, so the permits resolve without an executor. - let _store_permit = store - .limiter - .acquire(&payload) - .now_or_never() - .expect("ready") - .unwrap(); - let _retrieve_permit = retrieve - .limiter - .acquire(&claim) - .now_or_never() - .expect("ready") - .unwrap(); } #[test] diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index 6f38afab0..2473c386e 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -129,4 +129,4 @@ futures-util = { version = "0.3", default-features = false } opentelemetry_sdk = { version = "0.32", features = ["testing"] } rstest = "0.27" tempfile = "3.21" -tokio = { version = "1.47", features = ["macros", "rt"] } +tokio = { version = "1.47", features = ["macros", "rt", "sync", "time"] } diff --git a/crates/common/src/external_storage_references.rs b/crates/common/src/external_storage_references.rs index 8863d63e7..1ca208ffb 100644 --- a/crates/common/src/external_storage_references.rs +++ b/crates/common/src/external_storage_references.rs @@ -22,18 +22,22 @@ const REFERENCE_ENCODING: &str = "json/protobuf"; // `reference_message_type_matches_the_proto` instead of silently changing what we put on the wire. const REFERENCE_MESSAGE_TYPE: &str = "temporal.api.sdk.v1.ExternalStorageReference"; -/// Whether the given payload is a reference to externally stored data. -pub(crate) fn is_reference(payload: &Payload) -> bool { +/// Whether the given payload is an external storage reference and so has to be retrieved. Reads +/// metadata only, so one that needs no retrieval is not touched. +/// +/// Both halves are required: `json/protobuf` on its own is an ordinary encoding for any protobuf +/// payload. +pub(crate) fn is_plain_reference(payload: &Payload) -> bool { payload.metadata.get("encoding").map(Vec::as_slice) == Some(REFERENCE_ENCODING.as_bytes()) && payload.metadata.get("messageType").map(Vec::as_slice) == Some(REFERENCE_MESSAGE_TYPE.as_bytes()) } -/// Parse the reference from the given payload, which must be one. +/// Parse the reference from the given payload, which must be a plain reference. pub(crate) fn parse_reference( payload: &Payload, ) -> Result { - if !is_reference(payload) { + if !is_plain_reference(payload) { return Err(ExternalStorageError::InvalidReference( "payload is not an external storage reference".into(), )); @@ -101,7 +105,7 @@ mod tests { Some(REFERENCE_ENCODING.as_bytes()) ); assert_eq!(payload.external_payloads[0].size_bytes, 385); - assert!(is_reference(&payload)); + assert!(is_plain_reference(&payload)); let reference = parse_reference(&payload).unwrap(); assert_eq!(reference.driver_name, "aws.s3driver"); @@ -143,6 +147,52 @@ mod tests { ); } + #[test] + fn the_right_encoding_with_another_message_type_is_not_a_reference() { + let payload = Payload { + metadata: HashMap::from([ + ("encoding".to_owned(), b"json/protobuf".to_vec()), + ( + "messageType".to_owned(), + b"temporal.api.common.v1.Payload".to_vec(), + ), + ]), + data: b"{}".to_vec(), + external_payloads: vec![], + }; + + assert!(!is_plain_reference(&payload)); + } + + #[test] + fn the_reference_message_type_with_another_encoding_is_not_a_reference() { + let payload = Payload { + metadata: HashMap::from([ + ("encoding".to_owned(), b"binary/protobuf".to_vec()), + ( + "messageType".to_owned(), + b"temporal.api.sdk.v1.ExternalStorageReference".to_vec(), + ), + ]), + data: b"{}".to_vec(), + external_payloads: vec![], + }; + + assert!(!is_plain_reference(&payload)); + } + + #[test] + fn a_created_reference_is_detected() { + let payload = create_reference_payload( + "mem", + HashMap::from([("key".to_owned(), "k".to_owned())]), + 10, + ) + .unwrap(); + + assert!(is_plain_reference(&payload)); + } + #[test] fn ordinary_payload_is_not_a_reference() { let payload = Payload { @@ -151,7 +201,7 @@ mod tests { external_payloads: vec![], }; - assert!(!is_reference(&payload)); + assert!(!is_plain_reference(&payload)); assert!(matches!( parse_reference(&payload), Err(ExternalStorageError::InvalidReference(_)) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs new file mode 100644 index 000000000..3dcd281ee --- /dev/null +++ b/crates/common/src/external_storage_runner.rs @@ -0,0 +1,1081 @@ +//! Stores large payloads in external storage and fetches them back. +//! +//! Outbound, a payload at or above the configured size threshold is handed to a storage driver and +//! replaced by a reference to what the driver stored. Inbound, references are swapped back for the +//! payloads they point at. [`ExternalStorageRunner`] runs either direction over a whole message, +//! visiting every payload-bearing field except search attributes. +//! +//! References are stored and read in plain form, so a reference's claim data (its driver name and +//! claim keys) is written to history unencrypted. The payload codec is not applied to references, +//! and the configured concurrency limits are not enforced. + +// Crate-private until a worker or client calls it. +#![allow(dead_code)] + +use futures::future; +use std::{collections::HashMap, sync::Arc}; + +use prost::Message; +use temporalio_common_wasm::data_converters::{ + ExternalStorage, ExternalStorageError, PayloadConversionError, StorageDriver, + StorageDriverClaim, StorageDriverLimiter, StorageDriverRetrieveContext, + StorageDriverSelectContext, StorageDriverStoreContext, StorageDriverTargetInfo, +}; + +use crate::{ + external_storage_references::{create_reference_payload, is_plain_reference, parse_reference}, + payload_visitor::{ + AsyncPayloadVisitor, PayloadField, PayloadFieldData, PayloadVisitable, should_transform, + }, + protos::temporal::api::common::v1::Payload, +}; + +/// Drives external storage for the payloads reachable from one activation or request. +/// +/// A worker completion holds many commands that offload to different targets, so the caller +/// matches on the command and calls `store_in` once per command with that command's target. A +/// uniform message, such as a client request, is a single call. +pub(crate) struct ExternalStorageRunner<'a> { + storage: &'a ExternalStorage, + store_limiter: &'a Arc>, + retrieve_limiter: &'a Arc>, +} + +impl<'a> ExternalStorageRunner<'a> { + pub(crate) fn new( + storage: &'a ExternalStorage, + store_limiter: &'a Arc>, + retrieve_limiter: &'a Arc>, + ) -> Self { + Self { + storage, + store_limiter, + retrieve_limiter, + } + } + + /// Replace every eligible payload in a message with a reference. + pub(crate) async fn store_in( + &self, + message: &mut M, + target: Option<&StorageDriverTargetInfo>, + ) -> Result<(), ExternalStorageError> { + let mut store = StorePayloadsVisitor { + storage: self.storage, + target, + limiter: self.store_limiter, + error: None, + }; + let walked = message.visit_payloads_mut(&mut store).await; + finish(store.error, walked) + } + + /// Replace every reference in a message with the payload it points at. + pub(crate) async fn retrieve_in( + &self, + message: &mut M, + ) -> Result<(), ExternalStorageError> { + let mut retrieve = RetrievePayloadsVisitor { + storage: self.storage, + limiter: self.retrieve_limiter, + error: None, + }; + let walked = message.visit_payloads_mut(&mut retrieve).await; + finish(retrieve.error, walked) + } +} + +/// Fail if any payload needs retrieval. For a caller with no external storage configured, which +/// must not hand a reference to application code as though it were the application's own data. +pub(crate) fn reject_references(payloads: &[Payload]) -> Result<(), ExternalStorageError> { + if payloads.iter().any(is_plain_reference) { + return Err(ExternalStorageError::NotConfigured); + } + Ok(()) +} + +/// Replace every reference with the payload the driver returns for it. +/// +/// Each retrieved payload is left exactly as the driver returned it: the form `store` was given, +/// so it flows through the rest of the inbound path like any inline payload. +async fn retrieve_payloads( + payloads: &mut [Payload], + storage: &ExternalStorage, + limiter: &Arc>, +) -> Result<(), ExternalStorageError> { + let mut by_driver: HashMap> = HashMap::new(); + for (index, payload) in payloads.iter().enumerate() { + if !is_plain_reference(payload) { + continue; + } + let reference = parse_reference(payload)?; + by_driver + .entry(reference.driver_name) + .or_default() + .push((index, StorageDriverClaim::new(reference.claim_data))); + } + + // Resolve every driver first, so an unregistered name fails before any request is made. + let mut work = Vec::with_capacity(by_driver.len()); + for (driver_name, group) in by_driver { + let driver = storage + .driver(&driver_name) + .ok_or_else(|| ExternalStorageError::UnregisteredDriver(driver_name.clone()))?; + work.push((driver_name, driver, group)); + } + + let retrieved_groups = + future::try_join_all(work.into_iter().map(|(driver_name, driver, group)| { + let limiter = Arc::clone(limiter); + async move { + let context = StorageDriverRetrieveContext::builder() + .limiter(limiter) + .build(); + let claims: Vec<_> = group.iter().map(|(_, claim)| claim.clone()).collect(); + let retrieved = driver.retrieve(&context, claims).await?; + if retrieved.len() != group.len() { + return Err(ExternalStorageError::DriverContract { + driver: driver_name, + requested: group.len(), + returned: retrieved.len(), + }); + } + Ok::<_, ExternalStorageError>( + group + .into_iter() + .map(|(index, _)| index) + .zip(retrieved) + .collect::>(), + ) + } + })) + .await?; + + for group in retrieved_groups { + for (index, payload) in group { + payloads[index] = payload; + } + } + Ok(()) +} + +/// The visitor reports only that it failed, so the typed error it captured is preferred over the +/// placeholder that stopped the walk. +fn finish( + captured: Option, + walked: Result<(), PayloadConversionError>, +) -> Result<(), ExternalStorageError> { + if let Some(error) = captured { + return Err(error); + } + walked.map_err(|e| ExternalStorageError::Walk(Box::new(e))) +} + +fn stop_walk() -> PayloadConversionError { + PayloadConversionError::EncodingError("external storage pass failed".into()) +} + +struct StorePayloadsVisitor<'a> { + storage: &'a ExternalStorage, + target: Option<&'a StorageDriverTargetInfo>, + limiter: &'a Arc>, + error: Option, +} +impl AsyncPayloadVisitor for StorePayloadsVisitor<'_> { + fn visit<'a>( + &'a mut self, + field: PayloadField<'a>, + ) -> futures::future::BoxFuture<'a, Result<(), PayloadConversionError>> { + Box::pin(async move { + if !should_transform(field.path) { + return Ok(()); + } + let result = match field.data { + PayloadFieldData::Single(payload) => { + store_payloads( + std::slice::from_mut(payload), + self.storage, + self.target, + self.limiter, + ) + .await + } + PayloadFieldData::Repeated(payloads) => { + store_payloads(payloads, self.storage, self.target, self.limiter).await + } + PayloadFieldData::Payloads(message) => { + store_payloads( + &mut message.payloads, + self.storage, + self.target, + self.limiter, + ) + .await + } + }; + match result { + Ok(()) => Ok(()), + Err(e) => { + self.error = Some(e); + Err(stop_walk()) + } + } + }) + } +} + +struct RetrievePayloadsVisitor<'a> { + storage: &'a ExternalStorage, + limiter: &'a Arc>, + error: Option, +} + +impl AsyncPayloadVisitor for RetrievePayloadsVisitor<'_> { + fn visit<'a>( + &'a mut self, + field: PayloadField<'a>, + ) -> futures::future::BoxFuture<'a, Result<(), PayloadConversionError>> { + Box::pin(async move { + if !should_transform(field.path) { + return Ok(()); + } + let result = match field.data { + PayloadFieldData::Single(payload) => { + retrieve_payloads(std::slice::from_mut(payload), self.storage, self.limiter) + .await + } + PayloadFieldData::Repeated(payloads) => { + retrieve_payloads(payloads, self.storage, self.limiter).await + } + PayloadFieldData::Payloads(message) => { + retrieve_payloads(&mut message.payloads, self.storage, self.limiter).await + } + }; + match result { + Ok(()) => Ok(()), + Err(e) => { + self.error = Some(e); + Err(stop_walk()) + } + } + }) + } +} + +/// One driver and the payloads it takes: each payload's index, the payload, and its pre-offload +/// proto size. +type DriverGroup = (Arc, Vec<(usize, Payload, i64)>); + +/// Replace every eligible payload with a reference, marking each one it creates when writing is +/// enabled so [`encode_references`] can find them without re-examining encodings. +async fn store_payloads( + payloads: &mut [Payload], + storage: &ExternalStorage, + target: Option<&StorageDriverTargetInfo>, + limiter: &Arc>, +) -> Result<(), ExternalStorageError> { + let select_context = StorageDriverSelectContext::builder() + .maybe_target(target.cloned()) + .build(); + + let mut by_driver: HashMap = HashMap::new(); + for (index, payload) in payloads.iter().enumerate() { + let size = payload.encoded_len(); + if size < storage.payload_size_threshold() { + continue; + } + let Some(driver) = storage + .driver_selector() + .select_driver(&select_context, payload) + else { + continue; + }; + // Checked by identity, so a driver built to look like a registered one is rejected. + if !storage.drivers().iter().any(|d| Arc::ptr_eq(d, &driver)) { + return Err(ExternalStorageError::UnregisteredSelection( + driver.name().to_owned(), + )); + } + by_driver + .entry(driver.name().to_owned()) + .or_insert_with(|| (Arc::clone(&driver), Vec::new())) + .1 + .push((index, payload.clone(), size as i64)); + } + + let stored = + future::try_join_all(by_driver.into_iter().map(|(driver_name, (driver, group))| { + let limiter = Arc::clone(limiter); + let target = target.cloned(); + async move { + let context = StorageDriverStoreContext::builder() + .maybe_target(target) + .limiter(limiter) + .build(); + let to_store: Vec<_> = group + .iter() + .map(|(_, payload, _)| payload.clone()) + .collect(); + let claims = driver.store(&context, to_store).await?; + if claims.len() != group.len() { + return Err(ExternalStorageError::DriverContract { + driver: driver_name, + requested: group.len(), + returned: claims.len(), + }); + } + let mut replaced = Vec::with_capacity(group.len()); + for ((index, _, size), claim) in group.into_iter().zip(claims) { + let reference = create_reference_payload(&driver_name, claim.claim_data, size)?; + replaced.push((index, reference)); + } + Ok::<_, ExternalStorageError>(replaced) + } + })) + .await?; + + for group in stored { + for (index, reference) in group { + payloads[index] = reference; + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + external_storage_references::create_reference_payload, + protos::temporal::api::{ + common::v1::{Header, Memo, Payloads, SearchAttributes}, + workflowservice::v1::SignalWithStartWorkflowExecutionRequest, + }, + }; + use futures::future::{self, BoxFuture}; + use std::sync::Mutex; + use temporalio_common_wasm::data_converters::{ + StorageDriver, StorageDriverPermit, StorageDriverSelector, + }; + + /// Echoes each claim back as a payload (and each payload as a claim), recording the batch + /// sizes it was called with. The in-flight counters let one test observe overlap. + struct RecordingDriver { + name: &'static str, + batches: Arc>>, + // Released once every driver sharing it is in a request, so one test can prove overlap. + barrier: Arc, + } + + impl RecordingDriver { + // The common case: a one-party barrier that releases at once, so the driver does not wait. + // The overlap test shares a wider barrier across drivers instead. + fn new(name: &'static str, batches: Arc>>) -> Self { + Self { + name, + batches, + barrier: Arc::new(tokio::sync::Barrier::new(1)), + } + } + } + + impl StorageDriver for RecordingDriver { + fn name(&self) -> &str { + self.name + } + fn driver_type(&self) -> &str { + "test.recordingdriver" + } + fn store( + &self, + _context: &temporalio_common_wasm::data_converters::StorageDriverStoreContext, + payloads: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + self.batches.lock().unwrap().push(payloads.len()); + let barrier = Arc::clone(&self.barrier); + Box::pin(async move { + // Block until every driver sharing this barrier is also in a request, proving the + // requests overlap. A one-party barrier releases immediately. + barrier.wait().await; + Ok(payloads + .into_iter() + .map(|p| { + StorageDriverClaim::new(HashMap::from([( + "key".to_owned(), + String::from_utf8(p.data).unwrap(), + )])) + }) + .collect()) + }) + } + fn retrieve( + &self, + _context: &StorageDriverRetrieveContext, + claims: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + self.batches.lock().unwrap().push(claims.len()); + let barrier = Arc::clone(&self.barrier); + Box::pin(async move { + // Block until every driver sharing this barrier is also in a request, proving the + // requests overlap. A one-party barrier releases immediately. + barrier.wait().await; + Ok(claims + .into_iter() + .map(|c| Payload { + metadata: HashMap::from([("encoding".to_owned(), b"json/plain".to_vec())]), + data: c.claim_data["key"].clone().into_bytes(), + external_payloads: vec![], + }) + .collect()) + }) + } + } + + struct PassThroughLimiter; + impl StorageDriverLimiter for PassThroughLimiter { + fn acquire<'a>( + &'a self, + _item: &'a Item, + ) -> BoxFuture<'a, Result> { + Box::pin(future::ready(Ok(StorageDriverPermit::new(Box::new(()))))) + } + } + + fn storage(name: &'static str, batches: Arc>>) -> ExternalStorage { + ExternalStorage::new(Arc::new(RecordingDriver::new(name, batches))).unwrap() + } + + /// Never selects a driver. Retrieval ignores the selector, so this stands in where + /// `with_selector` needs an argument it will not use. + struct NeverSelects; + impl temporalio_common_wasm::data_converters::StorageDriverSelector for NeverSelects { + fn select_driver( + &self, + _context: &temporalio_common_wasm::data_converters::StorageDriverSelectContext, + _payload: &Payload, + ) -> Option> { + None + } + } + + fn limiter() -> Arc> { + Arc::new(PassThroughLimiter) + } + + fn reference(driver: &str, key: &str) -> Payload { + create_reference_payload( + driver, + HashMap::from([("key".to_owned(), key.to_owned())]), + 99, + ) + .unwrap() + } + + fn inline(data: &str) -> Payload { + Payload { + metadata: HashMap::from([("encoding".to_owned(), b"json/plain".to_vec())]), + data: data.as_bytes().to_vec(), + external_payloads: vec![], + } + } + + #[tokio::test] + async fn references_are_replaced_and_inline_payloads_untouched() { + let batches = Arc::new(Mutex::new(vec![])); + let storage = storage("mem", Arc::clone(&batches)); + let mut payloads = vec![ + inline("first"), + reference("mem", "a"), + inline("third"), + reference("mem", "b"), + ]; + + retrieve_payloads(&mut payloads, &storage, &limiter()) + .await + .unwrap(); + + assert_eq!(payloads[0].data, b"first", "inline payload untouched"); + assert_eq!(payloads[1].data, b"a"); + assert_eq!(payloads[2].data, b"third", "inline payload untouched"); + assert_eq!(payloads[3].data, b"b"); + // One call for the driver, carrying both of its claims, not one call per payload. + assert_eq!(*batches.lock().unwrap(), vec![2]); + } + + #[tokio::test] + async fn two_drivers_retrieve_concurrently() { + // Both drivers share this barrier, so each retrieve blocks until the other is also + // running — only possible if they overlap. The timeout fails a sequential regression + // instead of hanging on it. + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let make = |name: &'static str| -> Arc { + Arc::new(RecordingDriver { + name, + batches: Arc::new(Mutex::new(vec![])), + barrier: Arc::clone(&barrier), + }) + }; + let storage = + ExternalStorage::with_selector(vec![make("a"), make("b")], Arc::new(NeverSelects)) + .unwrap(); + let mut payloads = vec![reference("a", "one"), reference("b", "two")]; + + tokio::time::timeout( + std::time::Duration::from_secs(5), + retrieve_payloads(&mut payloads, &storage, &limiter()), + ) + .await + .expect("drivers did not run concurrently") + .unwrap(); + + assert_eq!(payloads[0].data, b"one"); + assert_eq!(payloads[1].data, b"two"); + } + + struct StoringDriver { + name: &'static str, + stored: Arc>>, + } + + impl StorageDriver for StoringDriver { + fn name(&self) -> &str { + self.name + } + fn driver_type(&self) -> &str { + "test.storingdriver" + } + fn store( + &self, + _context: &StorageDriverStoreContext, + payloads: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + let stored = Arc::clone(&self.stored); + Box::pin(async move { + let mut held = stored.lock().unwrap(); + let claims = payloads + .into_iter() + .map(|p| { + held.push(p); + StorageDriverClaim::new(HashMap::from([( + "key".to_owned(), + format!("k{}", held.len() - 1), + )])) + }) + .collect(); + Ok(claims) + }) + } + fn retrieve( + &self, + _context: &StorageDriverRetrieveContext, + _claims: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + Box::pin(future::ready(Ok(vec![]))) + } + } + + /// Returns an empty batch regardless of input, so store and retrieve both violate the + /// one-result-per-input contract. + struct MiscountingDriver { + name: &'static str, + } + impl StorageDriver for MiscountingDriver { + fn name(&self) -> &str { + self.name + } + fn driver_type(&self) -> &str { + "test.miscounting" + } + fn store( + &self, + _context: &StorageDriverStoreContext, + _payloads: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + Box::pin(future::ready(Ok(vec![]))) + } + fn retrieve( + &self, + _context: &StorageDriverRetrieveContext, + _claims: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + Box::pin(future::ready(Ok(vec![]))) + } + } + + /// Routes each payload to the driver keyed by its first data byte, declining any other. + struct SelectByFirstByte(HashMap>); + impl StorageDriverSelector for SelectByFirstByte { + fn select_driver( + &self, + _context: &StorageDriverSelectContext, + payload: &Payload, + ) -> Option> { + payload + .data + .first() + .and_then(|b| self.0.get(b)) + .map(Arc::clone) + } + } + + /// Returns the one driver it was built with, whether or not that driver is registered. + struct FixedSelector(Arc); + impl StorageDriverSelector for FixedSelector { + fn select_driver( + &self, + _context: &StorageDriverSelectContext, + _payload: &Payload, + ) -> Option> { + Some(Arc::clone(&self.0)) + } + } + + fn store_limiter() -> Arc> { + Arc::new(PassThroughLimiter) + } + + fn filled(byte: u8, bytes: usize) -> Payload { + Payload { + metadata: HashMap::from([("encoding".to_owned(), b"json/plain".to_vec())]), + data: vec![byte; bytes], + external_payloads: vec![], + } + } + + fn sized(bytes: usize) -> Payload { + filled(b'x', bytes) + } + + fn storage_with_threshold( + threshold: usize, + stored: Arc>>, + ) -> ExternalStorage { + ExternalStorage::new(Arc::new(StoringDriver { + name: "mem", + stored, + })) + .unwrap() + .with_payload_size_threshold(threshold) + } + + #[tokio::test] + async fn a_payload_at_the_threshold_is_offloaded_and_a_smaller_one_is_not() { + let small = sized(8); + let at_threshold = sized(64); + let threshold = at_threshold.encoded_len(); + assert!(small.encoded_len() < threshold); + let stored = Arc::new(Mutex::new(vec![])); + let storage = storage_with_threshold(threshold, Arc::clone(&stored)); + let mut payloads = vec![small.clone(), at_threshold.clone()]; + + store_payloads(&mut payloads, &storage, None, &store_limiter()) + .await + .unwrap(); + + assert_eq!(payloads[0], small, "below the threshold, left inline"); + assert!( + crate::external_storage_references::is_plain_reference(&payloads[1]), + "only the one at the threshold was offloaded" + ); + // What the driver was handed is byte-identical to what history would have held inline. + assert_eq!(stored.lock().unwrap()[0], at_threshold); + } + + #[tokio::test] + async fn the_reference_records_the_pre_offload_size() { + let stored = Arc::new(Mutex::new(vec![])); + let storage = storage_with_threshold(1, Arc::clone(&stored)); + let payload = sized(100); + let size = payload.encoded_len() as i64; + let mut payloads = vec![payload]; + + store_payloads(&mut payloads, &storage, None, &store_limiter()) + .await + .unwrap(); + + assert_eq!(payloads[0].external_payloads[0].size_bytes, size); + } + + #[tokio::test] + async fn a_driver_the_selector_invented_is_rejected() { + let stored = Arc::new(Mutex::new(vec![])); + let registered: Arc = Arc::new(StoringDriver { + name: "mem", + stored: Arc::clone(&stored), + }); + // Same name, different instance. + let impostor: Arc = Arc::new(StoringDriver { + name: "mem", + stored: Arc::clone(&stored), + }); + let storage = + ExternalStorage::with_selector(vec![registered], Arc::new(FixedSelector(impostor))) + .unwrap() + .with_payload_size_threshold(1); + let mut payloads = vec![sized(50)]; + + let err = store_payloads(&mut payloads, &storage, None, &store_limiter()) + .await + .unwrap_err(); + + assert!(matches!( + err, + ExternalStorageError::UnregisteredSelection(name) if name == "mem" + )); + } + + async fn a_store_driver_returning_the_wrong_number_of_claims_is_an_error() { + let driver: Arc = Arc::new(MiscountingDriver { name: "mem" }); + let storage = ExternalStorage::with_selector( + vec![Arc::clone(&driver)], + Arc::new(FixedSelector(driver)), + ) + .unwrap() + .with_payload_size_threshold(1); + let mut payloads = vec![sized(50)]; + + let err = store_payloads(&mut payloads, &storage, None, &store_limiter()) + .await + .unwrap_err(); + + assert!(matches!( + err, + ExternalStorageError::DriverContract { + driver, + requested: 1, + returned: 0, + } if driver == "mem" + )); + } + + #[tokio::test] + async fn a_retrieve_driver_returning_the_wrong_number_of_payloads_is_an_error() { + let storage = ExternalStorage::new(Arc::new(MiscountingDriver { name: "mem" })).unwrap(); + let mut payloads = vec![reference("mem", "a")]; + + let err = retrieve_payloads(&mut payloads, &storage, &limiter()) + .await + .unwrap_err(); + + assert!(matches!( + err, + ExternalStorageError::DriverContract { + driver, + requested: 1, + returned: 0, + } if driver == "mem" + )); + } + + #[tokio::test] + async fn a_payload_the_selector_declines_stays_inline() { + let stored = Arc::new(Mutex::new(vec![])); + let driver: Arc = Arc::new(StoringDriver { name: "a", stored }); + // Routes a payload beginning with b'a' and declines one beginning with b'z'. + let storage = ExternalStorage::with_selector( + vec![Arc::clone(&driver)], + Arc::new(SelectByFirstByte(HashMap::from([(b'a', driver)]))), + ) + .unwrap() + .with_payload_size_threshold(1); + let declined = filled(b'z', 50); + let mut payloads = vec![filled(b'a', 50), declined.clone()]; + + store_payloads(&mut payloads, &storage, None, &store_limiter()) + .await + .unwrap(); + + assert!( + is_plain_reference(&payloads[0]), + "the routed payload was offloaded" + ); + assert_eq!( + payloads[1], declined, + "the declined payload was left inline" + ); + } + + #[tokio::test] + async fn payloads_are_grouped_by_driver_and_restored_in_order() { + let batches_a = Arc::new(Mutex::new(vec![])); + let batches_b = Arc::new(Mutex::new(vec![])); + let make = + |name: &'static str, batches: &Arc>>| -> Arc { + Arc::new(RecordingDriver::new(name, Arc::clone(batches))) + }; + let a = make("a", &batches_a); + let b = make("b", &batches_b); + let storage = ExternalStorage::with_selector( + vec![Arc::clone(&a), Arc::clone(&b)], + Arc::new(SelectByFirstByte(HashMap::from([(b'a', a), (b'b', b)]))), + ) + .unwrap() + .with_payload_size_threshold(1); + // Interleaved, so grouping by driver and restoring to original order both matter. + let mut payloads = vec![ + filled(b'a', 50), + filled(b'b', 50), + filled(b'a', 50), + filled(b'b', 50), + ]; + + store_payloads(&mut payloads, &storage, None, &store_limiter()) + .await + .unwrap(); + + for p in &payloads { + assert!(is_plain_reference(p)); + } + // One call per driver, each carrying that driver's two payloads. + assert_eq!(*batches_a.lock().unwrap(), vec![2]); + assert_eq!(*batches_b.lock().unwrap(), vec![2]); + // References landed back on the positions they came from. + let driver_at = |i: usize| parse_reference(&payloads[i]).unwrap().driver_name; + assert_eq!(driver_at(0), "a"); + assert_eq!(driver_at(1), "b"); + assert_eq!(driver_at(2), "a"); + assert_eq!(driver_at(3), "b"); + } + + #[tokio::test] + async fn a_threshold_of_zero_offloads_every_payload() { + let stored = Arc::new(Mutex::new(vec![])); + let storage = storage_with_threshold(0, Arc::clone(&stored)); + let mut payloads = vec![sized(1), sized(8)]; + + store_payloads(&mut payloads, &storage, None, &store_limiter()) + .await + .unwrap(); + + assert!(is_plain_reference(&payloads[0])); + assert!(is_plain_reference(&payloads[1])); + assert_eq!(stored.lock().unwrap().len(), 2); + } + + struct RoundTripDriver { + stored: Arc>>, + } + + impl StorageDriver for RoundTripDriver { + fn name(&self) -> &str { + "mem" + } + fn driver_type(&self) -> &str { + "test.roundtrip" + } + fn store( + &self, + _context: &StorageDriverStoreContext, + payloads: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + let stored = Arc::clone(&self.stored); + Box::pin(async move { + let mut held = stored.lock().unwrap(); + Ok(payloads + .into_iter() + .map(|p| { + let key = format!("k{}", held.len()); + held.insert(key.clone(), p); + StorageDriverClaim::new(HashMap::from([("key".to_owned(), key)])) + }) + .collect()) + }) + } + fn retrieve( + &self, + _context: &StorageDriverRetrieveContext, + claims: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + let stored = Arc::clone(&self.stored); + Box::pin(async move { + let held = stored.lock().unwrap(); + claims + .into_iter() + .map(|c| { + held.get(&c.claim_data["key"]).cloned().ok_or_else(|| { + ExternalStorageError::InvalidReference( + "stored payload not found".into(), + ) + }) + }) + .collect() + }) + } + } + + #[tokio::test] + async fn a_message_makes_the_round_trip_through_the_runner() { + let stored = Arc::new(Mutex::new(HashMap::new())); + let driver: Arc = Arc::new(RoundTripDriver { + stored: Arc::clone(&stored), + }); + let storage = ExternalStorage::with_selector( + vec![Arc::clone(&driver)], + Arc::new(FixedSelector(driver)), + ) + .unwrap() + .with_payload_size_threshold(1); + let store_lim = store_limiter(); + let retrieve_lim = limiter(); + let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + let mut request = request_with_large_payloads(); + let before = request.clone(); + + runner.store_in(&mut request, None).await.unwrap(); + assert!(is_plain_reference( + &request.input.as_ref().unwrap().payloads[0] + )); + + runner.retrieve_in(&mut request).await.unwrap(); + + // Every field is back to what it was, search attributes included because they never left. + assert_eq!(request, before); + } + + #[tokio::test] + async fn the_runner_reports_a_driver_it_does_not_know() { + let batches = Arc::new(Mutex::new(vec![])); + let storage = storage("mem", Arc::clone(&batches)); + let store_lim = store_limiter(); + let retrieve_lim = limiter(); + let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + // A reference to the registered driver sits beside the unknown one. + let mut request = SignalWithStartWorkflowExecutionRequest { + input: Some(Payloads { + payloads: vec![reference("mem", "a"), reference("nope", "b")], + }), + ..Default::default() + }; + + let err = runner.retrieve_in(&mut request).await.unwrap_err(); + + assert!(matches!( + err, + ExternalStorageError::UnregisteredDriver(name) if name == "nope" + )); + // Drivers are resolved before any request, so the registered one was never called. + assert!(batches.lock().unwrap().is_empty()); + } + + // ---- over a real message graph ---- + + /// A request with a reference in every payload-bearing field the walk reaches, including + /// search attributes, which must be left alone. + fn request_with_references() -> SignalWithStartWorkflowExecutionRequest { + SignalWithStartWorkflowExecutionRequest { + input: Some(Payloads { + payloads: vec![reference("mem", "input-0"), reference("mem", "input-1")], + }), + signal_input: Some(Payloads { + payloads: vec![reference("mem", "signal-0")], + }), + memo: Some(Memo { + fields: HashMap::from([("m".to_owned(), reference("mem", "memo-0"))]), + }), + header: Some(Header { + fields: HashMap::from([("h".to_owned(), reference("mem", "header-0"))]), + }), + search_attributes: Some(SearchAttributes { + indexed_fields: HashMap::from([("s".to_owned(), reference("mem", "sa-0"))]), + }), + ..Default::default() + } + } + + /// The same request, but with payloads large enough to store rather than references. + fn request_with_large_payloads() -> SignalWithStartWorkflowExecutionRequest { + SignalWithStartWorkflowExecutionRequest { + input: Some(Payloads { + payloads: vec![filled(b'a', 500), filled(b'b', 500)], + }), + signal_input: Some(Payloads { + payloads: vec![filled(b'c', 500)], + }), + memo: Some(Memo { + fields: HashMap::from([("m".to_owned(), filled(b'd', 500))]), + }), + header: Some(Header { + fields: HashMap::from([("h".to_owned(), filled(b'e', 500))]), + }), + search_attributes: Some(SearchAttributes { + indexed_fields: HashMap::from([("s".to_owned(), filled(b'f', 500))]), + }), + ..Default::default() + } + } + + async fn run_store_pass( + request: &mut SignalWithStartWorkflowExecutionRequest, + stored: Arc>>, + ) { + let storage = storage_with_threshold(1, stored); + let store_lim = store_limiter(); + let retrieve_lim = limiter(); + let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + runner.store_in(request, None).await.unwrap(); + } + + #[tokio::test] + async fn every_other_payload_bearing_field_is_stored() { + let stored = Arc::new(Mutex::new(vec![])); + let mut request = request_with_large_payloads(); + let search_attributes = request.search_attributes.clone(); + + run_store_pass(&mut request, Arc::clone(&stored)).await; + + for payload in [ + &request.input.as_ref().unwrap().payloads[0], + &request.input.as_ref().unwrap().payloads[1], + &request.signal_input.as_ref().unwrap().payloads[0], + &request.memo.as_ref().unwrap().fields["m"], + &request.header.as_ref().unwrap().fields["h"], + ] { + assert!(is_plain_reference(payload)); + } + // The five stored payloads, and not the search attribute. + assert_eq!(stored.lock().unwrap().len(), 5); + // The server indexes search attributes literally, so they are left exactly as they were. + assert_eq!(request.search_attributes, search_attributes); + } + + #[tokio::test] + async fn every_other_payload_bearing_field_is_retrieved() { + let batches = Arc::new(Mutex::new(vec![])); + let storage = storage("mem", Arc::clone(&batches)); + let store_lim = store_limiter(); + let retrieve_lim = limiter(); + let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + let mut request = request_with_references(); + let search_attributes = request.search_attributes.clone(); + + runner.retrieve_in(&mut request).await.unwrap(); + + // Search attributes are indexed literally, so they are left untouched and still reference + // external storage afterward. + assert_eq!(request.search_attributes, search_attributes); + assert!(crate::external_storage_references::is_plain_reference( + &request.search_attributes.as_ref().unwrap().indexed_fields["s"] + )); + let input = request.input.unwrap().payloads; + assert_eq!(input[0].data, b"input-0"); + assert_eq!(input[1].data, b"input-1"); + assert_eq!(request.signal_input.unwrap().payloads[0].data, b"signal-0"); + assert_eq!(request.memo.unwrap().fields["m"].data, b"memo-0"); + assert_eq!(request.header.unwrap().fields["h"].data, b"header-0"); + + // One call per driver per payload-bearing field: input (2 claims), signal_input, memo, + // header. Search attributes are skipped, so there is no fifth call. + let mut sizes = batches.lock().unwrap().clone(); + sizes.sort_unstable(); + assert_eq!(sizes, vec![1, 1, 1, 2]); + } + + #[test] + fn rejecting_references_without_storage_is_tmprl1105() { + assert!(reject_references(&[inline("fine")]).is_ok()); + + let err = reject_references(&[reference("mem", "a")]).unwrap_err(); + assert!(matches!(err, ExternalStorageError::NotConfigured)); + assert!(err.to_string().contains("TMPRL1105")); + } +} diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 4809fb15c..dfa784d51 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -11,6 +11,8 @@ extern crate tracing; pub mod envconfig; #[cfg(feature = "experimental")] pub(crate) mod external_storage_references; +#[cfg(feature = "experimental")] +pub(crate) mod external_storage_runner; #[doc(hidden)] pub mod fsm_trait; #[doc(hidden)] diff --git a/crates/common/src/payload_visitor.rs b/crates/common/src/payload_visitor.rs index f9b1845a3..d9ae4126c 100644 --- a/crates/common/src/payload_visitor.rs +++ b/crates/common/src/payload_visitor.rs @@ -51,14 +51,15 @@ pub trait PayloadVisitable: Send { ) -> BoxFuture<'a, Result<(), PayloadConversionError>>; } -/// Check if a field path represents search attributes that should not be encoded. +/// Check if a field path represents search attributes that must not be transformed. /// Search attributes must remain server-readable for indexing. fn is_search_attributes_path(path: &str) -> bool { // All search attributes go through the SearchAttributes message which has indexed_fields path.contains("SearchAttributes.indexed_fields") } -fn should_encode(path: &str) -> bool { +/// Whether a field at this path may be transformed: encoded, decoded, or offloaded. +pub(crate) fn should_transform(path: &str) -> bool { !is_search_attributes_path(path) } @@ -74,7 +75,7 @@ impl AsyncPayloadVisitor for EncodeVisitor<'_> { field: PayloadField<'a>, ) -> BoxFuture<'a, Result<(), PayloadConversionError>> { Box::pin(async move { - if !should_encode(field.path) { + if !should_transform(field.path) { return Ok(()); } match field.data { @@ -117,7 +118,7 @@ impl AsyncPayloadVisitor for DecodeVisitor<'_> { field: PayloadField<'a>, ) -> BoxFuture<'a, Result<(), PayloadConversionError>> { Box::pin(async move { - if !should_encode(field.path) { + if !should_transform(field.path) { return Ok(()); } match field.data { From ed747ae1de09224894117011096cf6be5d2988d3 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:42 -0700 Subject: [PATCH 02/15] test: run the miscounting store driver test --- crates/common/src/external_storage_runner.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 3dcd281ee..20f4460bc 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -724,6 +724,7 @@ mod tests { )); } + #[tokio::test] async fn a_store_driver_returning_the_wrong_number_of_claims_is_an_error() { let driver: Arc = Arc::new(MiscountingDriver { name: "mem" }); let storage = ExternalStorage::with_selector( From 8c33b0b0626901d0d13b4c89e630127fda8d407f Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:42 -0700 Subject: [PATCH 03/15] docs: define a plain external storage reference --- crates/common/src/external_storage_references.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/crates/common/src/external_storage_references.rs b/crates/common/src/external_storage_references.rs index 1ca208ffb..4d5c89943 100644 --- a/crates/common/src/external_storage_references.rs +++ b/crates/common/src/external_storage_references.rs @@ -22,11 +22,14 @@ const REFERENCE_ENCODING: &str = "json/protobuf"; // `reference_message_type_matches_the_proto` instead of silently changing what we put on the wire. const REFERENCE_MESSAGE_TYPE: &str = "temporal.api.sdk.v1.ExternalStorageReference"; -/// Whether the given payload is an external storage reference and so has to be retrieved. Reads -/// metadata only, so one that needs no retrieval is not touched. +/// Whether the given payload is a plain external storage reference and so has to be retrieved. +/// Reads metadata only, so one that needs no retrieval is not touched. /// -/// Both halves are required: `json/protobuf` on its own is an ordinary encoding for any protobuf -/// payload. +/// Plain means the reference exactly as [`create_reference_payload`] built it, before any payload +/// codec has run over it. A codec rewrites the metadata, so an encoded reference is not recognized +/// here and has to be decoded first. +/// +/// Both the `encoding` and the `messageType` metadata must match. pub(crate) fn is_plain_reference(payload: &Payload) -> bool { payload.metadata.get("encoding").map(Vec::as_slice) == Some(REFERENCE_ENCODING.as_bytes()) && payload.metadata.get("messageType").map(Vec::as_slice) From 7d0cafa4e86c875b8d4f8b592fbcf79553f0e4de Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:43 -0700 Subject: [PATCH 04/15] docs: drop redundant note on retrieved payloads --- crates/common/src/external_storage_runner.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 20f4460bc..c38d40da6 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -95,9 +95,6 @@ pub(crate) fn reject_references(payloads: &[Payload]) -> Result<(), ExternalStor } /// Replace every reference with the payload the driver returns for it. -/// -/// Each retrieved payload is left exactly as the driver returned it: the form `store` was given, -/// so it flows through the rest of the inbound path like any inline payload. async fn retrieve_payloads( payloads: &mut [Payload], storage: &ExternalStorage, From c5920cc378cff882134188c52ce7ace056b9fe7a Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:43 -0700 Subject: [PATCH 05/15] docs: remove mention of nonexistent encode_references --- crates/common/src/external_storage_runner.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index c38d40da6..953c8e433 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -263,8 +263,7 @@ impl AsyncPayloadVisitor for RetrievePayloadsVisitor<'_> { /// proto size. type DriverGroup = (Arc, Vec<(usize, Payload, i64)>); -/// Replace every eligible payload with a reference, marking each one it creates when writing is -/// enabled so [`encode_references`] can find them without re-examining encodings. +/// Replace every eligible payload with a reference. async fn store_payloads( payloads: &mut [Payload], storage: &ExternalStorage, From ca76d318332fc227bde5a33be296c19033175c3d Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:43 -0700 Subject: [PATCH 06/15] docs: say why search attributes are not offloaded --- crates/common/src/external_storage_runner.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 953c8e433..299abff24 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -3,7 +3,8 @@ //! Outbound, a payload at or above the configured size threshold is handed to a storage driver and //! replaced by a reference to what the driver stored. Inbound, references are swapped back for the //! payloads they point at. [`ExternalStorageRunner`] runs either direction over a whole message, -//! visiting every payload-bearing field except search attributes. +//! visiting every payload-bearing field except search attributes, which the server has to be able +//! to read. //! //! References are stored and read in plain form, so a reference's claim data (its driver name and //! claim keys) is written to history unencrypted. The payload codec is not applied to references, From b6f0d379d02f00a2113878794f9ea10340286e3e Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:43 -0700 Subject: [PATCH 07/15] refactor: have the storage runner own its limiters --- crates/common/src/external_storage_runner.rs | 28 +++++++------------- 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 299abff24..dc243f65f 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -38,15 +38,15 @@ use crate::{ /// uniform message, such as a client request, is a single call. pub(crate) struct ExternalStorageRunner<'a> { storage: &'a ExternalStorage, - store_limiter: &'a Arc>, - retrieve_limiter: &'a Arc>, + store_limiter: Arc>, + retrieve_limiter: Arc>, } impl<'a> ExternalStorageRunner<'a> { pub(crate) fn new( storage: &'a ExternalStorage, - store_limiter: &'a Arc>, - retrieve_limiter: &'a Arc>, + store_limiter: Arc>, + retrieve_limiter: Arc>, ) -> Self { Self { storage, @@ -64,7 +64,7 @@ impl<'a> ExternalStorageRunner<'a> { let mut store = StorePayloadsVisitor { storage: self.storage, target, - limiter: self.store_limiter, + limiter: &self.store_limiter, error: None, }; let walked = message.visit_payloads_mut(&mut store).await; @@ -78,7 +78,7 @@ impl<'a> ExternalStorageRunner<'a> { ) -> Result<(), ExternalStorageError> { let mut retrieve = RetrievePayloadsVisitor { storage: self.storage, - limiter: self.retrieve_limiter, + limiter: &self.retrieve_limiter, error: None, }; let walked = message.visit_payloads_mut(&mut retrieve).await; @@ -913,9 +913,7 @@ mod tests { ) .unwrap() .with_payload_size_threshold(1); - let store_lim = store_limiter(); - let retrieve_lim = limiter(); - let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + let runner = ExternalStorageRunner::new(&storage, store_limiter(), limiter()); let mut request = request_with_large_payloads(); let before = request.clone(); @@ -934,9 +932,7 @@ mod tests { async fn the_runner_reports_a_driver_it_does_not_know() { let batches = Arc::new(Mutex::new(vec![])); let storage = storage("mem", Arc::clone(&batches)); - let store_lim = store_limiter(); - let retrieve_lim = limiter(); - let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + let runner = ExternalStorageRunner::new(&storage, store_limiter(), limiter()); // A reference to the registered driver sits beside the unknown one. let mut request = SignalWithStartWorkflowExecutionRequest { input: Some(Payloads { @@ -1007,9 +1003,7 @@ mod tests { stored: Arc>>, ) { let storage = storage_with_threshold(1, stored); - let store_lim = store_limiter(); - let retrieve_lim = limiter(); - let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + let runner = ExternalStorageRunner::new(&storage, store_limiter(), limiter()); runner.store_in(request, None).await.unwrap(); } @@ -1040,9 +1034,7 @@ mod tests { async fn every_other_payload_bearing_field_is_retrieved() { let batches = Arc::new(Mutex::new(vec![])); let storage = storage("mem", Arc::clone(&batches)); - let store_lim = store_limiter(); - let retrieve_lim = limiter(); - let runner = ExternalStorageRunner::new(&storage, &store_lim, &retrieve_lim); + let runner = ExternalStorageRunner::new(&storage, store_limiter(), limiter()); let mut request = request_with_references(); let search_attributes = request.search_attributes.clone(); From 70868eb106bce837594eae1dbbaa0111f1b09f05 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:43 -0700 Subject: [PATCH 08/15] perf: look up the selected driver by name --- crates/common/src/external_storage_runner.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index dc243f65f..8f630d572 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -288,7 +288,10 @@ async fn store_payloads( continue; }; // Checked by identity, so a driver built to look like a registered one is rejected. - if !storage.drivers().iter().any(|d| Arc::ptr_eq(d, &driver)) { + if !storage + .driver(driver.name()) + .is_some_and(|registered| Arc::ptr_eq(registered, &driver)) + { return Err(ExternalStorageError::UnregisteredSelection( driver.name().to_owned(), )); From 40de2a19a9c5724b6df9cbda6d5bcda3a855fccc Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:44 -0700 Subject: [PATCH 09/15] perf: await driver requests with FuturesUnordered --- crates/common/src/external_storage_runner.rs | 53 +++++++++----------- 1 file changed, 23 insertions(+), 30 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 8f630d572..33b3712e5 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -13,7 +13,7 @@ // Crate-private until a worker or client calls it. #![allow(dead_code)] -use futures::future; +use futures::stream::{FuturesUnordered, TryStreamExt}; use std::{collections::HashMap, sync::Arc}; use prost::Message; @@ -122,13 +122,13 @@ async fn retrieve_payloads( work.push((driver_name, driver, group)); } - let retrieved_groups = - future::try_join_all(work.into_iter().map(|(driver_name, driver, group)| { - let limiter = Arc::clone(limiter); + let mut pending: FuturesUnordered<_> = work + .into_iter() + .map(|(driver_name, driver, group)| { + let context = StorageDriverRetrieveContext::builder() + .limiter(Arc::clone(limiter)) + .build(); async move { - let context = StorageDriverRetrieveContext::builder() - .limiter(limiter) - .build(); let claims: Vec<_> = group.iter().map(|(_, claim)| claim.clone()).collect(); let retrieved = driver.retrieve(&context, claims).await?; if retrieved.len() != group.len() { @@ -138,19 +138,13 @@ async fn retrieve_payloads( returned: retrieved.len(), }); } - Ok::<_, ExternalStorageError>( - group - .into_iter() - .map(|(index, _)| index) - .zip(retrieved) - .collect::>(), - ) + Ok(group.into_iter().map(|(index, _)| index).zip(retrieved)) } - })) - .await?; + }) + .collect(); - for group in retrieved_groups { - for (index, payload) in group { + while let Some(retrieved) = pending.try_next().await? { + for (index, payload) in retrieved { payloads[index] = payload; } } @@ -303,15 +297,14 @@ async fn store_payloads( .push((index, payload.clone(), size as i64)); } - let stored = - future::try_join_all(by_driver.into_iter().map(|(driver_name, (driver, group))| { - let limiter = Arc::clone(limiter); - let target = target.cloned(); + let mut pending: FuturesUnordered<_> = by_driver + .into_iter() + .map(|(driver_name, (driver, group))| { + let context = StorageDriverStoreContext::builder() + .maybe_target(target.cloned()) + .limiter(Arc::clone(limiter)) + .build(); async move { - let context = StorageDriverStoreContext::builder() - .maybe_target(target) - .limiter(limiter) - .build(); let to_store: Vec<_> = group .iter() .map(|(_, payload, _)| payload.clone()) @@ -331,11 +324,11 @@ async fn store_payloads( } Ok::<_, ExternalStorageError>(replaced) } - })) - .await?; + }) + .collect(); - for group in stored { - for (index, reference) in group { + while let Some(replaced) = pending.try_next().await? { + for (index, reference) in replaced { payloads[index] = reference; } } From cfdd148be1926c68172b0d6db1f959c48e17eb32 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:44 -0700 Subject: [PATCH 10/15] perf: move payloads and claims to drivers instead of cloning --- crates/common/src/external_storage_runner.rs | 37 ++++++++++---------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 33b3712e5..79a33e5ec 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -101,16 +101,15 @@ async fn retrieve_payloads( storage: &ExternalStorage, limiter: &Arc>, ) -> Result<(), ExternalStorageError> { - let mut by_driver: HashMap> = HashMap::new(); + let mut by_driver: HashMap, Vec)> = HashMap::new(); for (index, payload) in payloads.iter().enumerate() { if !is_plain_reference(payload) { continue; } let reference = parse_reference(payload)?; - by_driver - .entry(reference.driver_name) - .or_default() - .push((index, StorageDriverClaim::new(reference.claim_data))); + let (indices, claims) = by_driver.entry(reference.driver_name).or_default(); + indices.push(index); + claims.push(StorageDriverClaim::new(reference.claim_data)); } // Resolve every driver first, so an unregistered name fails before any request is made. @@ -124,21 +123,20 @@ async fn retrieve_payloads( let mut pending: FuturesUnordered<_> = work .into_iter() - .map(|(driver_name, driver, group)| { + .map(|(driver_name, driver, (indices, claims))| { let context = StorageDriverRetrieveContext::builder() .limiter(Arc::clone(limiter)) .build(); async move { - let claims: Vec<_> = group.iter().map(|(_, claim)| claim.clone()).collect(); let retrieved = driver.retrieve(&context, claims).await?; - if retrieved.len() != group.len() { + if retrieved.len() != indices.len() { return Err(ExternalStorageError::DriverContract { driver: driver_name, - requested: group.len(), + requested: indices.len(), returned: retrieved.len(), }); } - Ok(group.into_iter().map(|(index, _)| index).zip(retrieved)) + Ok(indices.into_iter().zip(retrieved)) } }) .collect(); @@ -254,9 +252,8 @@ impl AsyncPayloadVisitor for RetrievePayloadsVisitor<'_> { } } -/// One driver and the payloads it takes: each payload's index, the payload, and its pre-offload -/// proto size. -type DriverGroup = (Arc, Vec<(usize, Payload, i64)>); +/// One driver and the payloads it takes: each payload's index and its pre-offload proto size. +type DriverGroup = (Arc, Vec<(usize, i64)>); /// Replace every eligible payload with a reference. async fn store_payloads( @@ -294,21 +291,23 @@ async fn store_payloads( .entry(driver.name().to_owned()) .or_insert_with(|| (Arc::clone(&driver), Vec::new())) .1 - .push((index, payload.clone(), size as i64)); + .push((index, size as i64)); } let mut pending: FuturesUnordered<_> = by_driver .into_iter() .map(|(driver_name, (driver, group))| { + // Moved out rather than cloned because payloads can be large. Selection is complete by + // now, so a selection error never leaves an emptied slot behind. + let to_store: Vec<_> = group + .iter() + .map(|&(index, _)| std::mem::take(&mut payloads[index])) + .collect(); let context = StorageDriverStoreContext::builder() .maybe_target(target.cloned()) .limiter(Arc::clone(limiter)) .build(); async move { - let to_store: Vec<_> = group - .iter() - .map(|(_, payload, _)| payload.clone()) - .collect(); let claims = driver.store(&context, to_store).await?; if claims.len() != group.len() { return Err(ExternalStorageError::DriverContract { @@ -318,7 +317,7 @@ async fn store_payloads( }); } let mut replaced = Vec::with_capacity(group.len()); - for ((index, _, size), claim) in group.into_iter().zip(claims) { + for ((index, size), claim) in group.into_iter().zip(claims) { let reference = create_reference_payload(&driver_name, claim.claim_data, size)?; replaced.push((index, reference)); } From be0d7528aba58382ae2e9dc4bd081d6bc12458c5 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:44 -0700 Subject: [PATCH 11/15] docs: warn that a failed pass leaves the message unusable --- crates/common/src/external_storage_runner.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 79a33e5ec..134e1dafc 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -56,6 +56,8 @@ impl<'a> ExternalStorageRunner<'a> { } /// Replace every eligible payload in a message with a reference. + /// + /// On error the message is left partly processed and must not be used. pub(crate) async fn store_in( &self, message: &mut M, @@ -72,6 +74,8 @@ impl<'a> ExternalStorageRunner<'a> { } /// Replace every reference in a message with the payload it points at. + /// + /// On error the message is left partly processed and must not be used. pub(crate) async fn retrieve_in( &self, message: &mut M, From 9a6a53b0efc5988727de673821fb86e944d1c2ab Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:42:44 -0700 Subject: [PATCH 12/15] refactor: drop reject_references until it has a caller --- crates/common/src/external_storage_runner.rs | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 134e1dafc..7df20419d 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -90,15 +90,6 @@ impl<'a> ExternalStorageRunner<'a> { } } -/// Fail if any payload needs retrieval. For a caller with no external storage configured, which -/// must not hand a reference to application code as though it were the application's own data. -pub(crate) fn reject_references(payloads: &[Payload]) -> Result<(), ExternalStorageError> { - if payloads.iter().any(is_plain_reference) { - return Err(ExternalStorageError::NotConfigured); - } - Ok(()) -} - /// Replace every reference with the payload the driver returns for it. async fn retrieve_payloads( payloads: &mut [Payload], @@ -1058,13 +1049,4 @@ mod tests { sizes.sort_unstable(); assert_eq!(sizes, vec![1, 1, 1, 2]); } - - #[test] - fn rejecting_references_without_storage_is_tmprl1105() { - assert!(reject_references(&[inline("fine")]).is_ok()); - - let err = reject_references(&[reference("mem", "a")]).unwrap_err(); - assert!(matches!(err, ExternalStorageError::NotConfigured)); - assert!(err.to_string().contains("TMPRL1105")); - } } From 26fbff5354df45010cb34d9b4b0182b280e690f9 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:36:01 -0700 Subject: [PATCH 13/15] refactor: look up drivers by scanning instead of a name map --- crates/common-wasm/src/external_storage.rs | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/crates/common-wasm/src/external_storage.rs b/crates/common-wasm/src/external_storage.rs index 1913921d9..00bd9e2b8 100644 --- a/crates/common-wasm/src/external_storage.rs +++ b/crates/common-wasm/src/external_storage.rs @@ -10,7 +10,11 @@ use crate::protos::temporal::api::common::v1::Payload; use futures::future::BoxFuture; -use std::{collections::HashMap, fmt, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + fmt, + sync::Arc, +}; const DEFAULT_PAYLOAD_SIZE_THRESHOLD: usize = 256 * 1024; @@ -306,7 +310,6 @@ pub enum ExternalStorageError { #[derive(Clone)] pub struct ExternalStorage { drivers: Vec>, - drivers_by_name: HashMap>, driver_selector: Arc, payload_size_threshold: usize, concurrency: ExternalStorageConcurrency, @@ -345,7 +348,7 @@ impl ExternalStorage { "at least one driver is required".to_owned(), )); } - let mut drivers_by_name = HashMap::with_capacity(drivers.len()); + let mut names = HashSet::with_capacity(drivers.len()); for driver in &drivers { let name = driver.name(); // The name is the routing key written into history, so a driver without one can never @@ -355,10 +358,7 @@ impl ExternalStorage { "driver name cannot be empty".to_owned(), )); } - if drivers_by_name - .insert(name.to_owned(), Arc::clone(driver)) - .is_some() - { + if !names.insert(name) { return Err(ExternalStorageError::InvalidConfiguration(format!( "multiple drivers given with name '{name}'" ))); @@ -366,7 +366,6 @@ impl ExternalStorage { } Ok(Self { drivers, - drivers_by_name, driver_selector, payload_size_threshold: DEFAULT_PAYLOAD_SIZE_THRESHOLD, concurrency: ExternalStorageConcurrency::default(), @@ -411,7 +410,7 @@ impl ExternalStorage { /// The driver registered under the given name, if any. pub fn driver(&self, name: &str) -> Option<&Arc> { - self.drivers_by_name.get(name) + self.drivers.iter().find(|d| d.name() == name) } } From 5bc63d7a0711130352868827d7237fdebaa0cca8 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:04:38 -0700 Subject: [PATCH 14/15] refactor: generate payload visitors with a macro --- crates/common/src/external_storage_runner.rs | 102 ++++------------ crates/common/src/payload_visitor.rs | 121 ++++++++----------- 2 files changed, 73 insertions(+), 150 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 7df20419d..0867cc260 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -25,9 +25,7 @@ use temporalio_common_wasm::data_converters::{ use crate::{ external_storage_references::{create_reference_payload, is_plain_reference, parse_reference}, - payload_visitor::{ - AsyncPayloadVisitor, PayloadField, PayloadFieldData, PayloadVisitable, should_transform, - }, + payload_visitor::{PayloadVisitable, payload_visitor}, protos::temporal::api::common::v1::Payload, }; @@ -156,8 +154,16 @@ fn finish( walked.map_err(|e| ExternalStorageError::Walk(Box::new(e))) } -fn stop_walk() -> PayloadConversionError { - PayloadConversionError::EncodingError("external storage pass failed".into()) +/// Keep a pass's typed error aside and stop the walk with a placeholder, since the walk can only +/// report a [`PayloadConversionError`]. +fn stop_walk( + slot: &mut Option, + result: Result<(), ExternalStorageError>, +) -> Result<(), PayloadConversionError> { + result.map_err(|e| { + *slot = Some(e); + PayloadConversionError::EncodingError("external storage pass failed".into()) + }) } struct StorePayloadsVisitor<'a> { @@ -166,48 +172,13 @@ struct StorePayloadsVisitor<'a> { limiter: &'a Arc>, error: Option, } -impl AsyncPayloadVisitor for StorePayloadsVisitor<'_> { - fn visit<'a>( - &'a mut self, - field: PayloadField<'a>, - ) -> futures::future::BoxFuture<'a, Result<(), PayloadConversionError>> { - Box::pin(async move { - if !should_transform(field.path) { - return Ok(()); - } - let result = match field.data { - PayloadFieldData::Single(payload) => { - store_payloads( - std::slice::from_mut(payload), - self.storage, - self.target, - self.limiter, - ) - .await - } - PayloadFieldData::Repeated(payloads) => { - store_payloads(payloads, self.storage, self.target, self.limiter).await - } - PayloadFieldData::Payloads(message) => { - store_payloads( - &mut message.payloads, - self.storage, - self.target, - self.limiter, - ) - .await - } - }; - match result { - Ok(()) => Ok(()), - Err(e) => { - self.error = Some(e); - Err(stop_walk()) - } - } - }) - } -} + +payload_visitor!(StorePayloadsVisitor<'_>, |this, payloads| { + let mut payloads = payloads; + let result = store_payloads(&mut payloads, this.storage, this.target, this.limiter).await; + stop_walk(&mut this.error, result)?; + payloads +}); struct RetrievePayloadsVisitor<'a> { storage: &'a ExternalStorage, @@ -215,37 +186,12 @@ struct RetrievePayloadsVisitor<'a> { error: Option, } -impl AsyncPayloadVisitor for RetrievePayloadsVisitor<'_> { - fn visit<'a>( - &'a mut self, - field: PayloadField<'a>, - ) -> futures::future::BoxFuture<'a, Result<(), PayloadConversionError>> { - Box::pin(async move { - if !should_transform(field.path) { - return Ok(()); - } - let result = match field.data { - PayloadFieldData::Single(payload) => { - retrieve_payloads(std::slice::from_mut(payload), self.storage, self.limiter) - .await - } - PayloadFieldData::Repeated(payloads) => { - retrieve_payloads(payloads, self.storage, self.limiter).await - } - PayloadFieldData::Payloads(message) => { - retrieve_payloads(&mut message.payloads, self.storage, self.limiter).await - } - }; - match result { - Ok(()) => Ok(()), - Err(e) => { - self.error = Some(e); - Err(stop_walk()) - } - } - }) - } -} +payload_visitor!(RetrievePayloadsVisitor<'_>, |this, payloads| { + let mut payloads = payloads; + let result = retrieve_payloads(&mut payloads, this.storage, this.limiter).await; + stop_walk(&mut this.error, result)?; + payloads +}); /// One driver and the payloads it takes: each payload's index and its pre-offload proto size. type DriverGroup = (Arc, Vec<(usize, i64)>); diff --git a/crates/common/src/payload_visitor.rs b/crates/common/src/payload_visitor.rs index d9ae4126c..9f3635620 100644 --- a/crates/common/src/payload_visitor.rs +++ b/crates/common/src/payload_visitor.rs @@ -63,48 +63,58 @@ pub(crate) fn should_transform(path: &str) -> bool { !is_search_attributes_path(path) } +/// Implement [`AsyncPayloadVisitor`] for a visitor that applies one batch operation to every +/// transformable field, given `|visitor, payloads| body`. The body takes the field's payloads as an +/// owned `Vec` and evaluates to the payloads that replace them. +macro_rules! payload_visitor { + ($visitor:ty, |$this:ident, $payloads:ident| $body:expr) => { + impl $crate::payload_visitor::AsyncPayloadVisitor for $visitor { + fn visit<'a>( + &'a mut self, + field: $crate::payload_visitor::PayloadField<'a>, + ) -> ::futures::future::BoxFuture< + 'a, + Result<(), $crate::data_converters::PayloadConversionError>, + > { + Box::pin(async move { + if !$crate::payload_visitor::should_transform(field.path) { + return Ok(()); + } + let $this = self; + match field.data { + $crate::payload_visitor::PayloadFieldData::Single(payload) => { + let $payloads = vec![::std::mem::take(payload)]; + if let Some(p) = $body.into_iter().next() { + *payload = p; + } + } + $crate::payload_visitor::PayloadFieldData::Repeated(payloads) => { + let $payloads = ::std::mem::take(payloads); + *payloads = $body; + } + $crate::payload_visitor::PayloadFieldData::Payloads(message) => { + let $payloads = ::std::mem::take(&mut message.payloads); + message.payloads = $body; + } + } + Ok(()) + }) + } + } + }; +} +#[cfg(feature = "experimental")] +pub(crate) use payload_visitor; + /// Visitor that encodes payloads using a codec. pub struct EncodeVisitor<'a> { codec: &'a (dyn PayloadCodec + Send + Sync), context: &'a SerializationContextData, } -impl AsyncPayloadVisitor for EncodeVisitor<'_> { - fn visit<'a>( - &'a mut self, - field: PayloadField<'a>, - ) -> BoxFuture<'a, Result<(), PayloadConversionError>> { - Box::pin(async move { - if !should_transform(field.path) { - return Ok(()); - } - match field.data { - PayloadFieldData::Single(payload) => { - let encoded = self - .codec - .encode(self.context, vec![std::mem::take(payload)]) - .await?; - if let Some(p) = encoded.into_iter().next() { - *payload = p; - } - } - PayloadFieldData::Repeated(payloads) => { - *payloads = self - .codec - .encode(self.context, std::mem::take(payloads)) - .await?; - } - PayloadFieldData::Payloads(payloads_msg) => { - payloads_msg.payloads = self - .codec - .encode(self.context, std::mem::take(&mut payloads_msg.payloads)) - .await?; - } - } - Ok(()) - }) - } -} +payload_visitor!(EncodeVisitor<'_>, |this, payloads| { + this.codec.encode(this.context, payloads).await? +}); /// Visitor that decodes payloads using a codec. pub struct DecodeVisitor<'a> { @@ -112,42 +122,9 @@ pub struct DecodeVisitor<'a> { context: &'a SerializationContextData, } -impl AsyncPayloadVisitor for DecodeVisitor<'_> { - fn visit<'a>( - &'a mut self, - field: PayloadField<'a>, - ) -> BoxFuture<'a, Result<(), PayloadConversionError>> { - Box::pin(async move { - if !should_transform(field.path) { - return Ok(()); - } - match field.data { - PayloadFieldData::Single(payload) => { - let decoded = self - .codec - .decode(self.context, vec![std::mem::take(payload)]) - .await?; - if let Some(p) = decoded.into_iter().next() { - *payload = p; - } - } - PayloadFieldData::Repeated(payloads) => { - *payloads = self - .codec - .decode(self.context, std::mem::take(payloads)) - .await?; - } - PayloadFieldData::Payloads(payloads_msg) => { - payloads_msg.payloads = self - .codec - .decode(self.context, std::mem::take(&mut payloads_msg.payloads)) - .await?; - } - } - Ok(()) - }) - } -} +payload_visitor!(DecodeVisitor<'_>, |this, payloads| { + this.codec.decode(this.context, payloads).await? +}); /// Encode all payloads in a message using the given codec. pub async fn encode_payloads( From 7d689fcc04265e6d04a4110e6899ea598b3617d9 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:07:32 -0700 Subject: [PATCH 15/15] refactor: make the payload walk a runner helper --- crates/common/src/external_storage_runner.rs | 38 ++++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/crates/common/src/external_storage_runner.rs b/crates/common/src/external_storage_runner.rs index 0867cc260..6c0710068 100644 --- a/crates/common/src/external_storage_runner.rs +++ b/crates/common/src/external_storage_runner.rs @@ -25,7 +25,7 @@ use temporalio_common_wasm::data_converters::{ use crate::{ external_storage_references::{create_reference_payload, is_plain_reference, parse_reference}, - payload_visitor::{PayloadVisitable, payload_visitor}, + payload_visitor::{AsyncPayloadVisitor, PayloadVisitable, payload_visitor}, protos::temporal::api::common::v1::Payload, }; @@ -61,14 +61,13 @@ impl<'a> ExternalStorageRunner<'a> { message: &mut M, target: Option<&StorageDriverTargetInfo>, ) -> Result<(), ExternalStorageError> { - let mut store = StorePayloadsVisitor { + let store = StorePayloadsVisitor { storage: self.storage, target, limiter: &self.store_limiter, error: None, }; - let walked = message.visit_payloads_mut(&mut store).await; - finish(store.error, walked) + Self::walk(message, store, |v| v.error.take()).await } /// Replace every reference in a message with the payload it points at. @@ -78,13 +77,26 @@ impl<'a> ExternalStorageRunner<'a> { &self, message: &mut M, ) -> Result<(), ExternalStorageError> { - let mut retrieve = RetrievePayloadsVisitor { + let retrieve = RetrievePayloadsVisitor { storage: self.storage, limiter: &self.retrieve_limiter, error: None, }; - let walked = message.visit_payloads_mut(&mut retrieve).await; - finish(retrieve.error, walked) + Self::walk(message, retrieve, |v| v.error.take()).await + } + + /// The visitor reports only that it failed, so the typed error it captured is preferred over the + /// placeholder that stopped the walk. + async fn walk( + message: &mut M, + mut visitor: V, + take_error: impl FnOnce(&mut V) -> Option, + ) -> Result<(), ExternalStorageError> { + let walked = message.visit_payloads_mut(&mut visitor).await; + if let Some(error) = take_error(&mut visitor) { + return Err(error); + } + walked.map_err(|e| ExternalStorageError::Walk(Box::new(e))) } } @@ -142,18 +154,6 @@ async fn retrieve_payloads( Ok(()) } -/// The visitor reports only that it failed, so the typed error it captured is preferred over the -/// placeholder that stopped the walk. -fn finish( - captured: Option, - walked: Result<(), PayloadConversionError>, -) -> Result<(), ExternalStorageError> { - if let Some(error) = captured { - return Err(error); - } - walked.map_err(|e| ExternalStorageError::Walk(Box::new(e))) -} - /// Keep a pass's typed error aside and stop the walk with a placeholder, since the walk can only /// report a [`PayloadConversionError`]. fn stop_walk(