yeti_types/backend/

traits.rs

//! Storage backend traits: `KvBackend`, `SnapshotTransaction`, and the
//! `WriteOp` batch operation type.

use async_trait::async_trait;

use crate::error::{Result, YetiError};

// ============================================================================
// WriteOp
// ============================================================================

/// A single write operation in a batch.
///
/// Serializable so it can be embedded in [`crate::backend::log::LogEntry`]
/// and replicated over the wire alongside HLC + originator metadata.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum WriteOp {
    /// Insert or update a key-value pair
    Put {
        /// Key bytes
        key: Vec<u8>,
        /// Value bytes
        value: Vec<u8>,
    },
    /// Delete a key
    Delete {
        /// Key bytes to delete
        key: Vec<u8>,
    },
}

impl WriteOp {
    /// The key this op targets (regardless of variant).
    #[must_use]
    pub fn key(&self) -> &[u8] {
        match self {
            Self::Put { key, .. } | Self::Delete { key } => key,
        }
    }

    /// True if this is a [`WriteOp::Delete`].
    #[must_use]
    pub const fn is_delete(&self) -> bool {
        matches!(self, Self::Delete { .. })
    }
}

// ============================================================================
// KvBackend trait
// ============================================================================

/// Drive a backend future to completion from a synchronous context, working on
/// both runtime flavors. On a multi-threaded runtime we use `block_in_place` so
/// the worker parks-and-migrates (other tasks keep running); on a current-thread
/// runtime (tests, the wasm benches) `block_in_place` would panic, so we fall
/// back to the `futures` executor.
///
/// This is only the *fallback* path for the `*_sync` trait methods — hot
/// backends (`RocksDB`, in-memory) override those with genuinely synchronous
/// implementations and never reach this helper. It must therefore never be
/// applied to a future that can only make progress on the very runtime we're
/// blocking (e.g. a full bounded-channel send); the only in-tree user is the
/// read passthrough, which is channel-free.
fn block_on_backend<F: std::future::Future>(fut: F) -> F::Output {
    // Native: prefer block_in_place on a multi-threaded runtime (parks +
    // migrates the worker); fall back to the futures executor on current-thread
    // (where block_in_place panics).
    #[cfg(not(target_arch = "wasm32"))]
    {
        match tokio::runtime::Handle::try_current() {
            Ok(h)
                if matches!(
                    h.runtime_flavor(),
                    tokio::runtime::RuntimeFlavor::MultiThread
                ) =>
            {
                tokio::task::block_in_place(|| h.block_on(fut))
            },
            _ => futures::executor::block_on(fut),
        }
    }
    // wasm32 guest: no tokio threading (block_in_place doesn't exist), and the
    // guest never drives a host backend through this path anyway.
    #[cfg(target_arch = "wasm32")]
    {
        futures::executor::block_on(fut)
    }
}

/// Trait that all storage backends must implement.
#[async_trait]
pub trait KvBackend: Send + Sync {
    /// Put a key-value pair.
    async fn put(&self, key: &[u8], value: &[u8]) -> Result<()>;

    /// Get a value by key.
    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>>;

    /// Synchronous get — backends whose reads complete inline (e.g. `RocksDB`
    /// via `with_cf_inline`, in-memory `HashMap`) should override this to skip
    /// the `async_trait` `Box::pin` overhead. The default falls back to the
    /// async variant via `block_on_backend` so any backend is correct without
    /// an override; hot backends should override for zero-cost dispatch.
    ///
    /// # Errors
    /// Propagates any storage error from the underlying backend.
    fn get_sync(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        block_on_backend(self.get(key))
    }

    /// Synchronous put — same rationale as `get_sync`.
    ///
    /// # Errors
    /// Propagates any storage error from the underlying backend.
    fn put_sync(&self, key: &[u8], value: &[u8]) -> Result<()> {
        block_on_backend(self.put(key, value))
    }

    /// Synchronous delete — same rationale as `get_sync`.
    ///
    /// # Errors
    /// Propagates any storage error from the underlying backend.
    fn delete_sync(&self, key: &[u8]) -> Result<()> {
        block_on_backend(self.delete(key))
    }

    /// Check if a key exists without loading the full value.
    async fn exists(&self, key: &[u8]) -> Result<bool> {
        Ok(self.get(key).await?.is_some())
    }

