From 353f0e2fcffa4dbf6d7828a7979d85652af1d660 Mon Sep 17 00:00:00 2001 From: Andrey Chudnovskiy Date: Tue, 14 Jul 2026 14:15:23 +0000 Subject: [PATCH] Expose replication + object-storage primitives for building WAL-processing tools - Open more of tls.rs and conn.rs as pub - Add PutIfAbsent method to storage interface, S3 and GCS The change allows to use wal-rus PG connection and storage code as a library --- src/pg/replication/tls.rs | 2 +- src/storage/gcs.rs | 46 +++++++++- src/storage/mod.rs | 28 ++++++ src/storage/retrying.rs | 107 ++++++++++++++++++++- src/storage/s3.rs | 189 +++++++++++++++++++++++++++++++++++++- 5 files changed, 368 insertions(+), 4 deletions(-) diff --git a/src/pg/replication/tls.rs b/src/pg/replication/tls.rs index 3661552..cac4b07 100644 --- a/src/pg/replication/tls.rs +++ b/src/pg/replication/tls.rs @@ -156,7 +156,7 @@ pub async fn maybe_upgrade( } } -fn build_client_config(sslmode: SslMode, tls: &TlsParams) -> Result { +pub fn build_client_config(sslmode: SslMode, tls: &TlsParams) -> Result { let provider = rustls::crypto::aws_lc_rs::default_provider(); let builder = ClientConfig::builder_with_provider(Arc::new(provider)) .with_safe_default_protocol_versions() diff --git a/src/storage/gcs.rs b/src/storage/gcs.rs index 2c28b9f..84a914d 100644 --- a/src/storage/gcs.rs +++ b/src/storage/gcs.rs @@ -22,7 +22,10 @@ use serde::Deserialize; use tokio::sync::Mutex; use tokio_util::io::ReaderStream; -use super::{AsyncReader, CopySource, ObjectMeta, ObjectStream, Result, Storage, StorageError}; +use super::{ + AsyncReader, CopySource, ObjectMeta, ObjectStream, PutIfAbsentOutcome, Result, Storage, + StorageError, +}; const TOKEN_URL: &str = "https://oauth2.googleapis.com/token"; const STORAGE_HOST: &str = "https://storage.googleapis.com"; @@ -248,6 +251,47 @@ impl Storage for GcsStorage { Ok(()) } + async fn put_if_absent( + &self, + key: &str, + body: AsyncReader, + _size_hint: Option, + ) -> Result { + let token = self.access_token().await?; + let full = self.full_key(key); + // ifGenerationMatch=0 makes the upload a create-if-absent: it succeeds + // only when no live object exists, and returns 412 otherwise. + let url = format!( + "{}/upload/storage/v1/b/{}/o?uploadType=media&ifGenerationMatch=0&name={}", + self.host, + self.cfg.bucket, + utf8_percent_encode(&full, NON_ALPHANUMERIC), + ); + let stream = ReaderStream::new(body); + let resp = self + .client + .post(&url) + .bearer_auth(token) + .header("content-type", "application/octet-stream") + .body(Body::wrap_stream(stream)) + .send() + .await?; + let st = resp.status(); + if st.is_success() { + Ok(PutIfAbsentOutcome::Created) + } else if st == reqwest::StatusCode::PRECONDITION_FAILED + || st == reqwest::StatusCode::CONFLICT + { + Ok(PutIfAbsentOutcome::AlreadyExists) + } else { + let body = resp.text().await.unwrap_or_default(); + Err(StorageError::Http { + status: st.as_u16(), + body: format!("gcs put_if_absent: {body}"), + }) + } + } + async fn get(&self, key: &str) -> Result { let token = self.access_token().await?; let url = format!("{}?alt=media", self.object_url(key)); diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 3c01b52..a69c76e 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -107,6 +107,16 @@ impl From for StorageError { pub type Result = std::result::Result; +/// Outcome of a create-if-absent PUT ([`Storage::put_if_absent`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PutIfAbsentOutcome { + /// No object existed at the key; the body was written. + Created, + /// An object already existed; nothing was written — a benign skip. The + /// existing object is authoritative and is never overwritten. + AlreadyExists, +} + /// Object storage backend /// /// Implementations stream uploads & downloads, no full-segment buffering @@ -118,6 +128,24 @@ pub trait Storage: Send + Sync { /// Upload object. `size_hint` lets s3 backend pick single-PUT vs multipart async fn put(&self, key: &str, body: AsyncReader, size_hint: Option) -> Result<()>; + /// Upload `body` only if no object exists at `key`, via a backend-native + /// atomic create-if-absent write (S3 `If-None-Match: *`, GCS + /// `ifGenerationMatch=0`, fs `O_EXCL`). Returns [`PutIfAbsentOutcome::Created`] + /// on a fresh write and [`PutIfAbsentOutcome::AlreadyExists`] when the object + /// is already present (a benign skip — the existing object is never + /// overwritten). Backends without conditional-write support return + /// `Err(StorageError::Unimplemented(_))` so callers can fall back to a plain + /// [`Storage::put`]. `size_hint`, when known, sizes the buffer. + async fn put_if_absent( + &self, + key: &str, + body: AsyncReader, + size_hint: Option, + ) -> Result { + let _ = (key, body, size_hint); + Err(StorageError::Unimplemented("put_if_absent")) + } + /// Download object as streaming reader async fn get(&self, key: &str) -> Result; diff --git a/src/storage/retrying.rs b/src/storage/retrying.rs index c19cfd7..fc14287 100644 --- a/src/storage/retrying.rs +++ b/src/storage/retrying.rs @@ -10,7 +10,9 @@ use std::io::Cursor; use crate::retry::{RetryPolicy, with_retry}; -use super::{AsyncReader, CopySource, ObjectStream, Result, Storage, StorageError}; +use super::{ + AsyncReader, CopySource, ObjectStream, PutIfAbsentOutcome, Result, Storage, StorageError, +}; /// Buffer-then-retry threshold for put bodies (matches wal-g's small-object /// path: sentinels, history files, manifest fragments) @@ -54,6 +56,31 @@ impl Storage for RetryingStorage { .await } + async fn put_if_absent( + &self, + key: &str, + mut body: AsyncReader, + size_hint: Option, + ) -> Result { + // Only retry small known-size bodies; + let bufferable = matches!(size_hint, Some(s) if s <= PUT_RETRY_BUFFER_THRESHOLD); + if !bufferable { + return self.inner.put_if_absent(key, body, size_hint).await; + } + let mut buf = Vec::with_capacity(size_hint.unwrap_or(0) as usize); + tokio::io::copy(&mut body, &mut buf).await?; + let bytes = Bytes::from(buf); + let len = bytes.len() as u64; + with_retry(&self.policy, StorageError::is_transient, || { + let bytes = bytes.clone(); + async move { + let reader: AsyncReader = Box::pin(Cursor::new(bytes)); + self.inner.put_if_absent(key, reader, Some(len)).await + } + }) + .await + } + async fn get(&self, key: &str) -> Result { with_retry(&self.policy, StorageError::is_transient, || async { self.inner.get(key).await @@ -113,6 +140,8 @@ mod tests { put_bodies: Mutex>>, get_script: Mutex>, put_script: Mutex>, + put_if_absent_calls: AtomicU32, + put_if_absent_script: Mutex>, } enum StubResult { @@ -121,6 +150,12 @@ mod tests { Ok, } + enum StubIfAbsent { + Transient, + Created, + AlreadyExists, + } + impl StubStorage { fn new() -> Self { Self { @@ -129,6 +164,8 @@ mod tests { put_bodies: Mutex::new(Vec::new()), get_script: Mutex::new(Vec::new()), put_script: Mutex::new(Vec::new()), + put_if_absent_calls: AtomicU32::new(0), + put_if_absent_script: Mutex::new(Vec::new()), } } } @@ -181,6 +218,25 @@ mod tests { async fn delete(&self, _key: &str) -> Result<()> { Ok(()) } + async fn put_if_absent( + &self, + _key: &str, + mut body: AsyncReader, + _size_hint: Option, + ) -> Result { + self.put_if_absent_calls.fetch_add(1, Ordering::SeqCst); + let mut buf = Vec::new(); + body.read_to_end(&mut buf).await?; + self.put_bodies.lock().unwrap().push(buf); + match self.put_if_absent_script.lock().unwrap().remove(0) { + StubIfAbsent::Transient => Err(StorageError::Http { + status: 503, + body: "stub down".into(), + }), + StubIfAbsent::Created => Ok(PutIfAbsentOutcome::Created), + StubIfAbsent::AlreadyExists => Ok(PutIfAbsentOutcome::AlreadyExists), + } + } } fn fast_policy() -> RetryPolicy { @@ -275,4 +331,53 @@ mod tests { assert!(matches!(r, Err(StorageError::Http { status: 503, .. }))); assert_eq!(retry.inner.put_calls.load(Ordering::SeqCst), 1); } + + #[tokio::test] + async fn put_if_absent_retries_transient_then_returns_already_exists() { + let stub = StubStorage::new(); + stub.put_if_absent_script.lock().unwrap().extend([ + StubIfAbsent::Transient, + StubIfAbsent::Transient, + StubIfAbsent::AlreadyExists, + ]); + let retry = RetryingStorage::new(stub, fast_policy()); + let body: AsyncReader = Box::pin(Cursor::new(b"seg".to_vec())); + let outcome = retry.put_if_absent("k", body, Some(3)).await.unwrap(); + // AlreadyExists on a retry counts as success (object exists), not an error. + assert!(matches!(outcome, PutIfAbsentOutcome::AlreadyExists)); + assert_eq!(retry.inner.put_if_absent_calls.load(Ordering::SeqCst), 3); + // body replayed byte-identically across attempts + let bodies = retry.inner.put_bodies.lock().unwrap(); + assert_eq!(bodies.len(), 3); + assert!(bodies.iter().all(|b| b == b"seg")); + } + + #[tokio::test] + async fn put_if_absent_retries_transient_then_created() { + let stub = StubStorage::new(); + stub.put_if_absent_script + .lock() + .unwrap() + .extend([StubIfAbsent::Transient, StubIfAbsent::Created]); + let retry = RetryingStorage::new(stub, fast_policy()); + let body: AsyncReader = Box::pin(Cursor::new(b"seg".to_vec())); + let outcome = retry.put_if_absent("k", body, Some(3)).await.unwrap(); + assert!(matches!(outcome, PutIfAbsentOutcome::Created)); + assert_eq!(retry.inner.put_if_absent_calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn put_if_absent_bypasses_retry_when_size_unknown() { + let stub = StubStorage::new(); + stub.put_if_absent_script + .lock() + .unwrap() + .push(StubIfAbsent::Transient); + let retry = RetryingStorage::new(stub, fast_policy()); + // size_hint = None → not bufferable → single pass-through attempt (mirrors put). + let body: AsyncReader = Box::pin(Cursor::new(b"seg".to_vec())); + let r = retry.put_if_absent("k", body, None).await; + assert!(matches!(r, Err(StorageError::Http { status: 503, .. }))); + assert_eq!(retry.inner.put_if_absent_calls.load(Ordering::SeqCst), 1); + } } diff --git a/src/storage/s3.rs b/src/storage/s3.rs index f6aaa48..ea129ff 100644 --- a/src/storage/s3.rs +++ b/src/storage/s3.rs @@ -27,7 +27,10 @@ use tokio_util::io::StreamReader; use url::Url; pub use super::creds::{CredentialSource, Credentials, ImdsProvider}; -use super::{AsyncReader, CopySource, ObjectMeta, ObjectStream, Result, Storage, StorageError}; +use super::{ + AsyncReader, CopySource, ObjectMeta, ObjectStream, PutIfAbsentOutcome, Result, Storage, + StorageError, +}; use crate::retry::{RetryPolicy, with_retry}; const MULTIPART_THRESHOLD: u64 = 32 * 1024 * 1024; @@ -214,6 +217,27 @@ impl S3Storage { .await } + /// One conditional (create-if-absent) PUT of an already-buffered body. + async fn put_if_absent_once(&self, key: &str, body: Bytes) -> Result { + let resp = self + .signed_request( + "PUT", + &self.full_key(key), + &[], + body, + &[("if-none-match", "*")], + ) + .await?; + if matches!( + resp.status(), + reqwest::StatusCode::PRECONDITION_FAILED | reqwest::StatusCode::CONFLICT + ) { + return Ok(PutIfAbsentOutcome::AlreadyExists); + } + check_status(resp).await?; + Ok(PutIfAbsentOutcome::Created) + } + /// PUT one already-buffered part, retrying transients in place (the buffer /// is owned so the body replays without re-reading source). Returns the /// part's ETag for the completion manifest @@ -427,6 +451,33 @@ impl Storage for S3Storage { } } + async fn put_if_absent( + &self, + key: &str, + mut body: AsyncReader, + size_hint: Option, + ) -> Result { + // A conditional PUT needs a rewindable body and bypasses the multipart + // manager, so buffer it. WAL segments compress well under the multipart + // threshold; a body over it is rejected rather than silently split (a + // multipart create-if-absent write is not atomic). + let mut buf = Vec::with_capacity(size_hint.unwrap_or(0) as usize); + body.read_to_end(&mut buf).await?; + if buf.len() as u64 > MULTIPART_THRESHOLD { + return Err(StorageError::Config(format!( + "put_if_absent body {} B exceeds single-PUT limit {} B", + buf.len(), + MULTIPART_THRESHOLD + ))); + } + let body = Bytes::from(buf); + with_retry(&self.retry_policy, StorageError::is_transient, || { + let body = body.clone(); + async move { self.put_if_absent_once(key, body).await } + }) + .await + } + async fn get(&self, key: &str) -> Result { let resp = self .signed_request("GET", &self.full_key(key), &[], Bytes::new(), &[]) @@ -1273,6 +1324,142 @@ mod tests { assert!(uploads.lock().unwrap().is_empty(), "abort must clean up"); } + /// Conditional (create-if-absent) PUT roundtrip against the in-process mock. + /// The mock tracks presence and rejects a second write of the same key with + /// 412, asserting the two mappings: fresh key → Created, existing key → + /// AlreadyExists (never an overwrite). Also checks the `If-None-Match: *` + /// header actually reaches the wire. + #[tokio::test] + async fn s3_put_if_absent_maps_created_and_already_exists() { + use crate::storage::test_http::{Req, Resp, reader, serve}; + use std::collections::BTreeSet; + use std::sync::{Arc, Mutex}; + + let present: Arc>> = Arc::new(Mutex::new(BTreeSet::new())); + let saw_header = Arc::new(Mutex::new(false)); + let (p, sh) = (present.clone(), saw_header.clone()); + let base = serve(move |req: &Req| { + let rest = req.path.trim_start_matches('/'); + let key = rest + .split_once('/') + .map(|(_, k)| k) + .unwrap_or("") + .to_string(); + match req.method.as_str() { + "PUT" => { + if req.headers.get("if-none-match").map(String::as_str) == Some("*") { + *sh.lock().unwrap() = true; + } + let mut set = p.lock().unwrap(); + if set.contains(&key) { + // create-if-absent rejected: object already present + Resp::new(412) + .body(b"PreconditionFailed".to_vec()) + } else { + set.insert(key); + Resp::new(200) + } + } + _ => Resp::new(400), + } + }) + .await; + + let cfg = S3Config { + bucket: "bkt".into(), + prefix: "p".into(), + region: "us-east-1".into(), + creds: CredentialSource::Static(Credentials { + access_key: "AKID".into(), + secret_key: "sek".into(), + session_token: None, + expires_at: None, + }), + endpoint: Some(base), + force_path_style: true, + }; + let policy = RetryPolicy { + max_attempts: 2, + base_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + jitter: false, + }; + let s = S3Storage::with_retry_policy(cfg, policy).unwrap(); + + // absent key → written, reported Created + assert_eq!( + s.put_if_absent("wal_005/x.lz4", reader(b"hello"), Some(5)) + .await + .unwrap(), + PutIfAbsentOutcome::Created + ); + assert!(*saw_header.lock().unwrap(), "If-None-Match: * must be sent"); + // same key again → server 412, reported AlreadyExists (no overwrite) + assert_eq!( + s.put_if_absent("wal_005/x.lz4", reader(b"world"), Some(5)) + .await + .unwrap(), + PutIfAbsentOutcome::AlreadyExists + ); + } + + #[tokio::test] + async fn s3_put_if_absent_retries_transient_then_created() { + use crate::storage::test_http::{Req, Resp, reader, serve}; + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + + // First PUT → transient 503; the retry (in s3's put_if_absent: buffer-once + // + with_retry over put_if_absent_once) → 200 Created. + let attempts = Arc::new(AtomicU32::new(0)); + let a = attempts.clone(); + let base = serve(move |req: &Req| { + if req.method == "PUT" { + if a.fetch_add(1, Ordering::SeqCst) == 0 { + Resp::new(503).body(b"SlowDown".to_vec()) + } else { + Resp::new(200) + } + } else { + Resp::new(400) + } + }) + .await; + + let cfg = S3Config { + bucket: "bkt".into(), + prefix: "p".into(), + region: "us-east-1".into(), + creds: CredentialSource::Static(Credentials { + access_key: "AKID".into(), + secret_key: "sek".into(), + session_token: None, + expires_at: None, + }), + endpoint: Some(base), + force_path_style: true, + }; + let policy = RetryPolicy { + max_attempts: 3, + base_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + jitter: false, + }; + let s = S3Storage::with_retry_policy(cfg, policy).unwrap(); + + assert_eq!( + s.put_if_absent("wal_005/x.lz4", reader(b"hello"), Some(5)) + .await + .unwrap(), + PutIfAbsentOutcome::Created + ); + assert_eq!( + attempts.load(Ordering::SeqCst), + 2, + "one transient retry then success" + ); + } + /// Pipelined multipart keeps several part PUTs in flight, so they finish /// out of order; CompleteMultipartUpload must still list every part /// ascending by partNumber with its matching ETag (S3 rejects unsorted