diff --git a/src/storage/metrics_layer.rs b/src/storage/metrics_layer.rs index 297e65eec..ee47a37cf 100644 --- a/src/storage/metrics_layer.rs +++ b/src/storage/metrics_layer.rs @@ -26,8 +26,8 @@ use async_trait::async_trait; use bytes::Bytes; use futures_util::{Stream, StreamExt, stream::BoxStream}; use object_store::{ - CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, - PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, + Attribute, CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as ObjectStoreResult, path::Path, }; @@ -93,6 +93,7 @@ pub fn error_to_status_code(err: &object_store::Error) -> &'static str { pub struct MetricLayer { inner: T, provider: String, + cache_control_no_store: bool, } impl MetricLayer { @@ -100,6 +101,20 @@ impl MetricLayer { Self { inner, provider: provider.to_string(), + cache_control_no_store: false, + } + } + + pub fn with_cache_control_no_store(mut self, enabled: bool) -> Self { + self.cache_control_no_store = enabled; + self + } + + fn set_cache_control(&self, location: &Path, attributes: &mut object_store::Attributes) { + // Parquet data is immutable and may be cached, while mutable metadata + // (manifests, schemas, stream metadata, etc.) must not be cached. + if self.cache_control_no_store && !location.as_ref().ends_with(".parquet") { + attributes.insert(Attribute::CacheControl, "no-store".into()); } } } @@ -116,8 +131,9 @@ impl ObjectStore for MetricLayer { &self, location: &Path, payload: PutPayload, - opts: PutOptions, + mut opts: PutOptions, ) -> ObjectStoreResult { + self.set_cache_control(location, &mut opts.attributes); let _guard = InflightGuard::new(&self.provider, "PUT"); let time = time::Instant::now(); let put_result = self.inner.put_opts(location, payload, opts).await; @@ -137,8 +153,9 @@ impl ObjectStore for MetricLayer { async fn put_multipart_opts( &self, location: &Path, - opts: PutMultipartOptions, + mut opts: PutMultipartOptions, ) -> ObjectStoreResult> { + self.set_cache_control(location, &mut opts.attributes); let _guard = InflightGuard::new(&self.provider, "PUT_MULTIPART"); let time = time::Instant::now(); let result = self.inner.put_multipart_opts(location, opts).await; @@ -324,3 +341,43 @@ impl Stream for StreamMetricWrapper<'_, T> { } } } + +#[cfg(test)] +mod tests { + use object_store::{Attribute, Attributes, memory::InMemory, path::Path}; + + use super::MetricLayer; + + #[test] + fn s3_metadata_uploads_disable_caching() { + let layer = MetricLayer::new(InMemory::new(), "s3").with_cache_control_no_store(true); + let mut attributes = Attributes::new(); + + layer.set_cache_control(&Path::from("stream/stream.json"), &mut attributes); + + assert_eq!( + attributes.get(&Attribute::CacheControl).map(AsRef::as_ref), + Some("no-store") + ); + } + + #[test] + fn s3_parquet_uploads_remain_cacheable() { + let layer = MetricLayer::new(InMemory::new(), "s3").with_cache_control_no_store(true); + let mut attributes = Attributes::new(); + + layer.set_cache_control(&Path::from("stream/events.parquet"), &mut attributes); + + assert_eq!(attributes.get(&Attribute::CacheControl), None); + } + + #[test] + fn cache_control_is_disabled_by_default() { + let layer = MetricLayer::new(InMemory::new(), "s3"); + let mut attributes = Attributes::new(); + + layer.set_cache_control(&Path::from("stream/stream.json"), &mut attributes); + + assert_eq!(attributes.get(&Attribute::CacheControl), None); + } +} diff --git a/src/storage/s3.rs b/src/storage/s3.rs index 6d47c3da9..24388dba1 100644 --- a/src/storage/s3.rs +++ b/src/storage/s3.rs @@ -122,6 +122,15 @@ pub struct S3Config { )] pub set_checksum: bool, + /// Add `Cache-Control: no-store` to non-Parquet object uploads. + #[arg( + long, + env = "P_S3_CACHE_CONTROL_NO_STORE", + value_name = "bool", + default_value = "false" + )] + pub cache_control_no_store: bool, + /// Set client to use virtual hosted style acess #[arg( long, @@ -335,7 +344,8 @@ impl ObjectStorageProvider for S3Config { // limit objectstore to a concurrent request limit let s3 = LimitStore::new(s3, super::MAX_OBJECT_STORE_REQUESTS); - let s3 = MetricLayer::new(s3, "s3"); + let s3 = + MetricLayer::new(s3, "s3").with_cache_control_no_store(self.cache_control_no_store); let object_store_registry = DefaultObjectStoreRegistry::new(); let url = ObjectStoreUrl::parse(format!("s3://{}", self.bucket_name)).unwrap(); @@ -348,7 +358,8 @@ impl ObjectStorageProvider for S3Config { let s3 = self.get_default_builder().build().unwrap(); // limit objectstore to a concurrent request limit let s3 = LimitStore::new(s3, super::MAX_OBJECT_STORE_REQUESTS); - let s3 = MetricLayer::new(s3, "s3"); + let s3 = + MetricLayer::new(s3, "s3").with_cache_control_no_store(self.cache_control_no_store); Arc::new(S3 { client: Arc::new(s3), bucket: self.bucket_name.clone(),