    /// Get multiple values by keys (batch get).
    async fn get_batch(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>> {
        let mut results = Vec::with_capacity(keys.len());
        for key in keys {
            results.push(self.get(key).await?);
        }
        Ok(results)
    }

    /// Delete a key.
    async fn delete(&self, key: &[u8]) -> Result<()>;

    /// Scan keys with a given prefix, returning all matching key-value pairs.
    async fn scan_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>>;

    /// Scan keys with a given prefix, returning only keys.
    async fn scan_keys(&self, prefix: &[u8]) -> Result<Vec<Vec<u8>>> {
        let results = self.scan_prefix(prefix).await?;
        Ok(results.into_iter().map(|(k, _)| k).collect())
    }

    /// Count keys with a given prefix.
    async fn count_prefix(&self, prefix: &[u8]) -> Result<usize>;

    /// Batch write operations (atomic).
    async fn write_batch(&self, ops: Vec<WriteOp>) -> Result<()>;

    /// Synchronous `write_batch` — same rationale as `get_sync`. Backends
    /// whose small-batch commit is inline (`RocksDB`, in-memory) should override
    /// to skip the `async_trait` `Box::pin`. The default falls back to the
    /// async variant via `block_in_place`.
    ///
    /// # Errors
    /// Propagates any storage error from the underlying backend.
    fn write_batch_sync(&self, ops: Vec<WriteOp>) -> Result<()> {
        block_on_backend(self.write_batch(ops))
    }

    /// Flush any pending writes to disk.
    async fn flush(&self) -> Result<()>;

    /// Scan keys with a given prefix in reverse order.
    async fn scan_prefix_reverse(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let mut results = self.scan_prefix(prefix).await?;
        results.reverse();
        Ok(results)
    }

    /// Scan a key range [start, end) returning all key-value pairs.
    async fn scan_range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let all = self.scan_prefix(&[]).await?;
        Ok(all
            .into_iter()
            .filter(|(k, _)| k.as_slice() >= start && k.as_slice() < end)
            .collect())
    }

    /// Get storage statistics as a JSON object.
    fn storage_stats(&self) -> Option<serde_json::Value> {
        None
    }

    /// Begin a snapshot-isolated transaction.
    ///
    /// Returns a transaction handle with read-your-writes consistency and
    /// optimistic conflict detection on commit. Not all backends support this.
    async fn begin_transaction(&self) -> Result<Box<dyn SnapshotTransaction>> {
        Err(YetiError::Internal(
            "Snapshot transactions not supported by this backend".into(),
        ))
    }

    /// Submit a raw JSON record for eventual processing (WAL consumer path).
    ///
    /// Strong consistency backends return an error (caller falls back to sync write).
    async fn submit_record(&self, _primary_key: &str, _json: &[u8], _delete: bool) -> Result<()> {
        Err(YetiError::Internal(
            "submit_record not supported by this backend".into(),
        ))
    }

    /// Delete all keys in this backend (truncate table).
    async fn truncate(&self) -> Result<()> {
        let keys = self.scan_keys(&[]).await?;
        if keys.is_empty() {
            return Ok(());
        }
        let ops: Vec<WriteOp> = keys
            .into_iter()
            .map(|k| WriteOp::Delete { key: k })
            .collect();
        self.write_batch(ops).await
    }

    /// Compare-and-swap put (durable work queue).
    ///
    /// Writes `new` for `key` only if the current on-disk bytes match
    /// `expected` (`None` meaning "the key must not exist"). Returns
    /// `YetiError::Cas { reason: CasReason::Mismatch }` when the expected
    /// value does not match.
    ///
    /// Default implementation is a non-atomic read+write — usable but racy
    /// under concurrent load. Real backends override this with a `RocksDB`
    /// `OptimisticTransactionDB` round-trip (see `yeti-store`).
    async fn put_if(&self, key: &[u8], expected: Option<&[u8]>, new: &[u8]) -> Result<()> {
        let current = self.get(key).await?;
        let matches = match (expected, current.as_deref()) {
            (None, None) => true,
            (Some(e), Some(c)) => e == c,
            _ => false,
        };
        if !matches {
            return Err(YetiError::Cas {
                reason: crate::error::CasReason::Mismatch,
            });
        }
        self.put(key, new).await
    }
}

// ============================================================================
// SnapshotTransaction
// ============================================================================

/// A snapshot-isolated transaction handle.
///
/// Provides read-your-writes consistency within the transaction scope.
/// On commit, optimistic conflict detection ensures serializability.
#[async_trait]
pub trait SnapshotTransaction: Send + Sync {
    /// Read a value within the transaction's snapshot.
    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>>;
    /// Write a value (visible to subsequent reads within this transaction).
    async fn put(&self, key: &[u8], value: &[u8]) -> Result<()>;
    /// Delete a key within the transaction.
    async fn delete(&self, key: &[u8]) -> Result<()>;
    /// Commit the transaction (optimistic conflict detection).
    async fn commit(self: Box<Self>) -> Result<()>;
    /// Rollback the transaction, discarding all changes.
    async fn rollback(self: Box<Self>) -> Result<()>;
}