feat: ttl support for atomic operations

Author: evansims

crates/storage-ledger/src/backend.rs

@@ -395,7 +395,7 @@ impl LedgerBackend {
     ///
     /// Returns [`StorageError::Internal`] if the system clock is set before
     /// the Unix epoch.
-    fn compute_expiration_timestamp(ttl: Duration) -> StorageResult<u64> {
+    pub(crate) fn compute_expiration_timestamp(ttl: Duration) -> StorageResult<u64> {
         std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .map(|d| d.as_secs() + ttl.as_secs())
@@ -809,6 +809,70 @@ impl StorageBackend for LedgerBackend {
         result
     }
 
+    /// Performs a compare-and-set operation with TTL, retrying on transient errors.
+    ///
+    /// Combines CAS precondition checking with automatic key expiration.
+    /// Conflict errors (`FailedPrecondition`) are **not** retried.
+    #[tracing::instrument(skip(self, key, expected, new_value), fields(key_len = key.len(), ttl_ms = ttl.as_millis() as u64))]
+    async fn compare_and_set_with_ttl(
+        &self,
+        key: &[u8],
+        expected: Option<&[u8]>,
+        new_value: Vec<u8>,
+        ttl: Duration,
+    ) -> StorageResult<()> {
+        self.check_cancelled()?;
+        self.check_circuit()?;
+        self.check_sizes(key, &new_value)?;
+        let start = std::time::Instant::now();
+        let encoded_key = encode_key(key);
+        let expires_at = Self::compute_expiration_timestamp(ttl)?;
+
+        let condition = match expected {
+            None => SetCondition::NotExists,
+            Some(expected_value) => SetCondition::ValueEquals(expected_value.to_vec()),
+        };
+
+        use inferadb_ledger_sdk::SdkError;
+        use tonic::Code;
+
+        let result = with_retry_timeout(
+            &self.retry_config,
+            self.timeout_config.write_timeout,
+            None,
+            "compare_and_set_with_ttl",
+            || async {
+                match self
+                    .client
+                    .write(
+                        self.organization,
+                        self.vault,
+                        vec![Operation::SetEntity {
+                            key: encoded_key.clone(),
+                            value: new_value.clone(),
+                            expires_at: Some(expires_at),
+                            condition: Some(condition.clone()),
+                        }],
+                    )
+                    .await
+                {
+                    Ok(_) => Ok(()),
+                    Err(SdkError::Rpc { code: Code::FailedPrecondition, .. }) => {
+                        Err(StorageError::conflict())
+                    },
+                    Err(e) => Err(StorageError::from(LedgerStorageError::from(e))),
+                }
+            },
+        )
+        .await;
+        self.record_circuit_result(&result);
+        self.metrics.record_set_org(start.elapsed(), &self.org_str());
+        if result.is_err() {
+            self.metrics.record_error_org(&self.org_str());
+        }
+        result
+    }
+
     /// Creates a new [`LedgerTransaction`] for buffered atomic writes.
     #[tracing::instrument(skip(self))]
     async fn transaction(&self) -> StorageResult<Box<dyn Transaction>> {

crates/storage-ledger/src/transaction.rs

@@ -7,6 +7,7 @@
 use std::{
     collections::{HashMap, HashSet},
     sync::Arc,
+    time::Duration,
 };
 
 use async_trait::async_trait;
@@ -27,6 +28,8 @@ struct CasOperation {
     expected: Option<Vec<u8>>,
     /// New value to set if condition is met.
     new_value: Vec<u8>,
+    /// Optional TTL for the new value.
+    ttl: Option<Duration>,
 }
 
 /// Transaction for atomic operations on the Ledger backend.
@@ -120,6 +123,9 @@ pub struct LedgerTransaction {
     /// Pending delete operations: hex-encoded keys.
     pending_deletes: HashSet<String>,
 
+    /// Pending TTLs for set operations: hex-encoded key -> duration.
+    pending_ttls: HashMap<String, Duration>,
+
     /// Pending compare-and-set operations.
     pending_cas: Vec<CasOperation>,
 }
@@ -151,6 +157,7 @@ impl LedgerTransaction {
             read_consistency,
             pending_sets: HashMap::new(),
             pending_deletes: HashSet::new(),
+            pending_ttls: HashMap::new(),
             pending_cas: Vec::new(),
         }
     }
@@ -199,6 +206,9 @@ impl Transaction for LedgerTransaction {
         // Remove from pending deletes if it was marked for deletion
         self.pending_deletes.remove(&encoded_key);
 
+        // Clear any pending TTL — a plain set produces a non-expiring key
+        self.pending_ttls.remove(&encoded_key);
+
         // Add to pending sets
         self.pending_sets.insert(encoded_key, value);
     }
@@ -207,6 +217,9 @@ impl Transaction for LedgerTransaction {
     fn delete(&mut self, key: Vec<u8>) {
         let encoded_key = encode_key(&key);
 
+        // Clear any pending TTL
+        self.pending_ttls.remove(&encoded_key);
+
         // Remove from pending sets if it was set in this transaction
         self.pending_sets.remove(&encoded_key);
 
@@ -224,7 +237,40 @@ impl Transaction for LedgerTransaction {
         let encoded_key = encode_key(&key);
 
         // Buffer the CAS operation - it will be applied at commit time
-        self.pending_cas.push(CasOperation { key: encoded_key, expected, new_value });
+        self.pending_cas.push(CasOperation { key: encoded_key, expected, new_value, ttl: None });
+        Ok(())
+    }
+
+    /// Buffers a set operation with TTL for atomic commit.
+    fn set_with_ttl(&mut self, key: Vec<u8>, value: Vec<u8>, ttl: Duration) {
+        let encoded_key = encode_key(&key);
+
+        // Remove from pending deletes if it was marked for deletion
+        self.pending_deletes.remove(&encoded_key);
+
+        // Track the TTL for this key
+        self.pending_ttls.insert(encoded_key.clone(), ttl);
+
+        // Add to pending sets
+        self.pending_sets.insert(encoded_key, value);
+    }
+
+    /// Buffers a compare-and-set operation with TTL for atomic commit.
+    fn compare_and_set_with_ttl(
+        &mut self,
+        key: Vec<u8>,
+        expected: Option<Vec<u8>>,
+        new_value: Vec<u8>,
+        ttl: Duration,
+    ) -> StorageResult<()> {
+        let encoded_key = encode_key(&key);
+
+        self.pending_cas.push(CasOperation {
+            key: encoded_key,
+            expected,
+            new_value,
+            ttl: Some(ttl),
+        });
         Ok(())
     }
 
@@ -253,12 +299,24 @@ impl Transaction for LedgerTransaction {
                 None => SetCondition::NotExists,
                 Some(expected_value) => SetCondition::ValueEquals(expected_value),
             };
-            operations.push(Operation::set_entity_if(cas.key, cas.new_value, condition));
+            let expires_at =
+                cas.ttl.map(crate::LedgerBackend::compute_expiration_timestamp).transpose()?;
+            operations.push(Operation::SetEntity {
+                key: cas.key,
+                value: cas.new_value,
+                expires_at,
+                condition: Some(condition),
+            });
         }
 
         // Add regular set operations
         for (key, value) in self.pending_sets {
-            operations.push(Operation::set_entity(key, value));
+            if let Some(ttl) = self.pending_ttls.get(&key) {
+                let expires_at = crate::LedgerBackend::compute_expiration_timestamp(*ttl)?;
+                operations.push(Operation::set_entity_with_expiry(key, value, expires_at));
+            } else {
+                operations.push(Operation::set_entity(key, value));
+            }
         }
 
         // Add delete operations

crates/storage/src/backend.rs

@@ -57,6 +57,7 @@ use crate::{
 /// | [`get_range`](StorageBackend::get_range) | Retrieve multiple keys in a range |
 /// | [`clear_range`](StorageBackend::clear_range) | Delete multiple keys in a range |
 /// | [`set_with_ttl`](StorageBackend::set_with_ttl) | Store with automatic expiration |
+/// | [`compare_and_set_with_ttl`](StorageBackend::compare_and_set_with_ttl) | Atomic compare-and-set with expiration |
 /// | [`transaction`](StorageBackend::transaction) | Begin an atomic transaction |
 /// | [`health_check`](StorageBackend::health_check) | Verify backend availability |
 ///
@@ -407,6 +408,42 @@ pub trait StorageBackend: Send + Sync {
     #[must_use = "storage operations may fail and errors must be handled"]
     async fn set_with_ttl(&self, key: Vec<u8>, value: Vec<u8>, ttl: Duration) -> StorageResult<()>;
 
+    /// Atomically sets a key's value if it matches the expected current value,
+    /// with automatic expiration.
+    ///
+    /// Combines the semantics of [`compare_and_set`](StorageBackend::compare_and_set)
+    /// with [`set_with_ttl`](StorageBackend::set_with_ttl): the CAS precondition
+    /// is evaluated first, and on success the new value is stored with the
+    /// given TTL.
+    ///
+    /// # Semantics
+    ///
+    /// See [`compare_and_set`](StorageBackend::compare_and_set) for the full
+    /// CAS precondition rules. The only difference is that on success the key
+    /// is given a TTL instead of becoming non-expiring.
+    ///
+    /// # Arguments
+    ///
+    /// * `key` — The key to update.
+    /// * `expected` — The expected current value. Use `None` for insert-if-absent.
+    /// * `new_value` — The new value to set if the comparison succeeds.
+    /// * `ttl` — Time-to-live duration after which the key expires.
+    ///
+    /// # Errors
+    ///
+    /// - [`StorageError::Conflict`](crate::StorageError) — the current value does not match
+    ///   `expected`.
+    /// - [`StorageError::SizeLimitExceeded`] — `key` or `new_value` exceeds the configured size
+    ///   limits.
+    #[must_use = "compare-and-set may fail with a conflict and errors must be handled"]
+    async fn compare_and_set_with_ttl(
+        &self,
+        key: &[u8],
+        expected: Option<&[u8]>,
+        new_value: Vec<u8>,
+        ttl: Duration,
+    ) -> StorageResult<()>;
+
     /// Begins a new transaction.
     ///
     /// Returns a [`Transaction`] handle that can be used to perform

crates/storage/src/batch.rs

@@ -1057,6 +1057,23 @@ mod tests {
             self.inner.get_mut().expect("lock poisoned").compare_and_set(key, expected, new_value)
         }
 
+        fn set_with_ttl(&mut self, key: Vec<u8>, value: Vec<u8>, ttl: std::time::Duration) {
+            self.inner.get_mut().expect("lock poisoned").set_with_ttl(key, value, ttl);
+        }
+
+        fn compare_and_set_with_ttl(
+            &mut self,
+            key: Vec<u8>,
+            expected: Option<Vec<u8>>,
+            new_value: Vec<u8>,
+            ttl: std::time::Duration,
+        ) -> StorageResult<()> {
+            self.inner
+                .get_mut()
+                .expect("lock poisoned")
+                .compare_and_set_with_ttl(key, expected, new_value, ttl)
+        }
+
         async fn commit(self: Box<Self>) -> StorageResult<()> {
             if self.should_fail {
                 Err(StorageError::connection("simulated commit failure"))
@@ -1085,6 +1102,16 @@ mod tests {
             self.inner.compare_and_set(key, expected, new_value).await
         }
 
+        async fn compare_and_set_with_ttl(
+            &self,
+            key: &[u8],
+            expected: Option<&[u8]>,
+            new_value: Vec<u8>,
+            ttl: std::time::Duration,
+        ) -> StorageResult<()> {
+            self.inner.compare_and_set_with_ttl(key, expected, new_value, ttl).await
+        }
+
         async fn delete(&self, key: &[u8]) -> StorageResult<()> {
             self.inner.delete(key).await
         }

crates/storage/src/conformance.rs

@@ -308,6 +308,72 @@ pub async fn tx_cas_conflict_rejects_commit<B: StorageBackend>(backend: &B) {
     assert_eq!(val, Some(Bytes::from("v_concurrent")));
 }
 
+// ============================================================================
+// Transaction TTL — set_with_ttl and compare_and_set_with_ttl in transactions (3 tests)
+// ============================================================================
+
+/// `set_with_ttl` in a transaction applies the TTL on commit, and the key
+/// expires after the TTL elapses.
+pub async fn tx_set_with_ttl_applies_on_commit<B: StorageBackend>(backend: &B) {
+    let mut tx = backend.transaction().await.expect("transaction");
+    tx.set_with_ttl(
+        b"txtl:expire".to_vec(),
+        b"ephemeral".to_vec(),
+        std::time::Duration::from_millis(50),
+    );
+    tx.commit().await.expect("commit");
+
+    // Immediately after commit the key should be readable
+    let val = backend.get(b"txtl:expire").await.expect("get");
+    assert_eq!(val, Some(Bytes::from("ephemeral")));
+
+    // After sleeping past the TTL the key should be gone
+    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    let val = backend.get(b"txtl:expire").await.expect("get after ttl");
+    assert_eq!(val, None);
+}
+
+/// `set_with_ttl` in a transaction supports read-your-writes: a `get` in the
+/// same transaction returns the buffered value.
+pub async fn tx_set_with_ttl_read_your_writes<B: StorageBackend>(backend: &B) {
+    let mut tx = backend.transaction().await.expect("transaction");
+    tx.set_with_ttl(b"txtl:ryw".to_vec(), b"buffered".to_vec(), std::time::Duration::from_secs(60));
+
+    let val = tx.get(b"txtl:ryw").await.expect("get");
+    assert_eq!(val, Some(Bytes::from("buffered")));
+
+    // Drop without commit — the key should not be visible
+    drop(tx);
+    let val = backend.get(b"txtl:ryw").await.expect("get after drop");
+    assert_eq!(val, None);
+}
+
+/// `compare_and_set_with_ttl` in a transaction applies both the CAS
+/// precondition and the TTL on commit.
+pub async fn tx_compare_and_set_with_ttl<B: StorageBackend>(backend: &B) {
+    // Insert a key so we can CAS against it
+    backend.set(b"txtl:cas".to_vec(), b"v1".to_vec()).await.expect("set");
+
+    let mut tx = backend.transaction().await.expect("transaction");
+    tx.compare_and_set_with_ttl(
+        b"txtl:cas".to_vec(),
+        Some(b"v1".to_vec()),
+        b"v2".to_vec(),
+        std::time::Duration::from_millis(50),
+    )
+    .expect("buffer CAS with TTL");
+    tx.commit().await.expect("commit");
+
+    // Immediately after commit the key should hold the new value
+    let val = backend.get(b"txtl:cas").await.expect("get");
+    assert_eq!(val, Some(Bytes::from("v2")));
+
+    // After sleeping past the TTL the key should be gone
+    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    let val = backend.get(b"txtl:cas").await.expect("get after ttl");
+    assert_eq!(val, None);
+}
+
 // ============================================================================
 // CAS — compare_and_set precondition checks (4 tests)
 // ============================================================================
@@ -517,6 +583,11 @@ pub async fn run_all<B: StorageBackend + 'static>(backend: Arc<B>) {
     tx_delete_then_get_returns_none(backend.as_ref()).await;
     tx_cas_conflict_rejects_commit(backend.as_ref()).await;
 
+    // Transaction TTL
+    tx_set_with_ttl_applies_on_commit(backend.as_ref()).await;
+    tx_set_with_ttl_read_your_writes(backend.as_ref()).await;
+    tx_compare_and_set_with_ttl(backend.as_ref()).await;
+
     // CAS
     cas_insert_if_absent(backend.as_ref()).await;
     cas_insert_if_absent_fails_when_key_exists(backend.as_ref()).await;