feoxdb/core/store/

atomic.rs

use std::sync::atomic::Ordering;
use std::sync::Arc;

use crate::constants::Operation;
use crate::core::record::Record;
use crate::error::{FeoxError, Result};

use super::FeoxStore;

impl FeoxStore {
    /// Atomically increment a numeric counter.
    ///
    /// The value must be stored as an 8-byte little-endian i64. If the key doesn't exist,
    /// it will be created with the given delta value. If it exists, the value will be
    /// incremented atomically.
    ///
    /// # Value Format
    ///
    /// The value MUST be exactly 8 bytes representing a little-endian i64.
    /// Use `i64::to_le_bytes()` to create the initial value:
    /// ```rust,ignore
    /// let zero: i64 = 0;
    /// store.insert(b"counter", &zero.to_le_bytes())?;
    /// ```
    ///
    /// # Arguments
    ///
    /// * `key` - The key of the counter
    /// * `delta` - The amount to increment by (can be negative for decrement)
    /// * `timestamp` - Optional timestamp for conflict resolution
    ///
    /// # Returns
    ///
    /// Returns the new value after incrementing.
    ///
    /// # Errors
    ///
    /// * `InvalidOperation` - Existing value is not exactly 8 bytes (not a valid i64)
    /// * `OlderTimestamp` - Timestamp is not newer than existing record
    ///
    /// # Example
    ///
    /// ```rust
    /// # use feoxdb::FeoxStore;
    /// # fn main() -> feoxdb::Result<()> {
    /// # let store = FeoxStore::new(None)?;
    /// // Initialize counter with proper binary format
    /// let initial: i64 = 0;
    /// store.insert(b"visits", &initial.to_le_bytes())?;
    ///
    /// // Increment atomically
    /// let val = store.atomic_increment(b"visits", 1)?;
    /// assert_eq!(val, 1);
    ///
    /// // Increment by 5
    /// let val = store.atomic_increment(b"visits", 5)?;
    /// assert_eq!(val, 6);
    ///
    /// // Decrement by 2
    /// let val = store.atomic_increment(b"visits", -2)?;
    /// assert_eq!(val, 4);
    ///
    /// // Or create new counter directly (starts at delta value)
    /// let downloads = store.atomic_increment(b"downloads", 100)?;
    /// assert_eq!(downloads, 100);
    /// # Ok(())
    /// # }
    /// ```
    pub fn atomic_increment(&self, key: &[u8], delta: i64) -> Result<i64> {
        self.atomic_increment_with_timestamp_and_ttl(key, delta, None, 0)
    }

    /// Atomically increment/decrement with explicit timestamp.
    ///
    /// This is the advanced version that allows manual timestamp control.
    /// Most users should use `atomic_increment()` instead.
    ///
    /// # Arguments
    ///
    /// * `key` - The key to increment/decrement
    /// * `delta` - Amount to add (negative to decrement)
    /// * `timestamp` - Optional timestamp. If `None`, uses current time.
    ///
    /// # Errors
    ///
    /// * `OlderTimestamp` - Timestamp is not newer than existing record
    pub fn atomic_increment_with_timestamp(
        &self,
        key: &[u8],
        delta: i64,
        timestamp: Option<u64>,
    ) -> Result<i64> {
        self.atomic_increment_with_timestamp_and_ttl(key, delta, timestamp, 0)
    }

    /// Atomically increment/decrement with TTL support.
    ///
    /// # Arguments
    ///
    /// * `key` - The key to increment/decrement
    /// * `delta` - Amount to add (negative to decrement)
    /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
    ///
    /// # Errors
    ///
    /// * `InvalidOperation` - Value is not a valid i64
    pub fn atomic_increment_with_ttl(
        &self,
        key: &[u8],
        delta: i64,
        ttl_seconds: u64,
    ) -> Result<i64> {
        self.atomic_increment_with_timestamp_and_ttl(key, delta, None, ttl_seconds)
    }

