diff --git a/Cargo.lock b/Cargo.lock index 668797c..b5ceb6b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -126,7 +126,7 @@ checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" dependencies = [ "proc-macro2", "quote", - "syn 3.0.2", + "syn 3.0.3", ] [[package]] @@ -348,9 +348,9 @@ checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "cc" -version = "1.3.0" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" +checksum = "5add81bb678e6cb321aff7fa0dc7689ad82b112dbc032cea19f91d6b8e3582b9" dependencies = [ "find-msvc-tools", "jobserver", @@ -424,9 +424,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.3" +version = "4.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fb99565819980999fb7b4a1796046a5c949e6d4ff132cf5fadf5a641e20d776" +checksum = "d91e0c145792ef73a6ad36d27c75ac09f1832222a3c209689d90f534685ee5b7" dependencies = [ "clap_builder", "clap_derive", @@ -446,14 +446,14 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.6.3" +version = "4.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32f2392eae7f16557a3d727ef3a12e57b2b2ca6f98566a5f4fb41ffe305df077" +checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061" dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -738,9 +738,9 @@ checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" [[package]] name = "either" -version = "1.16.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" dependencies = [ "serde", ] @@ -779,11 +779,10 @@ dependencies = [ [[package]] name = "event-listener" -version = "5.4.1" +version = "5.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" dependencies = [ - "concurrent-queue", "parking", "pin-project-lite", ] @@ -1556,9 +1555,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.188" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22053b6a34f84abc97f9129e61334f40174659a1b9bd18c970b83db6a9a6348b" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "libsqlite3-sys" @@ -2755,9 +2754,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.15.0" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" dependencies = [ "web-time", "zeroize", @@ -3009,7 +3008,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.2", + "syn 3.0.3", ] [[package]] @@ -3416,9 +3415,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.2" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a207d6d6a2b7fc470b80443726053f18a2481b7e1eee970597051596567987a3" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" dependencies = [ "proc-macro2", "quote", @@ -3487,7 +3486,7 @@ checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" dependencies = [ "proc-macro2", "quote", - "syn 3.0.2", + "syn 3.0.3", ] [[package]] @@ -3593,9 +3592,9 @@ dependencies = [ [[package]] name = "tokio-stream" -version = "0.1.18" +version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" dependencies = [ "futures-core", "pin-project-lite", @@ -4397,7 +4396,7 @@ dependencies = [ "smallvec", "sqlx-core", "sqlx-sqlite", - "syn 3.0.2", + "syn 3.0.3", "thiserror", "tokio", "tokio-stream", diff --git a/Cargo.toml b/Cargo.toml index 7fac6ef..d4c72dd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,6 +60,7 @@ serde_json = { version = "1.0.149" } snowflaked = { version = "1.0.3", features = ["sync"] } sqlformat = { version = "0.5.0" } sqlx = { version = "0.9.0", features = ["chrono", "derive", "json", "macros", "migrate", "runtime-tokio", "sqlite"] } +tempfile = "3.14" thiserror = { version = "2.0.18" } tokio = { version = "1.52.1", features = ["macros", "signal", "rt-multi-thread"] } tokio-tungstenite = { version = "0.30.0", features = ["connect", "stream"] } diff --git a/opsqueue/src/object_store/mod.rs b/opsqueue/src/object_store/mod.rs index 12ee43e..dbc83f0 100644 --- a/opsqueue/src/object_store/mod.rs +++ b/opsqueue/src/object_store/mod.rs @@ -8,6 +8,13 @@ use object_store::path::Path; use reqwest::Url; use ux::u63; +// Re-export the `object_store` types that appear in this module's public API +// (`new_with_retry`'s parameters and `ChunkRetrievalError`/`ChunkStorageError`'s +// `source` field). Downstream crates should name these through opsqueue instead +// of taking their own direct dependency on the `object_store` crate, which would +// have to be kept version-locked with the one opsqueue links against. +pub use object_store::{BackoffConfig, Error, RetryConfig}; + /// A client for interacting with an object store. /// /// This exists as a separate type, so we can build it _once_ @@ -106,6 +113,9 @@ impl ObjectStoreClient { /// The given `object_store_url` recognizes the formats detailed [here](https://docs.rs/object_store/0.11.1/object_store/enum.ObjectStoreScheme.html#method.parse). /// Most importantly, we support GCS (for production usage) and local file systems (for testing). /// + /// Uses the object store's default transport-level retry configuration; use + /// [`ObjectStoreClient::new_with_retry`] to override it. + /// /// # Errors /// /// Returns an error if the URL cannot be parsed or if object store initialization fails. @@ -113,8 +123,77 @@ impl ObjectStoreClient { object_store_url: &str, options: Vec<(String, String)>, ) -> Result { + Self::new_with_retry(object_store_url, options, RetryConfig::default()) + } + + /// Like [`ObjectStoreClient::new`] but overrides the transport-level + /// retry configuration used by the underlying object store client. + /// + /// The retry configuration is applied for backends whose builder exposes + /// [`with_retry`](object_store::gcp::GoogleCloudStorageBuilder::with_retry) + /// (currently GCS and HTTP). For local and in-memory backends the retry + /// config has no effect and is silently ignored, matching the behaviour + /// of [`ObjectStoreClient::new`]. + /// + /// # Errors + /// + /// Returns an error if the URL cannot be parsed or if object store + /// initialization fails. + pub fn new_with_retry( + object_store_url: &str, + options: Vec<(String, String)>, + retry_config: RetryConfig, + ) -> Result { + use object_store::ClientConfigKey; + use object_store::ObjectStoreScheme; + use object_store::gcp::{GoogleCloudStorageBuilder, GoogleConfigKey}; + use object_store::http::HttpBuilder; + use std::str::FromStr; + let url = Url::parse(object_store_url)?; - let (object_store, base_path) = object_store::parse_url_opts(&url, options)?; + let (scheme, raw_path) = + ObjectStoreScheme::parse(&url).map_err(object_store::Error::from)?; + let base_path = Path::parse(raw_path).map_err(object_store::Error::from)?; + + let object_store: Box = match scheme { + ObjectStoreScheme::GoogleCloudStorage => { + let builder = options.into_iter().fold( + GoogleCloudStorageBuilder::new() + .with_url(object_store_url.to_string()) + .with_retry(retry_config), + |builder, (key, value)| match GoogleConfigKey::from_str( + &key.to_ascii_lowercase(), + ) { + Ok(config_key) => builder.with_config(config_key, value), + Err(_) => builder, + }, + ); + Box::new(builder.build()?) + } + ObjectStoreScheme::Http => { + let url_without_path = &url[..url::Position::BeforePath]; + let builder = options.into_iter().fold( + HttpBuilder::new() + .with_url(url_without_path) + .with_retry(retry_config), + |builder, (key, value)| match ClientConfigKey::from_str( + &key.to_ascii_lowercase(), + ) { + Ok(config_key) => builder.with_config(config_key, value), + Err(_) => builder, + }, + ); + Box::new(builder.build()?) + } + _ => { + // Retry configuration is not meaningful for local/in-memory + // backends; fall back to the default construction path so + // that the `options` list is still honoured for those cases. + let (store, _) = object_store::parse_url_opts(&url, options)?; + store + } + }; + Ok(ObjectStoreClient(Arc::new(ObjectStoreClientInner { url: object_store_url.into(), object_store, @@ -264,3 +343,151 @@ impl ObjectStoreClient { &self.0.url } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::chunk::ChunkIndex; + use object_store::{BackoffConfig, RetryConfig}; + use std::time::Duration; + use ux::u63; + + fn tight_retry_config() -> RetryConfig { + RetryConfig { + backoff: BackoffConfig { + init_backoff: Duration::from_millis(1), + max_backoff: Duration::from_millis(2), + base: 2.0, + }, + max_retries: 3, + retry_timeout: Duration::from_secs(1), + } + } + + #[tokio::test] + async fn new_with_retry_supports_in_memory_backend() { + // Non-cloud backends must keep working: `new_with_retry` should build + // successfully and behave identically to `new` for `memory://` URLs. + let client = + ObjectStoreClient::new_with_retry("memory:///", Vec::new(), tight_retry_config()) + .expect("memory:// URL should build with new_with_retry"); + + let chunk_index: ChunkIndex = u63::new(0).into(); + client + .store_chunk("test-prefix", chunk_index, ChunkType::Input, b"hi".to_vec()) + .await + .expect("store on in-memory backend should succeed"); + + let retrieved = client + .retrieve_chunk("test-prefix", chunk_index, ChunkType::Input) + .await + .expect("retrieve on in-memory backend should succeed"); + assert_eq!(retrieved, b"hi"); + } + + #[tokio::test] + async fn new_with_retry_supports_local_filesystem_backend() { + // Local filesystem paths follow the same fallback path as memory://: + // retry config is silently ignored but everything else must keep + // working. This guards the fallback branch of `new_with_retry`. + let tmp = tempfile::tempdir().expect("failed to create tmp dir"); + let url = format!("file://{}", tmp.path().display()); + + let client = ObjectStoreClient::new_with_retry(&url, Vec::new(), tight_retry_config()) + .expect("file:// URL should build with new_with_retry"); + + let chunk_index: ChunkIndex = u63::new(0).into(); + client + .store_chunk("test-prefix", chunk_index, ChunkType::Input, b"hi".to_vec()) + .await + .expect("store on local backend should succeed"); + + let retrieved = client + .retrieve_chunk("test-prefix", chunk_index, ChunkType::Input) + .await + .expect("retrieve on local backend should succeed"); + assert_eq!(retrieved, b"hi"); + } + + #[test] + fn new_with_retry_rejects_malformed_url() { + // Regression guard: bad URLs must still surface as + // `UrlParseFailure` and not, e.g., panic while classifying the + // scheme. + let err = ObjectStoreClient::new_with_retry("not a url", Vec::new(), tight_retry_config()) + .expect_err("malformed URL should not build"); + assert!( + matches!(err, NewObjectStoreClientError::UrlParseFailure(_)), + "unexpected error variant: {err:?}" + ); + } + + /// A minimal service account JSON that carries `disable_oauth: true`, so + /// `GoogleCloudStorageBuilder::build()` succeeds without talking to any + /// external auth service. Mirrors the fake key used inside `object_store`'s + /// own test suite. + const FAKE_GCS_SERVICE_ACCOUNT_KEY: &str = r#"{ + "private_key": "private_key", + "private_key_id": "private_key_id", + "client_email": "client_email", + "disable_oauth": true + }"#; + + #[test] + fn new_with_retry_builds_http_backend_without_network_io() { + // The HTTP arm of `new_with_retry` goes through `HttpBuilder`, which + // is a different code path from the `parse_url_opts` fallback covered + // by the memory:// and file:// tests. Building the client does not + // touch the network, so we can point it at an unreachable localhost + // port and still assert that construction succeeds. + let client = ObjectStoreClient::new_with_retry( + "http://127.0.0.1:1/some/prefix", + Vec::new(), + tight_retry_config(), + ) + .expect("http:// URL should build with new_with_retry"); + // Sanity-check the base path was captured (the URL segment after the + // authority). + assert!(format!("{client:?}").contains("some/prefix")); + } + + #[test] + fn new_with_retry_builds_gcs_backend_with_disabled_oauth() { + // The GCS arm of `new_with_retry` goes through `GoogleCloudStorageBuilder`, + // which is the branch that motivated this whole change (issue #3883). + // We feed it a service-account JSON with `disable_oauth: true` so the + // builder skips real authentication and `build()` completes without + // network I/O. + let options = vec![( + "service_account_key".to_string(), + FAKE_GCS_SERVICE_ACCOUNT_KEY.to_string(), + )]; + let client = ObjectStoreClient::new_with_retry( + "gs://some-bucket/some/prefix", + options, + tight_retry_config(), + ) + .expect("gs:// URL should build with new_with_retry"); + assert!(format!("{client:?}").contains("some/prefix")); + } + + #[test] + fn new_with_retry_forwards_options_to_gcs_builder_and_ignores_unknown_keys() { + // Regression guard for the `options.into_iter().fold(...)` in the GCS + // arm: if a future refactor stops forwarding options to the builder, + // `service_account_key` below would no longer disable OAuth and + // `build()` would fall back to ADC lookup — which we can't rely on in + // CI. If instead the fold turned unknown keys into hard errors, + // `totally_unknown_key` would fail the build. Both directions of + // regression are caught by asserting `Ok(_)` here. + let options = vec![ + ( + "service_account_key".to_string(), + FAKE_GCS_SERVICE_ACCOUNT_KEY.to_string(), + ), + ("totally_unknown_key".to_string(), "some-value".to_string()), + ]; + ObjectStoreClient::new_with_retry("gs://some-bucket/", options, tight_retry_config()) + .expect("gs:// URL with mixed known+unknown options should build"); + } +}