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

Author: achudnovskij

src/pg/replication/tls.rs

@@ -156,7 +156,7 @@ pub async fn maybe_upgrade(
     }
 }
 
-fn build_client_config(sslmode: SslMode, tls: &TlsParams) -> Result<ClientConfig> {
+pub fn build_client_config(sslmode: SslMode, tls: &TlsParams) -> Result<ClientConfig> {
     let provider = rustls::crypto::aws_lc_rs::default_provider();
     let builder = ClientConfig::builder_with_provider(Arc::new(provider))
         .with_safe_default_protocol_versions()

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<u64>,
+    ) -> Result<PutIfAbsentOutcome> {
+        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<AsyncReader> {
         let token = self.access_token().await?;
         let url = format!("{}?alt=media", self.object_url(key));

src/storage/mod.rs

@@ -107,6 +107,16 @@ impl From<reqwest::Error> for StorageError {
 
 pub type Result<T> = std::result::Result<T, StorageError>;
 
+/// 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<u64>) -> 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<u64>,
+    ) -> Result<PutIfAbsentOutcome> {
+        let _ = (key, body, size_hint);
+        Err(StorageError::Unimplemented("put_if_absent"))
+    }
+
     /// Download object as streaming reader
     async fn get(&self, key: &str) -> Result<AsyncReader>;
 

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<S: Storage + 'static> Storage for RetryingStorage<S> {
         .await
     }
 
+    async fn put_if_absent(
+        &self,
+        key: &str,
+        mut body: AsyncReader,
+        size_hint: Option<u64>,
+    ) -> Result<PutIfAbsentOutcome> {
+        // 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<AsyncReader> {
         with_retry(&self.policy, StorageError::is_transient, || async {
             self.inner.get(key).await
@@ -113,6 +140,8 @@ mod tests {
         put_bodies: Mutex<Vec<Vec<u8>>>,
         get_script: Mutex<Vec<StubResult>>,
         put_script: Mutex<Vec<StubResult>>,
+        put_if_absent_calls: AtomicU32,
+        put_if_absent_script: Mutex<Vec<StubIfAbsent>>,
     }
 
     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<u64>,
+        ) -> Result<PutIfAbsentOutcome> {
+            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);
+    }
 }