    /// Atomically increment/decrement with explicit timestamp and TTL.
    ///
    /// # Arguments
    ///
    /// * `key` - The key to increment/decrement
    /// * `delta` - Amount to add (negative to decrement)
    /// * `timestamp` - Optional timestamp. If `None`, uses current time.
    /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
    ///
    /// # Errors
    ///
    /// * `OlderTimestamp` - Timestamp is not newer than existing record
    pub fn atomic_increment_with_timestamp_and_ttl(
        &self,
        key: &[u8],
        delta: i64,
        timestamp: Option<u64>,
        ttl_seconds: u64,
    ) -> Result<i64> {
        self.validate_key(key)?;

        let key_vec = key.to_vec();

        let result = match self.hash_table.entry(key_vec.clone()) {
            scc::hash_map::Entry::Occupied(mut entry) => {
                let old_record = entry.get();

                // Get timestamp inside the critical section to ensure it's always newer
                let timestamp = match timestamp {
                    Some(0) | None => self.get_timestamp(),
                    Some(ts) => ts,
                };

                // Check if timestamp is valid
                if timestamp < old_record.timestamp {
                    return Err(FeoxError::OlderTimestamp);
                }

                // Load value from memory or disk
                let value = if let Some(val) = old_record.get_value() {
                    val.to_vec()
                } else {
                    // Try loading from disk if not in memory
                    self.load_value_from_disk(old_record)?
                };

                let current_val = if value.len() == 8 {
                    let bytes = value
                        .get(..8)
                        .and_then(|slice| slice.try_into().ok())
                        .ok_or(FeoxError::InvalidNumericValue)?;
                    i64::from_le_bytes(bytes)
                } else {
                    return Err(FeoxError::InvalidOperation);
                };

                let new_val = current_val.saturating_add(delta);
                let new_value = new_val.to_le_bytes().to_vec();

                // Create new record with TTL if specified
                let new_record = if ttl_seconds > 0 {
                    let ttl_expiry = timestamp + (ttl_seconds * 1_000_000_000); // Convert to nanoseconds
                    Arc::new(Record::new_with_timestamp_ttl(
                        old_record.key.clone(),
                        new_value,
                        timestamp,
                        ttl_expiry,
                    ))
                } else {
                    Arc::new(Record::new(old_record.key.clone(), new_value, timestamp))
                };

                let old_value_len = old_record.value_len;
                let old_size = old_record.calculate_size();
                let new_size = self.calculate_record_size(old_record.key.len(), 8);
                let old_record_arc = Arc::clone(old_record);

                // Atomically update the entry
                entry.insert(Arc::clone(&new_record));

                // Update skip list as well
                self.tree.insert(key_vec.clone(), Arc::clone(&new_record));

                // Update memory usage
                if new_size > old_size {
                    self.stats
                        .memory_usage
                        .fetch_add(new_size - old_size, Ordering::AcqRel);
                } else {
                    self.stats
                        .memory_usage
                        .fetch_sub(old_size - new_size, Ordering::AcqRel);
                }

                // Only do cache and persistence operations if not in memory-only mode
                if !self.memory_only {
                    if self.enable_caching {
                        if let Some(ref cache) = self.cache {
                            cache.remove(&key_vec);
                        }
                    }

                    if let Some(ref wb) = self.write_buffer {
                        if let Err(e) =
                            wb.add_write(Operation::Update, Arc::clone(&new_record), old_value_len)
                        {
                            // Atomic operation succeeded in memory
                            let _ = e;
                        }

                        if let Err(e) =
                            wb.add_write(Operation::Delete, old_record_arc, old_value_len)
                        {
                            // Atomic operation succeeded in memory
                            let _ = e;
                        }
                    }
                }

                Ok(new_val)
            }
            scc::hash_map::Entry::Vacant(entry) => {
                // Key doesn't exist, create it with initial value
                // Get timestamp inside the critical section
                let timestamp = match timestamp {
                    Some(0) | None => self.get_timestamp(),
                    Some(ts) => ts,
                };

                let initial_val = delta;
                let value = initial_val.to_le_bytes().to_vec();

                // Create new record with TTL if specified
                let new_record = if ttl_seconds > 0 {
                    let ttl_expiry = timestamp + (ttl_seconds * 1_000_000_000); // Convert to nanoseconds
                    Arc::new(Record::new_with_timestamp_ttl(
                        key_vec.clone(),
                        value,
                        timestamp,
                        ttl_expiry,
                    ))
                } else {
                    Arc::new(Record::new(key_vec.clone(), value, timestamp))
                };

                let _ = entry.insert_entry(Arc::clone(&new_record));

                // Update skip list
                self.tree.insert(key_vec.clone(), Arc::clone(&new_record));

                // Update statistics
                self.stats.record_count.fetch_add(1, Ordering::AcqRel);
                let record_size = self.calculate_record_size(key.len(), 8);
                self.stats
                    .memory_usage
                    .fetch_add(record_size, Ordering::AcqRel);

                // Handle persistence if needed
                if !self.memory_only {
                    if let Some(ref wb) = self.write_buffer {
                        if let Err(e) = wb.add_write(Operation::Insert, Arc::clone(&new_record), 0)
                        {
                            // Operation succeeded in memory
                            let _ = e;
                        }
                    }
                }

                Ok(initial_val)
            }
        };

        result
    }

