From 4bd02785e909afd9b51b68b864dec0ebb30c3232 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 28 Oct 2026 12:48:56 -0400 Subject: [PATCH 1/4] feat(extstore): add aws-s3driver crate with S3 driver for external storage --- .github/workflows/per-pr.yml | 2 +- Cargo.toml | 1 + changelog/added/solemn-simulacrum.md | 1 + crates/aws-s3driver/Cargo.toml | 44 + crates/aws-s3driver/README.md | 135 +++ crates/aws-s3driver/src/aws_sdk_client.rs | 194 ++++ crates/aws-s3driver/src/client.rs | 38 + crates/aws-s3driver/src/driver.rs | 1259 +++++++++++++++++++++ crates/aws-s3driver/src/lib.rs | 21 + 9 files changed, 1694 insertions(+), 1 deletion(-) create mode 100644 changelog/added/solemn-simulacrum.md create mode 100644 crates/aws-s3driver/Cargo.toml create mode 100644 crates/aws-s3driver/README.md create mode 100644 crates/aws-s3driver/src/aws_sdk_client.rs create mode 100644 crates/aws-s3driver/src/client.rs create mode 100644 crates/aws-s3driver/src/driver.rs create mode 100644 crates/aws-s3driver/src/lib.rs diff --git a/.github/workflows/per-pr.yml b/.github/workflows/per-pr.yml index e046f770f..507842f13 100644 --- a/.github/workflows/per-pr.yml +++ b/.github/workflows/per-pr.yml @@ -106,7 +106,7 @@ jobs: - group: rust-1.88 crates: client common common-wasm macros protos sdk-core workflow - group: rust-1.92 - crates: sdk + crates: aws-s3driver sdk steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - uses: ./.github/actions/setup diff --git a/Cargo.toml b/Cargo.toml index 76bdd46f6..419d80247 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "crates/client", "crates/common", "crates/common-wasm", + "crates/aws-s3driver", "crates/protos", "crates/macros", "crates/sdk", diff --git a/changelog/added/solemn-simulacrum.md b/changelog/added/solemn-simulacrum.md new file mode 100644 index 000000000..486eae2b5 --- /dev/null +++ b/changelog/added/solemn-simulacrum.md @@ -0,0 +1 @@ +Added experimental `temporalio-aws-s3driver` crate with an S3 storage driver for external storage payloads, and an `aws-sdk` feature, enabled by default, with an `aws-sdk-s3` backed client. diff --git a/crates/aws-s3driver/Cargo.toml b/crates/aws-s3driver/Cargo.toml new file mode 100644 index 000000000..d42d391f4 --- /dev/null +++ b/crates/aws-s3driver/Cargo.toml @@ -0,0 +1,44 @@ +[package] +name = "temporalio-aws-s3driver" +version = "1.1.0" +edition = "2024" +rust-version = "1.92.0" +authors = ["Temporal Technologies Inc. "] +license-file = { workspace = true } +description = "Amazon S3 external storage driver for the Temporal Rust SDK" +homepage = "https://temporal.io/" +repository = "https://github.com/temporalio/sdk-rust" +keywords = ["temporal", "workflow", "s3"] +categories = ["development-tools"] +readme = "README.md" + +[package.metadata.docs.rs] +all-features = true + +[features] +default = ["aws-sdk"] +aws-sdk = ["dep:aws-sdk-s3"] + +[dependencies] +# The application configures the S3 client, including its TLS stack and async runtime. +aws-sdk-s3 = { version = "1", default-features = false, optional = true } +futures-util = { version = "0.3", default-features = false, features = ["alloc"] } +prost = { workspace = true } +sha2 = "0.11" +thiserror = { workspace = true } + +[dependencies.temporalio-common] +path = "../common" +version = "~1.1.0" +default-features = false +features = ["experimental"] + +[dev-dependencies] +aws-sdk-s3 = { version = "1", default-features = false, features = ["rt-tokio"] } +aws-smithy-http-client = { version = "1", features = ["test-util"] } +aws-smithy-types = "1" +http = "1" +tokio = { version = "1.47", features = ["macros", "rt", "sync", "time"] } + +[lints] +workspace = true diff --git a/crates/aws-s3driver/README.md b/crates/aws-s3driver/README.md new file mode 100644 index 000000000..f544d63d0 --- /dev/null +++ b/crates/aws-s3driver/README.md @@ -0,0 +1,135 @@ +# Amazon S3 External Storage Driver + +> [!WARNING] +> This crate is experimental and may change in the future. + +`temporalio-aws-s3driver` stores and retrieves Temporal payloads in Amazon S3 using external +storage. + +`S3StorageDriver` has no AWS dependency. It makes requests through the `S3StorageDriverClient` +trait. The `aws-sdk` feature, enabled by default, adds `AwsSdkS3StorageDriverClient`, backed by the +official [`aws-sdk-s3`](https://crates.io/crates/aws-sdk-s3) crate. To bring your own client without +depending on the AWS SDK, disable default features and implement `S3StorageDriverClient`. + +## Quick Start + +```rust,ignore +use std::sync::Arc; +use temporalio_aws_s3driver::{AwsSdkS3StorageDriverClient, S3StorageDriver}; + +let config = aws_config::load_from_env().await; +let driver = S3StorageDriver::new( + Arc::new(AwsSdkS3StorageDriverClient::new(aws_sdk_s3::Client::new(&config))), + "my-temporal-payloads", +); +``` + +To choose the bucket per payload, use `S3StorageDriver::with_bucket_selector`. + +## Required IAM Permissions + +Before uploading a payload, the driver checks whether the object already exists with a `HeadObject` +request. That check needs `s3:GetObject`. It also needs `s3:ListBucket`, because without it S3 +reports a missing object as `403 Forbidden` instead of `404 Not Found`, and the upload fails. + +```json +[ + { + "Effect": "Allow", + "Action": ["s3:PutObject", "s3:GetObject"], + "Resource": "arn:aws:s3:::my-temporal-payloads/*" + }, + { + "Effect": "Allow", + "Action": "s3:ListBucket", + "Resource": "arn:aws:s3:::my-temporal-payloads" + } +] +``` + +Components that store payloads need all three permissions. Components that only retrieve payloads +need only `s3:GetObject`. + +## S3 Storage Key Specification + +All Temporal S3 drivers generate S3 keys in a consistent manner. + +### Key format + +Workflow key: + +```text +v0/ns/{namespace}/wt/{workflow-type}/wi/{workflow-id}/ri/{run-id}/d/{hash-algorithm}/{hex-digest} +``` + +Activity key: + +```text +v0/ns/{namespace}/at/{activity-type}/ai/{activity-id}/ri/{run-id}/d/{hash-algorithm}/{hex-digest} +``` + +Fallback key (unknown target): + +```text +v0/d/{hash-algorithm}/{hex-digest} +``` + +- If no namespace, workflow, or activity information is available, the fallback is used. +- Dynamic path segments are percent-encoded (rules below). +- Missing values (including a missing `run-id`) are encoded as `null`. +- `hex-digest` is lower-case SHA-256 hex (64 characters). + +### Percent-encoding rules + +The Temporal SDKs escape anything that isn't listed in S3's safe character set: https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html + +Safe Characters: + +```text +Alphanumeric characters + 0-9 + a-z + A-Z + +Special characters + Exclamation point (!) + Hyphen (-) + Underscore (_) + Period (.) + Asterisk (*) + Single quotation mark (') + Opening parenthesis (() + Closing parenthesis ()) +``` + +### Examples + +Workflow key example: + +```text +input: + namespace=payments prod + workflow-type=ChargeWorkflow + workflow-id=order+123=abc + run-id=3f1d6c7a-8b2e-4f7a-9d0a-87a6f95e4d31 + hash-algorithm=sha256 + hex-digest=9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08 + +output: + v0/ns/payments%20prod/wt/ChargeWorkflow/wi/order%2B123%3Dabc/ri/3f1d6c7a-8b2e-4f7a-9d0a-87a6f95e4d31/d/sha256/9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08 +``` + +Activity key example: + +```text +input: + namespace=payments prod + activity-type=Capture/Charge + activity-id=activity id+42 + run-id=9e1d1fd9-2f8a-4c40-93e2-731f31b9268b + hash-algorithm=sha256 + hex-digest=2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824 + +output: + v0/ns/payments%20prod/at/Capture%2FCharge/ai/activity%20id%2B42/ri/9e1d1fd9-2f8a-4c40-93e2-731f31b9268b/d/sha256/2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824 +``` diff --git a/crates/aws-s3driver/src/aws_sdk_client.rs b/crates/aws-s3driver/src/aws_sdk_client.rs new file mode 100644 index 000000000..43856d8c6 --- /dev/null +++ b/crates/aws-s3driver/src/aws_sdk_client.rs @@ -0,0 +1,194 @@ +use crate::client::{BoxError, S3StorageDriverClient}; +use aws_sdk_s3::{Client, primitives::ByteStream}; +use futures_util::future::BoxFuture; +use std::collections::HashMap; + +/// [`S3StorageDriverClient`] backed by an [`aws_sdk_s3::Client`] from the official AWS SDK. +/// +/// This API is experimental and may change in the future. +#[derive(Clone, Debug)] +pub struct AwsSdkS3StorageDriverClient { + client: Client, +} + +impl AwsSdkS3StorageDriverClient { + /// Create a client that makes requests with the given AWS S3 client. + pub fn new(client: Client) -> Self { + Self { client } + } +} + +impl S3StorageDriverClient for AwsSdkS3StorageDriverClient { + fn put_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + data: Vec, + ) -> BoxFuture<'a, Result<(), BoxError>> { + Box::pin(async move { + self.client + .put_object() + .bucket(bucket) + .key(key) + .body(ByteStream::from(data)) + .send() + .await?; + Ok(()) + }) + } + + fn object_exists<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + match self + .client + .head_object() + .bucket(bucket) + .key(key) + .send() + .await + { + Ok(_) => Ok(true), + Err(err) if err.as_service_error().is_some_and(|err| err.is_not_found()) => { + Ok(false) + } + Err(err) => Err(err.into()), + } + }) + } + + fn get_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result, BoxError>> { + Box::pin(async move { + let output = self + .client + .get_object() + .bucket(bucket) + .key(key) + .send() + .await?; + Ok(output.body.collect().await?.into_bytes().to_vec()) + }) + } + + fn describe(&self) -> HashMap { + self.client + .config() + .region() + .map(|region| HashMap::from([("client_region".to_owned(), region.to_string())])) + .unwrap_or_default() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aws_sdk_s3::{ + Config, + config::{BehaviorVersion, Credentials, Region}, + }; + use aws_smithy_http_client::test_util::infallible_client_fn; + use aws_smithy_types::body::SdkBody; + use std::sync::{Arc, Mutex}; + + fn client( + respond: impl Fn(http::Request) -> http::Response + Send + Sync + 'static, + ) -> AwsSdkS3StorageDriverClient { + let config = Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-west-2")) + .credentials_provider(Credentials::new("id", "secret", None, None, "test")) + .force_path_style(true) + .http_client(infallible_client_fn(respond)) + .build(); + AwsSdkS3StorageDriverClient::new(Client::from_conf(config)) + } + + fn status(code: u16) -> http::Response { + http::Response::builder() + .status(code) + .body(SdkBody::empty()) + .unwrap() + } + + #[tokio::test] + async fn object_exists_returns_true_for_200_and_false_for_404() { + assert!( + client(|_| status(200)) + .object_exists("my-bucket", "my-key") + .await + .unwrap() + ); + assert!( + !client(|_| status(404)) + .object_exists("my-bucket", "my-key") + .await + .unwrap() + ); + } + + #[tokio::test] + async fn object_exists_fails_for_errors_other_than_404() { + client(|_| status(403)) + .object_exists("my-bucket", "my-key") + .await + .unwrap_err(); + } + + #[tokio::test] + async fn put_object_sends_data_to_bucket_and_key() { + let requests = Arc::new(Mutex::new(Vec::new())); + let recorded = requests.clone(); + let client = client(move |request| { + recorded.lock().unwrap().push(( + request.method().clone(), + request.uri().path().to_owned(), + request.body().bytes().map(<[u8]>::to_vec), + )); + status(200) + }); + + client + .put_object("my-bucket", "my-key", vec![0, 1]) + .await + .unwrap(); + + assert_eq!( + *requests.lock().unwrap(), + [( + http::Method::PUT, + "/my-bucket/my-key".to_owned(), + Some(vec![0, 1]) + )] + ); + } + + #[tokio::test] + async fn get_object_returns_the_response_body() { + let client = client(|_| { + http::Response::builder() + .status(200) + .body(SdkBody::from(vec![0, 1])) + .unwrap() + }); + + assert_eq!( + client.get_object("my-bucket", "my-key").await.unwrap(), + [0, 1] + ); + } + + #[test] + fn describe_returns_the_client_region() { + assert_eq!( + client(|_| status(200)).describe(), + HashMap::from([("client_region".to_owned(), "us-west-2".to_owned())]) + ); + } +} diff --git a/crates/aws-s3driver/src/client.rs b/crates/aws-s3driver/src/client.rs new file mode 100644 index 000000000..69c3bca85 --- /dev/null +++ b/crates/aws-s3driver/src/client.rs @@ -0,0 +1,38 @@ +use futures_util::future::BoxFuture; +use std::collections::HashMap; + +/// Error returned by an [`S3StorageDriverClient`]. +pub type BoxError = Box; + +/// S3 object operations used by [`S3StorageDriver`](crate::S3StorageDriver). Implement this to use +/// an S3 library other than the AWS SDK. +/// +/// Implementations are called concurrently and must be thread safe. +pub trait S3StorageDriverClient: Send + Sync { + /// Upload data to the given bucket and key. + fn put_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + data: Vec, + ) -> BoxFuture<'a, Result<(), BoxError>>; + + /// Check whether an object exists at the given bucket and key. + fn object_exists<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result>; + + /// Download the object at the given bucket and key. + fn get_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result, BoxError>>; + + /// Client diagnostic information, such as the region, that the driver adds to error messages. + fn describe(&self) -> HashMap { + HashMap::new() + } +} diff --git a/crates/aws-s3driver/src/driver.rs b/crates/aws-s3driver/src/driver.rs new file mode 100644 index 000000000..fc4f611f5 --- /dev/null +++ b/crates/aws-s3driver/src/driver.rs @@ -0,0 +1,1259 @@ +use crate::client::{BoxError, S3StorageDriverClient}; +use futures_util::future::{BoxFuture, try_join_all}; +use prost::Message; +use sha2::{Digest, Sha256}; +use std::{collections::HashSet, sync::Arc}; +use temporalio_common::{ + data_converters::{ + ExternalStorageError, StorageDriver, StorageDriverClaim, StorageDriverLimiter, + StorageDriverRetrieveContext, StorageDriverStoreContext, StorageDriverTargetInfo, + }, + protos::temporal::api::common::v1::Payload, +}; + +const DRIVER_TYPE: &str = "aws.s3driver"; +const DEFAULT_MAX_PAYLOAD_SIZE: usize = 50 * 1024 * 1024; +const CLAIM_BUCKET: &str = "bucket"; +const CLAIM_KEY: &str = "key"; +const CLAIM_HASH_ALGORITHM: &str = "hash_algorithm"; +const CLAIM_HASH_VALUE: &str = "hash_value"; +const NULL_SEGMENT: &str = "null"; +const SAFE_SPECIAL_CHARACTERS: &[u8] = b"!-_.*'()"; + +type BucketSelector = dyn Fn(&StorageDriverStoreContext, &Payload) -> String + Send + Sync; + +/// Errors raised by [`S3StorageDriver`], returned inside [`ExternalStorageError::Driver`]. +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum S3StorageDriverError { + /// A payload is larger than the configured maximum. + #[error("Payload size {size} bytes exceeds the configured max payload size of {max} bytes")] + PayloadTooLarge { + /// Serialized payload size, in bytes. + size: usize, + /// Configured maximum, in bytes. + max: usize, + }, + /// A claim has no bucket or key. + #[error( + "S3StorageDriver claim is missing required location information: claim data must \ + contain 'bucket' and 'key'" + )] + MissingLocation, + /// A claim has no content hash. + #[error( + "S3StorageDriver claim is missing required content hash information [bucket={bucket}, \ + key={key}]: claim data must contain 'hash_algorithm' and 'hash_value'" + )] + MissingHash { + /// Bucket from the claim. + bucket: String, + /// Key from the claim. + key: String, + }, + /// A claim uses a hash algorithm other than SHA-256. + #[error( + "S3StorageDriver unsupported hash algorithm [bucket={bucket}, key={key}]: expected \ + sha256, got {hash_algorithm}" + )] + UnsupportedHashAlgorithm { + /// Bucket from the claim. + bucket: String, + /// Key from the claim. + key: String, + /// Hash algorithm from the claim. + hash_algorithm: String, + }, + /// The client failed to store an object. + #[error("S3StorageDriver store failed [bucket={bucket}, key={key}{client_context}]")] + Store { + /// Bucket of the object. + bucket: String, + /// Key of the object. + key: String, + /// Client diagnostic information, formatted for the message. + client_context: String, + /// The client error. + #[source] + source: BoxError, + }, + /// The client failed to retrieve an object. + #[error("S3StorageDriver retrieve failed [bucket={bucket}, key={key}{client_context}]")] + Retrieve { + /// Bucket of the object. + bucket: String, + /// Key of the object. + key: String, + /// Client diagnostic information, formatted for the message. + client_context: String, + /// The client error. + #[source] + source: BoxError, + }, + /// A retrieved object does not match the hash in its claim. + #[error( + "S3StorageDriver integrity check failed [bucket={bucket}, key={key}]: expected \ + sha256:{expected}, got sha256:{actual}" + )] + IntegrityCheck { + /// Bucket of the object. + bucket: String, + /// Key of the object. + key: String, + /// Hash from the claim. + expected: String, + /// Hash of the retrieved data. + actual: String, + }, +} + +impl From for ExternalStorageError { + fn from(err: S3StorageDriverError) -> Self { + ExternalStorageError::Driver(Box::new(err)) + } +} + +/// Storage driver that stores and retrieves payloads in Amazon S3. +/// +/// Payloads are content addressed by a SHA-256 hash of their serialized bytes, so storing the same +/// payload twice for the same execution uploads it once. Object keys are built from the namespace +/// and the workflow or activity the payload is stored for, percent-encoded to S3's safe character +/// set. +/// +/// This API is experimental and may change in the future. +#[derive(Clone)] +pub struct S3StorageDriver { + client: Arc, + bucket_selector: Arc, + name: String, + max_payload_size: usize, +} + +impl S3StorageDriver { + /// Create a driver that stores every payload in one bucket. + pub fn new(client: Arc, bucket: impl Into) -> Self { + let bucket = bucket.into(); + Self::with_bucket_selector(client, move |_, _| bucket.clone()) + } + + /// Create a driver that chooses the bucket for each payload. + pub fn with_bucket_selector( + client: Arc, + bucket_selector: impl Fn(&StorageDriverStoreContext, &Payload) -> String + Send + Sync + 'static, + ) -> Self { + Self { + client, + bucket_selector: Arc::new(bucket_selector), + name: DRIVER_TYPE.to_owned(), + max_payload_size: DEFAULT_MAX_PAYLOAD_SIZE, + } + } + + /// Set the driver instance name written into history. Defaults to `aws.s3driver`. Only change + /// this when registering multiple S3 drivers with different configurations on the same + /// external storage. + pub fn with_driver_name(mut self, name: impl Into) -> Self { + self.name = name.into(); + self + } + + /// Set the maximum serialized payload size, in bytes. Storing a larger payload fails. Defaults + /// to 50 MiB. + pub fn with_max_payload_size( + mut self, + max_payload_size: usize, + ) -> Result { + if max_payload_size == 0 { + return Err(ExternalStorageError::InvalidConfiguration( + "max payload size must be at least 1".to_owned(), + )); + } + self.max_payload_size = max_payload_size; + Ok(self) + } + + fn stored_object( + &self, + context: &StorageDriverStoreContext, + payload: Payload, + context_segments: &str, + ) -> Result { + let bucket = (self.bucket_selector)(context, &payload); + let data = payload.encode_to_vec(); + if data.len() > self.max_payload_size { + return Err(S3StorageDriverError::PayloadTooLarge { + size: data.len(), + max: self.max_payload_size, + }); + } + let hash_value = sha256_hex(&data); + Ok(StoredObject { + key: format!("v0{context_segments}/d/sha256/{hash_value}"), + payload, + bucket, + hash_value, + data, + }) + } +} + +impl StorageDriver for S3StorageDriver { + fn name(&self) -> &str { + &self.name + } + + fn driver_type(&self) -> &str { + DRIVER_TYPE + } + + fn store( + &self, + context: &StorageDriverStoreContext, + payloads: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + let context_segments = context_key_segments(context.target.as_ref()); + let objects: Result, _> = payloads + .into_iter() + .map(|payload| self.stored_object(context, payload, &context_segments)) + .collect(); + let client = self.client.clone(); + let limiter = context.limiter.clone(); + Box::pin(async move { + let objects = objects?; + let claims = objects.iter().map(StoredObject::claim).collect(); + let mut seen = HashSet::new(); + let uploads = objects + .into_iter() + .filter(|object| seen.insert((object.bucket.clone(), object.key.clone()))) + .map(|object| upload_if_absent(client.as_ref(), limiter.as_ref(), object)); + try_join_all(uploads).await?; + Ok(claims) + }) + } + + fn retrieve( + &self, + context: &StorageDriverRetrieveContext, + claims: Vec, + ) -> BoxFuture<'static, Result, ExternalStorageError>> { + let client = self.client.clone(); + let limiter = context.limiter.clone(); + Box::pin(async move { + for claim in &claims { + validate_claim(claim)?; + } + try_join_all( + claims + .iter() + .map(|claim| download(client.as_ref(), limiter.as_ref(), claim)), + ) + .await + }) + } +} + +struct StoredObject { + payload: Payload, + bucket: String, + key: String, + hash_value: String, + data: Vec, +} + +impl StoredObject { + fn claim(&self) -> StorageDriverClaim { + StorageDriverClaim::new( + [ + (CLAIM_BUCKET, self.bucket.as_str()), + (CLAIM_KEY, self.key.as_str()), + (CLAIM_HASH_ALGORITHM, "sha256"), + (CLAIM_HASH_VALUE, self.hash_value.as_str()), + ] + .into_iter() + .map(|(name, value)| (name.to_owned(), value.to_owned())) + .collect(), + ) + } +} + +async fn upload_if_absent( + client: &dyn S3StorageDriverClient, + limiter: &dyn StorageDriverLimiter, + object: StoredObject, +) -> Result<(), ExternalStorageError> { + let _permit = limiter.acquire(&object.payload).await?; + let StoredObject { + bucket, key, data, .. + } = object; + let result: Result<(), BoxError> = async { + if !client.object_exists(&bucket, &key).await? { + client.put_object(&bucket, &key, data).await?; + } + Ok(()) + } + .await; + result.map_err(|source| { + S3StorageDriverError::Store { + client_context: client_context(client), + bucket, + key, + source, + } + .into() + }) +} + +async fn download( + client: &dyn S3StorageDriverClient, + limiter: &dyn StorageDriverLimiter, + claim: &StorageDriverClaim, +) -> Result { + let bucket = &claim.claim_data[CLAIM_BUCKET]; + let key = &claim.claim_data[CLAIM_KEY]; + let data = { + let _permit = limiter.acquire(claim).await?; + client + .get_object(bucket, key) + .await + .map_err(|source| S3StorageDriverError::Retrieve { + bucket: bucket.clone(), + key: key.clone(), + client_context: client_context(client), + source, + })? + }; + let expected = &claim.claim_data[CLAIM_HASH_VALUE]; + let actual = sha256_hex(&data); + if actual != *expected { + return Err(S3StorageDriverError::IntegrityCheck { + bucket: bucket.clone(), + key: key.clone(), + expected: expected.clone(), + actual, + } + .into()); + } + Payload::decode(data.as_slice()).map_err(|err| ExternalStorageError::Driver(Box::new(err))) +} + +fn validate_claim(claim: &StorageDriverClaim) -> Result<(), S3StorageDriverError> { + let value = |name| claim.claim_data.get(name).filter(|value| !value.is_empty()); + let (Some(bucket), Some(key)) = (value(CLAIM_BUCKET), value(CLAIM_KEY)) else { + return Err(S3StorageDriverError::MissingLocation); + }; + let (Some(hash_algorithm), Some(_)) = (value(CLAIM_HASH_ALGORITHM), value(CLAIM_HASH_VALUE)) + else { + return Err(S3StorageDriverError::MissingHash { + bucket: bucket.clone(), + key: key.clone(), + }); + }; + if hash_algorithm != "sha256" { + return Err(S3StorageDriverError::UnsupportedHashAlgorithm { + bucket: bucket.clone(), + key: key.clone(), + hash_algorithm: hash_algorithm.clone(), + }); + } + Ok(()) +} + +fn sha256_hex(data: &[u8]) -> String { + Sha256::digest(data) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn context_key_segments(target: Option<&StorageDriverTargetInfo>) -> String { + match target { + Some(StorageDriverTargetInfo::Workflow(workflow)) => format!( + "/ns/{}/wt/{}/wi/{}/ri/{}", + encode_key_segment(Some(workflow.namespace.as_str())), + encode_key_segment(workflow.workflow_type.as_deref()), + encode_key_segment(workflow.id.as_deref()), + encode_key_segment(workflow.run_id.as_deref()), + ), + Some(StorageDriverTargetInfo::Activity(activity)) => format!( + "/ns/{}/at/{}/ai/{}/ri/{}", + encode_key_segment(Some(activity.namespace.as_str())), + encode_key_segment(activity.activity_type.as_deref()), + encode_key_segment(activity.id.as_deref()), + encode_key_segment(activity.run_id.as_deref()), + ), + _ => String::new(), + } +} + +/// Escapes every UTF-8 byte outside S3's safe character set, so keys are portable across S3 +/// implementations and identical to the keys other SDK languages build. +fn encode_key_segment(value: Option<&str>) -> String { + match value { + None | Some("") => NULL_SEGMENT.to_owned(), + Some(value) => value + .bytes() + .map(|byte| { + if byte.is_ascii_alphanumeric() || SAFE_SPECIAL_CHARACTERS.contains(&byte) { + char::from(byte).to_string() + } else { + format!("%{byte:02X}") + } + }) + .collect(), + } +} + +fn client_context(client: &dyn S3StorageDriverClient) -> String { + client + .describe() + .iter() + .map(|(key, value)| format!(", {key}={value}")) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use futures_util::future::{self, FutureExt}; + use std::{ + collections::HashMap, + sync::{ + Mutex, + atomic::{AtomicUsize, Ordering}, + }, + }; + use temporalio_common::data_converters::{ + StorageDriverActivityInfo, StorageDriverPermit, StorageDriverWorkflowInfo, + }; + use tokio::{ + sync::{OwnedSemaphorePermit, Semaphore, watch}, + time::{Duration, timeout}, + }; + + struct RecordingLimiter { + items: Mutex>, + } + + impl RecordingLimiter { + fn new() -> Self { + Self { + items: Mutex::new(Vec::new()), + } + } + } + + impl StorageDriverLimiter for RecordingLimiter { + fn acquire<'a>( + &'a self, + item: &'a Item, + ) -> BoxFuture<'a, Result> { + self.items.lock().unwrap().push(item.clone()); + future::ready(Ok(StorageDriverPermit::new(Box::new(())))).boxed() + } + } + + /// Counts the permits it has handed out that have not yet been dropped. + struct LivePermitCountingLimiter(Arc); + + struct LivePermit(Arc); + + impl Drop for LivePermit { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } + } + + impl StorageDriverLimiter for LivePermitCountingLimiter { + fn acquire<'a>( + &'a self, + _: &'a Item, + ) -> BoxFuture<'a, Result> { + self.0.fetch_add(1, Ordering::SeqCst); + let permit = StorageDriverPermit::new(Box::new(LivePermit(self.0.clone()))); + future::ready(Ok(permit)).boxed() + } + } + + struct PermitCheckingS3Client { + inner: FakeS3Client, + live_permits: Arc, + request_count: AtomicUsize, + requests_without_permit_count: AtomicUsize, + } + + impl PermitCheckingS3Client { + fn record_request(&self) { + self.request_count.fetch_add(1, Ordering::SeqCst); + if self.live_permits.load(Ordering::SeqCst) == 0 { + self.requests_without_permit_count + .fetch_add(1, Ordering::SeqCst); + } + } + } + + impl S3StorageDriverClient for PermitCheckingS3Client { + fn put_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + data: Vec, + ) -> BoxFuture<'a, Result<(), BoxError>> { + self.record_request(); + self.inner.put_object(bucket, key, data) + } + + fn object_exists<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result> { + self.record_request(); + self.inner.object_exists(bucket, key) + } + + fn get_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result, BoxError>> { + self.record_request(); + self.inner.get_object(bucket, key) + } + } + + struct SemaphoreLimiter(Arc); + + impl StorageDriverLimiter for SemaphoreLimiter { + fn acquire<'a>( + &'a self, + _: &'a Payload, + ) -> BoxFuture<'a, Result> { + async move { + let permit: OwnedSemaphorePermit = self.0.clone().acquire_owned().await.unwrap(); + Ok(StorageDriverPermit::new(Box::new(permit))) + } + .boxed() + } + } + + /// Records the most existence checks in flight at once, holding each one until the expected + /// number are in flight together. + struct ConcurrencyTrackingS3Client { + expected_concurrency: usize, + expected_concurrency_reached: watch::Sender, + in_flight: AtomicUsize, + max_in_flight: AtomicUsize, + request_count: AtomicUsize, + } + + impl S3StorageDriverClient for ConcurrencyTrackingS3Client { + fn put_object<'a>( + &'a self, + _: &'a str, + _: &'a str, + _: Vec, + ) -> BoxFuture<'a, Result<(), BoxError>> { + future::ready(Ok(())).boxed() + } + + fn object_exists<'a>( + &'a self, + _: &'a str, + _: &'a str, + ) -> BoxFuture<'a, Result> { + async move { + self.request_count.fetch_add(1, Ordering::SeqCst); + let in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(in_flight, Ordering::SeqCst); + if in_flight == self.expected_concurrency { + self.expected_concurrency_reached.send_replace(true); + } + let _ = self + .expected_concurrency_reached + .subscribe() + .wait_for(|reached| *reached) + .await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + Ok(false) + } + .boxed() + } + + fn get_object<'a>( + &'a self, + _: &'a str, + _: &'a str, + ) -> BoxFuture<'a, Result, BoxError>> { + unimplemented!() + } + } + + #[derive(Default)] + struct FakeS3Client { + objects: Mutex>>, + put_count: AtomicUsize, + } + + impl S3StorageDriverClient for FakeS3Client { + fn put_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + data: Vec, + ) -> BoxFuture<'a, Result<(), BoxError>> { + self.put_count.fetch_add(1, Ordering::SeqCst); + self.objects + .lock() + .unwrap() + .insert((bucket.to_owned(), key.to_owned()), data); + future::ready(Ok(())).boxed() + } + + fn object_exists<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result> { + let exists = self + .objects + .lock() + .unwrap() + .contains_key(&(bucket.to_owned(), key.to_owned())); + future::ready(Ok(exists)).boxed() + } + + fn get_object<'a>( + &'a self, + bucket: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result, BoxError>> { + let data = self.objects.lock().unwrap()[&(bucket.to_owned(), key.to_owned())].clone(); + future::ready(Ok(data)).boxed() + } + } + + #[derive(Default)] + struct FailingS3Client { + request_count: AtomicUsize, + } + + impl FailingS3Client { + fn fail<'a, T: Send + 'a>(&self) -> BoxFuture<'a, Result> { + self.request_count.fetch_add(1, Ordering::SeqCst); + future::ready(Err("access denied".into())).boxed() + } + } + + impl S3StorageDriverClient for FailingS3Client { + fn put_object<'a>( + &'a self, + _: &'a str, + _: &'a str, + _: Vec, + ) -> BoxFuture<'a, Result<(), BoxError>> { + self.fail() + } + + fn object_exists<'a>( + &'a self, + _: &'a str, + _: &'a str, + ) -> BoxFuture<'a, Result> { + self.fail() + } + + fn get_object<'a>( + &'a self, + _: &'a str, + _: &'a str, + ) -> BoxFuture<'a, Result, BoxError>> { + self.fail() + } + + fn describe(&self) -> HashMap { + HashMap::from([("client_region".to_owned(), "us-west-2".to_owned())]) + } + } + + /// Fails the request for one key once every other request is in flight, and counts the other + /// requests that are dropped before they finish. + struct FailOneS3Client { + failing_key_suffix: String, + other_requests_started: Semaphore, + canceled_count: Arc, + } + + struct CountOnDrop(Arc); + + impl Drop for CountOnDrop { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + impl S3StorageDriverClient for FailOneS3Client { + fn put_object<'a>( + &'a self, + _: &'a str, + _: &'a str, + _: Vec, + ) -> BoxFuture<'a, Result<(), BoxError>> { + future::ready(Ok(())).boxed() + } + + fn object_exists<'a>( + &'a self, + _: &'a str, + key: &'a str, + ) -> BoxFuture<'a, Result> { + async move { + if key.ends_with(&self.failing_key_suffix) { + let _ = self.other_requests_started.acquire_many(2).await; + return Err("fail".into()); + } + let _guard = CountOnDrop(self.canceled_count.clone()); + self.other_requests_started.add_permits(1); + future::pending().await + } + .boxed() + } + + fn get_object<'a>( + &'a self, + _: &'a str, + _: &'a str, + ) -> BoxFuture<'a, Result, BoxError>> { + unimplemented!() + } + } + + fn payload(data: &str) -> Payload { + Payload { + metadata: HashMap::from([("encoding".to_owned(), b"json/plain".to_vec())]), + data: data.as_bytes().to_vec(), + ..Default::default() + } + } + + fn hash_of(data: &str) -> String { + sha256_hex(&payload(data).encode_to_vec()) + } + + fn workflow_target( + namespace: &str, + id: Option<&str>, + run_id: Option<&str>, + workflow_type: Option<&str>, + ) -> StorageDriverTargetInfo { + let mut info = StorageDriverWorkflowInfo::default(); + info.namespace = namespace.to_owned(); + info.id = id.map(str::to_owned); + info.run_id = run_id.map(str::to_owned); + info.workflow_type = workflow_type.map(str::to_owned); + StorageDriverTargetInfo::Workflow(info) + } + + fn my_workflow() -> StorageDriverTargetInfo { + workflow_target( + "my-ns", + Some("my-workflow-id"), + Some("my-run-id"), + Some("MyWorkflow"), + ) + } + + fn store_context(target: Option) -> StorageDriverStoreContext { + store_context_with_limiter(target, Arc::new(RecordingLimiter::new())) + } + + fn store_context_with_limiter( + target: Option, + limiter: Arc>, + ) -> StorageDriverStoreContext { + StorageDriverStoreContext::builder() + .maybe_target(target) + .limiter(limiter) + .build() + } + + fn retrieve_context() -> StorageDriverRetrieveContext { + StorageDriverRetrieveContext::builder() + .limiter(Arc::new(RecordingLimiter::new())) + .build() + } + + fn claim( + bucket: Option<&str>, + key: Option<&str>, + hash_algorithm: Option<&str>, + hash_value: Option<&str>, + ) -> StorageDriverClaim { + StorageDriverClaim::new( + [ + (CLAIM_BUCKET, bucket), + (CLAIM_KEY, key), + (CLAIM_HASH_ALGORITHM, hash_algorithm), + (CLAIM_HASH_VALUE, hash_value), + ] + .into_iter() + .filter_map(|(name, value)| Some((name.to_owned(), value?.to_owned()))) + .collect(), + ) + } + + fn valid_claim() -> StorageDriverClaim { + claim( + Some("my-bucket"), + Some("my-key"), + Some("sha256"), + Some("abc"), + ) + } + + async fn store_key(target: Option, data: &str) -> String { + let driver = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "my-bucket"); + let claims = driver + .store(&store_context(target), vec![payload(data)]) + .await + .unwrap(); + claims[0].claim_data[CLAIM_KEY].clone() + } + + fn driver_error(err: ExternalStorageError) -> S3StorageDriverError { + match err { + ExternalStorageError::Driver(err) => *err.downcast().unwrap(), + err => panic!("expected a driver error, got {err}"), + } + } + + #[test] + fn default_name_is_aws_s3driver() { + // Retrieval is routed by the name recorded in history, so this must stay identical across + // SDKs. The literal is deliberate: a test against the constant would pass if it changed. + let driver = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "b"); + assert_eq!(driver.name(), "aws.s3driver"); + } + + #[test] + fn max_payload_size_of_zero_is_rejected() { + let result = + S3StorageDriver::new(Arc::new(FakeS3Client::default()), "b").with_max_payload_size(0); + assert!(matches!( + result, + Err(ExternalStorageError::InvalidConfiguration(_)) + )); + } + + #[tokio::test] + async fn retrieve_of_stored_claims_returns_the_original_payloads() { + let driver = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "my-bucket"); + let payloads = vec![payload("one"), payload("two")]; + + let claims = driver + .store(&store_context(Some(my_workflow())), payloads.clone()) + .await + .unwrap(); + + assert_eq!( + driver.retrieve(&retrieve_context(), claims).await.unwrap(), + payloads + ); + } + + #[tokio::test] + async fn store_claim_holds_bucket_key_and_sha256_hash() { + let driver = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "my-bucket"); + + let claims = driver + .store(&store_context(Some(my_workflow())), vec![payload("one")]) + .await + .unwrap(); + + let hash_value = hash_of("one"); + assert_eq!( + claims[0].claim_data, + HashMap::from([ + ("bucket".to_owned(), "my-bucket".to_owned()), + ( + "key".to_owned(), + format!( + "v0/ns/my-ns/wt/MyWorkflow/wi/my-workflow-id/ri/my-run-id/d/sha256/{hash_value}" + ) + ), + ("hash_algorithm".to_owned(), "sha256".to_owned()), + ("hash_value".to_owned(), hash_value), + ]) + ); + } + + #[tokio::test] + async fn store_workflow_target_builds_the_spec_example_key() { + let target = workflow_target( + "payments prod", + Some("order+123=abc"), + Some("3f1d6c7a-8b2e-4f7a-9d0a-87a6f95e4d31"), + Some("ChargeWorkflow"), + ); + + assert_eq!( + store_key(Some(target), "one").await, + format!( + "v0/ns/payments%20prod/wt/ChargeWorkflow/wi/order%2B123%3Dabc/ri/\ + 3f1d6c7a-8b2e-4f7a-9d0a-87a6f95e4d31/d/sha256/{}", + hash_of("one") + ) + ); + } + + #[tokio::test] + async fn store_activity_target_builds_the_spec_example_key() { + let mut info = StorageDriverActivityInfo::default(); + info.namespace = "payments prod".to_owned(); + info.id = Some("activity id+42".to_owned()); + info.run_id = Some("9e1d1fd9-2f8a-4c40-93e2-731f31b9268b".to_owned()); + info.activity_type = Some("Capture/Charge".to_owned()); + + assert_eq!( + store_key(Some(StorageDriverTargetInfo::Activity(info)), "one").await, + format!( + "v0/ns/payments%20prod/at/Capture%2FCharge/ai/activity%20id%2B42/ri/\ + 9e1d1fd9-2f8a-4c40-93e2-731f31b9268b/d/sha256/{}", + hash_of("one") + ) + ); + } + + #[tokio::test] + async fn store_key_segment_escapes_every_byte_outside_s3_safe_characters() { + let target = workflow_target( + "AZaz09!-_.*'()", + Some("~é"), + Some("run-1"), + Some("MyWorkflow"), + ); + + assert_eq!( + store_key(Some(target), "one").await, + format!( + "v0/ns/AZaz09!-_.*'()/wt/MyWorkflow/wi/%7E%C3%A9/ri/run-1/d/sha256/{}", + hash_of("one") + ) + ); + } + + #[tokio::test] + async fn store_empty_or_missing_target_values_write_null_segments() { + let target = workflow_target("", None, Some(""), None); + + assert_eq!( + store_key(Some(target), "one").await, + format!( + "v0/ns/null/wt/null/wi/null/ri/null/d/sha256/{}", + hash_of("one") + ) + ); + } + + #[tokio::test] + async fn store_without_target_key_has_only_the_digest() { + assert_eq!( + store_key(None, "one").await, + format!("v0/d/sha256/{}", hash_of("one")) + ); + } + + #[tokio::test] + async fn store_same_payload_twice_in_one_call_uploads_once() { + let client = Arc::new(FakeS3Client::default()); + let driver = S3StorageDriver::new(client.clone(), "my-bucket"); + + let claims = driver + .store( + &store_context(Some(my_workflow())), + vec![payload("same"), payload("same")], + ) + .await + .unwrap(); + + assert_eq!(client.put_count.load(Ordering::SeqCst), 1); + assert_eq!(claims[0], claims[1]); + } + + #[tokio::test] + async fn store_object_already_in_s3_skips_the_upload() { + let client = Arc::new(FakeS3Client::default()); + let driver = S3StorageDriver::new(client.clone(), "my-bucket"); + + for _ in 0..2 { + driver + .store(&store_context(Some(my_workflow())), vec![payload("same")]) + .await + .unwrap(); + } + + assert_eq!(client.put_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn store_bucket_selector_chooses_the_bucket_for_each_payload() { + let client = Arc::new(FakeS3Client::default()); + let driver = S3StorageDriver::with_bucket_selector(client.clone(), |_, payload| { + format!("bucket-{}", String::from_utf8_lossy(&payload.data)) + }); + + let claims = driver + .store( + &store_context(Some(my_workflow())), + vec![payload("a"), payload("b")], + ) + .await + .unwrap(); + + let claimed_buckets: Vec<_> = claims + .iter() + .map(|claim| claim.claim_data[CLAIM_BUCKET].as_str()) + .collect(); + assert_eq!(claimed_buckets, ["bucket-a", "bucket-b"]); + let mut stored_buckets: Vec<_> = client + .objects + .lock() + .unwrap() + .keys() + .map(|(bucket, _)| bucket.clone()) + .collect(); + stored_buckets.sort(); + assert_eq!(stored_buckets, ["bucket-a", "bucket-b"]); + } + + #[tokio::test] + async fn each_s3_request_takes_one_limiter_permit() { + let driver = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "my-bucket"); + let store_limiter = Arc::new(RecordingLimiter::new()); + let retrieve_limiter = Arc::new(RecordingLimiter::new()); + + let claims = driver + .store( + &store_context_with_limiter(Some(my_workflow()), store_limiter.clone()), + vec![payload("one"), payload("two"), payload("one")], + ) + .await + .unwrap(); + driver + .retrieve( + &StorageDriverRetrieveContext::builder() + .limiter(retrieve_limiter.clone()) + .build(), + claims.clone(), + ) + .await + .unwrap(); + + let mut stored = store_limiter.items.lock().unwrap().clone(); + stored.sort_by(|a, b| a.data.cmp(&b.data)); + assert_eq!(stored, [payload("one"), payload("two")]); + assert_eq!(retrieve_limiter.items.lock().unwrap().len(), claims.len()); + } + + #[tokio::test] + async fn every_s3_request_runs_while_holding_a_permit() { + let live_permits = Arc::new(AtomicUsize::new(0)); + let client = Arc::new(PermitCheckingS3Client { + inner: FakeS3Client::default(), + live_permits: live_permits.clone(), + request_count: AtomicUsize::new(0), + requests_without_permit_count: AtomicUsize::new(0), + }); + let driver = S3StorageDriver::new(client.clone(), "my-bucket"); + let limiter = Arc::new(LivePermitCountingLimiter(live_permits)); + + // One payload at a time, so another request's permit cannot hide a request made without one + let claims = driver + .store( + &store_context_with_limiter(None, limiter.clone()), + vec![payload("one")], + ) + .await + .unwrap(); + driver + .retrieve( + &StorageDriverRetrieveContext::builder() + .limiter(limiter) + .build(), + claims, + ) + .await + .unwrap(); + + assert_eq!(client.request_count.load(Ordering::SeqCst), 3); + assert_eq!( + client.requests_without_permit_count.load(Ordering::SeqCst), + 0 + ); + } + + #[tokio::test] + async fn store_runs_as_many_requests_at_once_as_the_limiter_allows() { + let client = Arc::new(ConcurrencyTrackingS3Client { + expected_concurrency: 2, + expected_concurrency_reached: watch::Sender::new(false), + in_flight: AtomicUsize::new(0), + max_in_flight: AtomicUsize::new(0), + request_count: AtomicUsize::new(0), + }); + let driver = S3StorageDriver::new(client.clone(), "my-bucket"); + let context = store_context_with_limiter( + None, + Arc::new(SemaphoreLimiter(Arc::new(Semaphore::new(2)))), + ); + let payloads = ["a", "b", "c", "d"].map(payload).to_vec(); + + // The client holds requests until two are in flight, so a driver that did not fan out + // would never finish + timeout(Duration::from_secs(10), driver.store(&context, payloads)) + .await + .unwrap() + .unwrap(); + + assert_eq!(client.max_in_flight.load(Ordering::SeqCst), 2); + assert_eq!(client.request_count.load(Ordering::SeqCst), 4); + } + + #[tokio::test] + async fn store_payload_over_max_payload_size_fails_and_at_the_limit_succeeds() { + let size = payload("one").encode_to_vec().len(); + let at_limit = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "b") + .with_max_payload_size(size) + .unwrap(); + let below_limit = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "b") + .with_max_payload_size(size - 1) + .unwrap(); + + at_limit + .store(&store_context(None), vec![payload("one")]) + .await + .unwrap(); + let err = below_limit + .store(&store_context(None), vec![payload("one")]) + .await + .unwrap_err(); + assert!(matches!( + driver_error(err), + S3StorageDriverError::PayloadTooLarge { .. } + )); + } + + #[tokio::test] + async fn store_s3_request_failure_reports_bucket_key_and_region() { + let driver = S3StorageDriver::new(Arc::new(FailingS3Client::default()), "my-bucket"); + + let err = driver + .store(&store_context(None), vec![payload("one")]) + .await + .unwrap_err(); + + let err = driver_error(err); + assert_eq!( + err.to_string(), + format!( + "S3StorageDriver store failed [bucket=my-bucket, key=v0/d/sha256/{}, \ + client_region=us-west-2]", + hash_of("one") + ) + ); + assert_eq!( + std::error::Error::source(&err).unwrap().to_string(), + "access denied" + ); + } + + #[tokio::test] + async fn retrieve_s3_request_failure_reports_bucket_key_and_region() { + let driver = S3StorageDriver::new(Arc::new(FailingS3Client::default()), "my-bucket"); + + let err = driver + .retrieve(&retrieve_context(), vec![valid_claim()]) + .await + .unwrap_err(); + + let err = driver_error(err); + assert_eq!( + err.to_string(), + "S3StorageDriver retrieve failed [bucket=my-bucket, key=my-key, client_region=us-west-2]" + ); + assert_eq!( + std::error::Error::source(&err).unwrap().to_string(), + "access denied" + ); + } + + #[tokio::test] + async fn store_one_upload_failing_cancels_the_others_and_returns_the_cause() { + let canceled_count = Arc::new(AtomicUsize::new(0)); + let client = FailOneS3Client { + failing_key_suffix: hash_of("fail"), + other_requests_started: Semaphore::new(0), + canceled_count: canceled_count.clone(), + }; + let driver = S3StorageDriver::new(Arc::new(client), "my-bucket"); + + let err = driver + .store( + &store_context(None), + vec![payload("one"), payload("fail"), payload("two")], + ) + .await + .unwrap_err(); + + let err = driver_error(err); + assert_eq!(std::error::Error::source(&err).unwrap().to_string(), "fail"); + assert_eq!(canceled_count.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn retrieve_data_not_matching_claim_hash_fails() { + let client = Arc::new(FakeS3Client::default()); + let driver = S3StorageDriver::new(client.clone(), "my-bucket"); + let claims = driver + .store(&store_context(None), vec![payload("one")]) + .await + .unwrap(); + client.objects.lock().unwrap().insert( + ( + "my-bucket".to_owned(), + claims[0].claim_data[CLAIM_KEY].clone(), + ), + payload("tampered").encode_to_vec(), + ); + + let err = driver + .retrieve(&retrieve_context(), claims) + .await + .unwrap_err(); + + assert!(matches!( + driver_error(err), + S3StorageDriverError::IntegrityCheck { .. } + )); + } + + #[tokio::test] + async fn retrieve_incomplete_or_unsupported_claim_fails_before_any_s3_request() { + let client = Arc::new(FailingS3Client::default()); + let driver = S3StorageDriver::new(client.clone(), "my-bucket"); + let invalid_claims = [ + claim(None, Some("my-key"), Some("sha256"), Some("abc")), + claim(Some("my-bucket"), Some(""), Some("sha256"), Some("abc")), + claim(Some("my-bucket"), Some("my-key"), None, Some("abc")), + claim(Some("my-bucket"), Some("my-key"), Some("sha256"), None), + claim(Some("my-bucket"), Some("my-key"), Some("md5"), Some("abc")), + ]; + + for invalid_claim in invalid_claims { + driver + .retrieve(&retrieve_context(), vec![valid_claim(), invalid_claim]) + .await + .unwrap_err(); + } + assert_eq!(client.request_count.load(Ordering::SeqCst), 0); + } +} diff --git a/crates/aws-s3driver/src/lib.rs b/crates/aws-s3driver/src/lib.rs new file mode 100644 index 000000000..c4a9e6668 --- /dev/null +++ b/crates/aws-s3driver/src/lib.rs @@ -0,0 +1,21 @@ +#![warn(missing_docs)] // error if there are missing docs +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! Amazon S3 storage driver for Temporal external storage. +//! +//! [`S3StorageDriver`] has no AWS dependency. It performs requests through an +//! [`S3StorageDriverClient`], which you can implement for any S3 library. The `aws-sdk` feature, +//! enabled by default, adds [`AwsSdkS3StorageDriverClient`] for the official AWS SDK. Disable +//! default features to bring your own client without depending on the AWS SDK. +//! +//! This crate is experimental and may change in the future. + +#[cfg(feature = "aws-sdk")] +mod aws_sdk_client; +mod client; +mod driver; + +#[cfg(feature = "aws-sdk")] +pub use aws_sdk_client::AwsSdkS3StorageDriverClient; +pub use client::{BoxError, S3StorageDriverClient}; +pub use driver::{S3StorageDriver, S3StorageDriverError}; From 66f619925b3232ff2f78813a5da8fafe12f8ead5 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 28 Oct 2026 15:12:12 -0400 Subject: [PATCH 2/4] fix error types and tests --- crates/aws-s3driver/src/driver.rs | 65 +++++++++---------------------- 1 file changed, 18 insertions(+), 47 deletions(-) diff --git a/crates/aws-s3driver/src/driver.rs b/crates/aws-s3driver/src/driver.rs index fc4f611f5..ae9b1db36 100644 --- a/crates/aws-s3driver/src/driver.rs +++ b/crates/aws-s3driver/src/driver.rs @@ -2,10 +2,10 @@ use crate::client::{BoxError, S3StorageDriverClient}; use futures_util::future::{BoxFuture, try_join_all}; use prost::Message; use sha2::{Digest, Sha256}; -use std::{collections::HashSet, sync::Arc}; +use std::{collections::HashSet, num::NonZeroUsize, sync::Arc}; use temporalio_common::{ data_converters::{ - ExternalStorageError, StorageDriver, StorageDriverClaim, StorageDriverLimiter, + StorageDriver, StorageDriverClaim, StorageDriverError, StorageDriverLimiter, StorageDriverRetrieveContext, StorageDriverStoreContext, StorageDriverTargetInfo, }, protos::temporal::api::common::v1::Payload, @@ -22,7 +22,7 @@ const SAFE_SPECIAL_CHARACTERS: &[u8] = b"!-_.*'()"; type BucketSelector = dyn Fn(&StorageDriverStoreContext, &Payload) -> String + Send + Sync; -/// Errors raised by [`S3StorageDriver`], returned inside [`ExternalStorageError::Driver`]. +/// Errors raised by [`S3StorageDriver`], returned inside [`StorageDriverError`]. #[derive(Debug, thiserror::Error)] #[non_exhaustive] pub enum S3StorageDriverError { @@ -107,12 +107,6 @@ pub enum S3StorageDriverError { }, } -impl From for ExternalStorageError { - fn from(err: S3StorageDriverError) -> Self { - ExternalStorageError::Driver(Box::new(err)) - } -} - /// Storage driver that stores and retrieves payloads in Amazon S3. /// /// Payloads are content addressed by a SHA-256 hash of their serialized bytes, so storing the same @@ -159,17 +153,9 @@ impl S3StorageDriver { /// Set the maximum serialized payload size, in bytes. Storing a larger payload fails. Defaults /// to 50 MiB. - pub fn with_max_payload_size( - mut self, - max_payload_size: usize, - ) -> Result { - if max_payload_size == 0 { - return Err(ExternalStorageError::InvalidConfiguration( - "max payload size must be at least 1".to_owned(), - )); - } - self.max_payload_size = max_payload_size; - Ok(self) + pub fn with_max_payload_size(mut self, max_payload_size: NonZeroUsize) -> Self { + self.max_payload_size = max_payload_size.get(); + self } fn stored_object( @@ -210,7 +196,7 @@ impl StorageDriver for S3StorageDriver { &self, context: &StorageDriverStoreContext, payloads: Vec, - ) -> BoxFuture<'static, Result, ExternalStorageError>> { + ) -> BoxFuture<'static, Result, StorageDriverError>> { let context_segments = context_key_segments(context.target.as_ref()); let objects: Result, _> = payloads .into_iter() @@ -235,7 +221,7 @@ impl StorageDriver for S3StorageDriver { &self, context: &StorageDriverRetrieveContext, claims: Vec, - ) -> BoxFuture<'static, Result, ExternalStorageError>> { + ) -> BoxFuture<'static, Result, StorageDriverError>> { let client = self.client.clone(); let limiter = context.limiter.clone(); Box::pin(async move { @@ -280,7 +266,7 @@ async fn upload_if_absent( client: &dyn S3StorageDriverClient, limiter: &dyn StorageDriverLimiter, object: StoredObject, -) -> Result<(), ExternalStorageError> { +) -> Result<(), StorageDriverError> { let _permit = limiter.acquire(&object.payload).await?; let StoredObject { bucket, key, data, .. @@ -307,7 +293,7 @@ async fn download( client: &dyn S3StorageDriverClient, limiter: &dyn StorageDriverLimiter, claim: &StorageDriverClaim, -) -> Result { +) -> Result { let bucket = &claim.claim_data[CLAIM_BUCKET]; let key = &claim.claim_data[CLAIM_KEY]; let data = { @@ -333,7 +319,7 @@ async fn download( } .into()); } - Payload::decode(data.as_slice()).map_err(|err| ExternalStorageError::Driver(Box::new(err))) + Ok(Payload::decode(data.as_slice())?) } fn validate_claim(claim: &StorageDriverClaim) -> Result<(), S3StorageDriverError> { @@ -446,7 +432,7 @@ mod tests { fn acquire<'a>( &'a self, item: &'a Item, - ) -> BoxFuture<'a, Result> { + ) -> BoxFuture<'a, Result> { self.items.lock().unwrap().push(item.clone()); future::ready(Ok(StorageDriverPermit::new(Box::new(())))).boxed() } @@ -467,7 +453,7 @@ mod tests { fn acquire<'a>( &'a self, _: &'a Item, - ) -> BoxFuture<'a, Result> { + ) -> BoxFuture<'a, Result> { self.0.fetch_add(1, Ordering::SeqCst); let permit = StorageDriverPermit::new(Box::new(LivePermit(self.0.clone()))); future::ready(Ok(permit)).boxed() @@ -527,7 +513,7 @@ mod tests { fn acquire<'a>( &'a self, _: &'a Payload, - ) -> BoxFuture<'a, Result> { + ) -> BoxFuture<'a, Result> { async move { let permit: OwnedSemaphorePermit = self.0.clone().acquire_owned().await.unwrap(); Ok(StorageDriverPermit::new(Box::new(permit))) @@ -819,11 +805,8 @@ mod tests { claims[0].claim_data[CLAIM_KEY].clone() } - fn driver_error(err: ExternalStorageError) -> S3StorageDriverError { - match err { - ExternalStorageError::Driver(err) => *err.downcast().unwrap(), - err => panic!("expected a driver error, got {err}"), - } + fn driver_error(err: StorageDriverError) -> S3StorageDriverError { + err.into_source().downcast().unwrap() } #[test] @@ -834,16 +817,6 @@ mod tests { assert_eq!(driver.name(), "aws.s3driver"); } - #[test] - fn max_payload_size_of_zero_is_rejected() { - let result = - S3StorageDriver::new(Arc::new(FakeS3Client::default()), "b").with_max_payload_size(0); - assert!(matches!( - result, - Err(ExternalStorageError::InvalidConfiguration(_)) - )); - } - #[tokio::test] async fn retrieve_of_stored_claims_returns_the_original_payloads() { let driver = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "my-bucket"); @@ -1122,11 +1095,9 @@ mod tests { async fn store_payload_over_max_payload_size_fails_and_at_the_limit_succeeds() { let size = payload("one").encode_to_vec().len(); let at_limit = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "b") - .with_max_payload_size(size) - .unwrap(); + .with_max_payload_size(NonZeroUsize::new(size).unwrap()); let below_limit = S3StorageDriver::new(Arc::new(FakeS3Client::default()), "b") - .with_max_payload_size(size - 1) - .unwrap(); + .with_max_payload_size(NonZeroUsize::new(size - 1).unwrap()); at_limit .store(&store_context(None), vec![payload("one")]) From e1025cb11f03263629721c2d3a05ea5895ebdf78 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 28 Oct 2026 17:31:29 -0400 Subject: [PATCH 3/4] sort payload metadata by key before creating hash for storage key --- crates/aws-s3driver/src/driver.rs | 57 +++++++++++++++++++++++++++++-- 1 file changed, 54 insertions(+), 3 deletions(-) diff --git a/crates/aws-s3driver/src/driver.rs b/crates/aws-s3driver/src/driver.rs index ae9b1db36..c15d1e5ed 100644 --- a/crates/aws-s3driver/src/driver.rs +++ b/crates/aws-s3driver/src/driver.rs @@ -2,7 +2,11 @@ use crate::client::{BoxError, S3StorageDriverClient}; use futures_util::future::{BoxFuture, try_join_all}; use prost::Message; use sha2::{Digest, Sha256}; -use std::{collections::HashSet, num::NonZeroUsize, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + num::NonZeroUsize, + sync::Arc, +}; use temporalio_common::{ data_converters::{ StorageDriver, StorageDriverClaim, StorageDriverError, StorageDriverLimiter, @@ -161,11 +165,11 @@ impl S3StorageDriver { fn stored_object( &self, context: &StorageDriverStoreContext, - payload: Payload, + mut payload: Payload, context_segments: &str, ) -> Result { let bucket = (self.bucket_selector)(context, &payload); - let data = payload.encode_to_vec(); + let data = encode_payload(&mut payload); if data.len() > self.max_payload_size { return Err(S3StorageDriverError::PayloadTooLarge { size: data.len(), @@ -344,6 +348,27 @@ fn validate_claim(claim: &StorageDriverClaim) -> Result<(), S3StorageDriverError Ok(()) } +/// Encodes a payload with its metadata entries sorted by key. +fn encode_payload(payload: &mut Payload) -> Vec { + let mut metadata: Vec<_> = std::mem::take(&mut payload.metadata).into_iter().collect(); + metadata.sort_unstable_by(|(a, _), (b, _)| a.cmp(b)); + let mut data = Vec::new(); + for (key, value) in &metadata { + let entry = Payload { + metadata: HashMap::from([(key.clone(), value.clone())]), + ..Default::default() + }; + entry + .encode(&mut data) + .expect("encoding into a Vec cannot run out of space"); + } + payload + .encode(&mut data) + .expect("encoding into a Vec cannot run out of space"); + payload.metadata = metadata.into_iter().collect(); + data +} + fn sha256_hex(data: &[u8]) -> String { Sha256::digest(data) .iter() @@ -952,6 +977,32 @@ mod tests { assert_eq!(claims[0], claims[1]); } + #[tokio::test] + async fn store_equal_payloads_with_several_metadata_entries_uploads_once() { + // Each HashMap iterates in its own random order, so with this many entries two equal + // payloads would almost never encode the same way unless the driver sorts them + let entries: Vec<_> = (0..8u8).map(|i| (format!("key-{i}"), vec![i])).collect(); + let equal_payload = || Payload { + metadata: entries.iter().cloned().collect(), + data: b"same".to_vec(), + ..Default::default() + }; + let client = Arc::new(FakeS3Client::default()); + let driver = S3StorageDriver::new(client.clone(), "my-bucket"); + + let claims = driver + .store(&store_context(None), vec![equal_payload(), equal_payload()]) + .await + .unwrap(); + + assert_eq!(client.put_count.load(Ordering::SeqCst), 1); + assert_eq!(claims[0], claims[1]); + assert_eq!( + driver.retrieve(&retrieve_context(), claims).await.unwrap(), + [equal_payload(), equal_payload()] + ); + } + #[tokio::test] async fn store_object_already_in_s3_skips_the_upload() { let client = Arc::new(FakeS3Client::default()); From 66a6a6f05294bf2b331f1379a4d6cc23b28cf26a Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 28 Oct 2026 17:38:11 -0400 Subject: [PATCH 4/4] rename encode -> serialize --- crates/aws-s3driver/src/driver.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/crates/aws-s3driver/src/driver.rs b/crates/aws-s3driver/src/driver.rs index c15d1e5ed..5b683d3cf 100644 --- a/crates/aws-s3driver/src/driver.rs +++ b/crates/aws-s3driver/src/driver.rs @@ -169,7 +169,7 @@ impl S3StorageDriver { context_segments: &str, ) -> Result { let bucket = (self.bucket_selector)(context, &payload); - let data = encode_payload(&mut payload); + let data = serialize_payload(&mut payload); if data.len() > self.max_payload_size { return Err(S3StorageDriverError::PayloadTooLarge { size: data.len(), @@ -348,8 +348,8 @@ fn validate_claim(claim: &StorageDriverClaim) -> Result<(), S3StorageDriverError Ok(()) } -/// Encodes a payload with its metadata entries sorted by key. -fn encode_payload(payload: &mut Payload) -> Vec { +/// Serializes a payload with its metadata entries sorted by key. +fn serialize_payload(payload: &mut Payload) -> Vec { let mut metadata: Vec<_> = std::mem::take(&mut payload.metadata).into_iter().collect(); metadata.sort_unstable_by(|(a, _), (b, _)| a.cmp(b)); let mut data = Vec::new(); @@ -979,8 +979,6 @@ mod tests { #[tokio::test] async fn store_equal_payloads_with_several_metadata_entries_uploads_once() { - // Each HashMap iterates in its own random order, so with this many entries two equal - // payloads would almost never encode the same way unless the driver sorts them let entries: Vec<_> = (0..8u8).map(|i| (format!("key-{i}"), vec![i])).collect(); let equal_payload = || Payload { metadata: entries.iter().cloned().collect(),