diff --git a/crates/common-wasm/src/external_storage.rs b/crates/common-wasm/src/external_storage.rs index 398dcfed6..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; @@ -146,11 +150,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 +165,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 +187,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 +200,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 +276,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. @@ -300,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, @@ -339,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 @@ -349,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}'" ))); @@ -360,7 +366,6 @@ impl ExternalStorage { } Ok(Self { drivers, - drivers_by_name, driver_selector, payload_size_threshold: DEFAULT_PAYLOAD_SIZE_THRESHOLD, concurrency: ExternalStorageConcurrency::default(), @@ -405,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) } } @@ -424,8 +429,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 +472,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..4d5c89943 100644 --- a/crates/common/src/external_storage_references.rs +++ b/crates/common/src/external_storage_references.rs @@ -22,18 +22,25 @@ 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 a plain external storage reference and so has to be retrieved. +/// Reads metadata only, so one that needs no retrieval is not touched. +/// +/// 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) == 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 +108,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 +150,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 +204,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..6c0710068 --- /dev/null +++ b/crates/common/src/external_storage_runner.rs @@ -0,0 +1,998 @@ +//! 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, 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, +//! and the configured concurrency limits are not enforced. + +// Crate-private until a worker or client calls it. +#![allow(dead_code)] + +use futures::stream::{FuturesUnordered, TryStreamExt}; +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, PayloadVisitable, payload_visitor}, + 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: Arc>, + retrieve_limiter: Arc>, +} + +impl<'a> ExternalStorageRunner<'a> { + pub(crate) fn new( + storage: &'a ExternalStorage, + store_limiter: Arc>, + retrieve_limiter: Arc>, + ) -> Self { + Self { + storage, + store_limiter, + retrieve_limiter, + } + } + + /// 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, + target: Option<&StorageDriverTargetInfo>, + ) -> Result<(), ExternalStorageError> { + let store = StorePayloadsVisitor { + storage: self.storage, + target, + limiter: &self.store_limiter, + error: None, + }; + Self::walk(message, store, |v| v.error.take()).await + } + + /// 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, + ) -> Result<(), ExternalStorageError> { + let retrieve = RetrievePayloadsVisitor { + storage: self.storage, + limiter: &self.retrieve_limiter, + error: None, + }; + 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))) + } +} + +/// Replace every reference with the payload the driver returns for it. +async fn retrieve_payloads( + payloads: &mut [Payload], + storage: &ExternalStorage, + limiter: &Arc>, +) -> Result<(), ExternalStorageError> { + 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)?; + 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. + 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 mut pending: FuturesUnordered<_> = work + .into_iter() + .map(|(driver_name, driver, (indices, claims))| { + let context = StorageDriverRetrieveContext::builder() + .limiter(Arc::clone(limiter)) + .build(); + async move { + let retrieved = driver.retrieve(&context, claims).await?; + if retrieved.len() != indices.len() { + return Err(ExternalStorageError::DriverContract { + driver: driver_name, + requested: indices.len(), + returned: retrieved.len(), + }); + } + Ok(indices.into_iter().zip(retrieved)) + } + }) + .collect(); + + while let Some(retrieved) = pending.try_next().await? { + for (index, payload) in retrieved { + payloads[index] = payload; + } + } + Ok(()) +} + +/// 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> { + storage: &'a ExternalStorage, + target: Option<&'a StorageDriverTargetInfo>, + limiter: &'a Arc>, + error: Option, +} + +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, + limiter: &'a Arc>, + error: Option, +} + +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)>); + +/// Replace every eligible payload with a reference. +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 + .driver(driver.name()) + .is_some_and(|registered| Arc::ptr_eq(registered, &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, 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 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) + } + }) + .collect(); + + while let Some(replaced) = pending.try_next().await? { + for (index, reference) in replaced { + 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" + )); + } + + #[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( + 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 runner = ExternalStorageRunner::new(&storage, store_limiter(), limiter()); + 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 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 { + 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 runner = ExternalStorageRunner::new(&storage, store_limiter(), limiter()); + 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 runner = ExternalStorageRunner::new(&storage, store_limiter(), limiter()); + 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]); + } +} 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..9f3635620 100644 --- a/crates/common/src/payload_visitor.rs +++ b/crates/common/src/payload_visitor.rs @@ -51,59 +51,70 @@ 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) } +/// 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_encode(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> { @@ -111,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_encode(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(