    /// Atomically compare and swap a value.
    ///
    /// Compares the current value of a key with an expected value, and if they match,
    /// atomically replaces it with a new value. This operation is atomic within the
    /// HashMap shard, preventing race conditions.
    ///
    /// # Arguments
    ///
    /// * `key` - The key to check and potentially update
    /// * `expected` - The expected current value
    /// * `new_value` - The new value to set if comparison succeeds
    ///
    /// # Returns
    ///
    /// Returns `Ok(true)` if the swap succeeded (current value matched expected).
    /// Returns `Ok(false)` if the current value didn't match or key doesn't exist.
    ///
    /// # Errors
    ///
    /// * `InvalidKeySize` - Key is invalid
    /// * `InvalidValueSize` - New value is too large
    /// * `OutOfMemory` - Memory limit exceeded
    /// * `IoError` - Failed to read value from disk
    ///
    /// # Example
    ///
    /// ```rust
    /// # use feoxdb::FeoxStore;
    /// # fn main() -> feoxdb::Result<()> {
    /// # let store = FeoxStore::new(None)?;
    /// store.insert(b"config", b"v1")?;
    ///
    /// // Successful CAS - value matches
    /// let swapped = store.compare_and_swap(b"config", b"v1", b"v2")?;
    /// assert_eq!(swapped, true);
    ///
    /// // Failed CAS - value doesn't match
    /// let swapped = store.compare_and_swap(b"config", b"v1", b"v3")?;
    /// assert_eq!(swapped, false); // Value is now "v2", not "v1"
    ///
    /// // CAS on non-existent key
    /// let swapped = store.compare_and_swap(b"missing", b"any", b"new")?;
    /// assert_eq!(swapped, false);
    /// # Ok(())
    /// # }
    /// ```
    pub fn compare_and_swap(&self, key: &[u8], expected: &[u8], new_value: &[u8]) -> Result<bool> {
        self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, None, 0)
    }

    /// Compare and swap with explicit timestamp.
    ///
    /// This is the advanced version that allows manual timestamp control for
    /// conflict resolution. Most users should use `compare_and_swap()` instead.
    ///
    /// # Arguments
    ///
    /// * `key` - The key to check and potentially update
    /// * `expected` - The expected current value
    /// * `new_value` - The new value to set if comparison succeeds
    /// * `timestamp` - Optional timestamp. If `None`, uses current time.
    ///
    /// # Errors
    ///
    /// * `OlderTimestamp` - Timestamp is not newer than existing record
    pub fn compare_and_swap_with_timestamp(
        &self,
        key: &[u8],
        expected: &[u8],
        new_value: &[u8],
        timestamp: Option<u64>,
    ) -> Result<bool> {
        self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, timestamp, 0)
    }

    /// Compare and swap with TTL support.
    ///
    /// # Arguments
    ///
    /// * `key` - The key to check and potentially update
    /// * `expected` - The expected current value
    /// * `new_value` - The new value to set if comparison succeeds
    /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
    ///
    /// # Errors
    ///
    /// * `InvalidKeySize` - Key is invalid
    /// * `InvalidValueSize` - New value is too large
    pub fn compare_and_swap_with_ttl(
        &self,
        key: &[u8],
        expected: &[u8],
        new_value: &[u8],
        ttl_seconds: u64,
    ) -> Result<bool> {
        self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, None, ttl_seconds)
    }

    /// Compare and swap with explicit timestamp and TTL.
    ///
    /// # Arguments
    ///
    /// * `key` - The key to check and potentially update
    /// * `expected` - The expected current value
    /// * `new_value` - The new value to set if comparison succeeds
    /// * `timestamp` - Optional timestamp. If `None`, uses current time.
    /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
    ///
    /// # Errors
    ///
    /// * `OlderTimestamp` - Timestamp is not newer than existing record
    pub fn compare_and_swap_with_timestamp_and_ttl(
        &self,
        key: &[u8],
        expected: &[u8],
        new_value: &[u8],
        timestamp: Option<u64>,
        ttl_seconds: u64,
    ) -> Result<bool> {
        let start = std::time::Instant::now();
        self.validate_key_value(key, new_value)?;
        let key_vec = key.to_vec();

        // Phase 1: Check value and save record reference for version tracking
        let initial_record = {
            let entry = match self.hash_table.read(&key_vec, |_, v| v.clone()) {
                Some(e) => e,
                None => return Ok(false), // Key doesn't exist
            };

            let record_arc = entry;

            // Check if value matches expected
            let value_matches = if let Some(val) = record_arc.get_value() {
                // Fast path: value in memory
                val.as_ref() == expected
            } else {
                // Need disk I/O

                let disk_value = self.load_value_from_disk(&record_arc)?;
                disk_value == expected
            };

            if !value_matches {
                return Ok(false); // Value doesn't match expected
            }

            // Return the Arc pointer itself as our version identifier
            // NOTE: We can't use the stored timestamp for verification here because
            // SystemTime::now() resolution is 1us, which is too coarse for CAS operations.
            record_arc
        };

        // Phase 2: Acquire write lock and verify record hasn't changed
        match self.hash_table.entry(key_vec.clone()) {
            scc::hash_map::Entry::Occupied(mut entry) => {
                let old_record = entry.get();

                // Check if the record is still the same one we read earlier
                if !Arc::ptr_eq(old_record, &initial_record) {
                    // Record was modified between our check and acquiring lock
                    return Ok(false);
                }

                let timestamp = match timestamp {
                    Some(0) | None => self.get_timestamp(),
                    Some(ts) => ts,
                };

                if timestamp < old_record.timestamp {
                    return Err(FeoxError::OlderTimestamp);
                }

                let old_size = old_record.calculate_size();
                let new_size = self.calculate_record_size(key.len(), new_value.len());
                let old_value_len = old_record.value_len;
                let old_record_arc = Arc::clone(old_record);

                // Pre-check memory limit
                if new_size > old_size && !self.check_memory_limit(new_size - old_size) {
                    return Err(FeoxError::OutOfMemory);
                }

                // Create new record with TTL if specified
                let new_record = if ttl_seconds > 0 {
                    let ttl_expiry = timestamp + (ttl_seconds * 1_000_000_000); // Convert to nanoseconds
                    Arc::new(Record::new_with_timestamp_ttl(
                        key.to_vec(),
                        new_value.to_vec(),
                        timestamp,
                        ttl_expiry,
                    ))
                } else {
                    Arc::new(Record::new(key.to_vec(), new_value.to_vec(), timestamp))
                };

                entry.insert(Arc::clone(&new_record));

                self.tree.insert(key_vec.clone(), Arc::clone(&new_record));

                if new_size > old_size {
                    self.stats
                        .memory_usage
                        .fetch_add(new_size - old_size, Ordering::AcqRel);
                } else {
                    self.stats
                        .memory_usage
                        .fetch_sub(old_size - new_size, Ordering::AcqRel);
                }

                self.stats
                    .record_insert(start.elapsed().as_nanos() as u64, true);

                if !self.memory_only {
                    if self.enable_caching {
                        if let Some(ref cache) = self.cache {
                            cache.remove(&key_vec);
                        }
                    }

                    if let Some(ref wb) = self.write_buffer {
                        let _ = wb.add_write(Operation::Update, new_record, old_value_len);
                        let _ = wb.add_write(Operation::Delete, old_record_arc, old_value_len);
                    }
                }

                Ok(true)
            }
            scc::hash_map::Entry::Vacant(_) => Ok(false),
        }
    }
}