rustfs/crates/lock/src/fast_lock/state.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime};
use tokio::sync::Notify;

use crate::fast_lock::optimized_notify::OptimizedNotify;
use crate::fast_lock::types::{LockMode, LockPriority};

/// Optimized atomic lock state encoding in u64
/// Bits: [63:48] reserved | [47:32] writers_waiting | [31:16] readers_waiting | [15:8] readers_count | [7:1] flags | [0] writer_flag
const WRITER_FLAG_MASK: u64 = 0x1;
const READERS_SHIFT: u8 = 8;
const READERS_MASK: u64 = 0xFF << READERS_SHIFT; // Support up to 255 concurrent readers
const READERS_WAITING_SHIFT: u8 = 16;
const READERS_WAITING_MASK: u64 = 0xFFFF << READERS_WAITING_SHIFT;
const WRITERS_WAITING_SHIFT: u8 = 32;
const WRITERS_WAITING_MASK: u64 = 0xFFFF << WRITERS_WAITING_SHIFT;

// Fast path check masks
const NO_WRITER_AND_NO_WAITING_WRITERS: u64 = WRITER_FLAG_MASK | WRITERS_WAITING_MASK;
const COMPLETELY_UNLOCKED: u64 = 0;

/// Fast atomic lock state for single version
#[derive(Debug)]
pub struct AtomicLockState {
    state: AtomicU64,
    last_accessed: AtomicU64,
}

impl Default for AtomicLockState {
    fn default() -> Self {
        Self::new()
    }
}

impl AtomicLockState {
    pub fn new() -> Self {
        Self {
            state: AtomicU64::new(0),
            last_accessed: AtomicU64::new(
                SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .unwrap_or(Duration::ZERO)
                    .as_secs(),
            ),
        }
    }

    /// Check if fast path is available for given lock mode
    #[inline(always)]
    pub fn is_fast_path_available(&self, mode: LockMode) -> bool {
        let state = self.state.load(Ordering::Relaxed); // Use Relaxed for better performance
        match mode {
            LockMode::Shared => {
                // No writer and no waiting writers
                (state & NO_WRITER_AND_NO_WAITING_WRITERS) == 0
            }
            LockMode::Exclusive => {
                // Completely unlocked
                state == COMPLETELY_UNLOCKED
            }
        }
    }

    /// Try to acquire shared lock (fast path)
    pub fn try_acquire_shared(&self) -> bool {
        self.update_access_time();

        loop {
            let current = self.state.load(Ordering::Acquire);

            // Fast path check - cannot acquire if there's a writer or writers waiting
            if (current & NO_WRITER_AND_NO_WAITING_WRITERS) != 0 {
                return false;
            }

            let readers = self.readers_count(current);
            if readers == 0xFF {
                // Updated limit to 255
                return false; // Too many readers
            }

            let new_state = current + (1 << READERS_SHIFT);

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                return true;
            }
        }
    }

    /// Try to acquire exclusive lock (fast path)
    pub fn try_acquire_exclusive(&self) -> bool {
        self.update_access_time();

        loop {
            let current = self.state.load(Ordering::Acquire);

            // Only ownership bits may block acquisition: no writer flag, no
            // active readers. The waiting counters are preserved, not
            // required to be zero — demanding a fully-zero word means a lock
            // with registered waiters can be acquired by *no one*, including
            // the waiters themselves (each sees the others' registration),
            // so contended acquisition only succeeds in windows where every
            // waiter happens to be unregistered. `try_acquire_shared` above
            // already masks correctly; this mirrors it.
            if (current & (WRITER_FLAG_MASK | READERS_MASK)) != 0 {
                return false;
            }

            let new_state = current | WRITER_FLAG_MASK;

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                return true;
            }
        }
    }

    /// Release shared lock
    pub fn release_shared(&self) -> bool {
        loop {
            let current = self.state.load(Ordering::Acquire);
            let readers = self.readers_count(current);

            if readers == 0 {
                return false; // No shared lock to release
            }

            let new_state = current - (1 << READERS_SHIFT);

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                self.update_access_time();
                return true;
            }
        }
    }

    /// Release exclusive lock
    pub fn release_exclusive(&self) -> bool {
        loop {
            let current = self.state.load(Ordering::Acquire);

            if (current & WRITER_FLAG_MASK) == 0 {
                return false; // No exclusive lock to release
            }

            let new_state = current & !WRITER_FLAG_MASK;

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                self.update_access_time();
                return true;
            }
        }
    }

    /// Increment waiting readers count
    pub fn inc_readers_waiting(&self) -> bool {
        loop {
            let current = self.state.load(Ordering::Acquire);
            let waiting = self.readers_waiting(current);

            if waiting == 0xFFFF {
                return false; // Max waiting readers
            }

            let new_state = current + (1 << READERS_WAITING_SHIFT);

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                return true;
            }
        }
    }

    /// Decrement waiting readers count
    pub fn dec_readers_waiting(&self) {
        loop {
            let current = self.state.load(Ordering::Acquire);
            let waiting = self.readers_waiting(current);

            if waiting == 0 {
                break; // No waiting readers
            }

            let new_state = current - (1 << READERS_WAITING_SHIFT);

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                break;
            }
        }
    }

    /// Increment waiting writers count
    pub fn inc_writers_waiting(&self) -> bool {
        loop {
            let current = self.state.load(Ordering::Acquire);
            let waiting = self.writers_waiting(current);

            if waiting == 0xFFFF {
                return false; // Max waiting writers
            }

            let new_state = current + (1 << WRITERS_WAITING_SHIFT);

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                return true;
            }
        }
    }

    /// Decrement waiting writers count
    pub fn dec_writers_waiting(&self) {
        loop {
            let current = self.state.load(Ordering::Acquire);
            let waiting = self.writers_waiting(current);

            if waiting == 0 {
                break; // No waiting writers
            }

            let new_state = current - (1 << WRITERS_WAITING_SHIFT);

            if self
                .state
                .compare_exchange_weak(current, new_state, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                break;
            }
        }
    }

    /// Check if lock is completely free
    pub fn is_free(&self) -> bool {
        let state = self.state.load(Ordering::Acquire);
        state == 0
    }

    /// Check if anyone is waiting
    pub fn has_waiters(&self) -> bool {
        let state = self.state.load(Ordering::Acquire);
        self.readers_waiting(state) > 0 || self.writers_waiting(state) > 0
    }

    /// Get last access time
    pub fn last_accessed(&self) -> u64 {
        self.last_accessed.load(Ordering::Relaxed)
    }

    pub fn update_access_time(&self) {
        let now = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap_or(Duration::ZERO)
            .as_secs();
        self.last_accessed.store(now, Ordering::Relaxed);
    }

    fn readers_count(&self, state: u64) -> u8 {
        ((state & READERS_MASK) >> READERS_SHIFT) as u8
    }

    fn readers_waiting(&self, state: u64) -> u16 {
        ((state & READERS_WAITING_MASK) >> READERS_WAITING_SHIFT) as u16
    }

    fn writers_waiting(&self, state: u64) -> u16 {
        ((state & WRITERS_WAITING_MASK) >> WRITERS_WAITING_SHIFT) as u16
    }

    #[cfg(test)]
    pub fn readers_waiting_count(&self) -> u16 {
        let state = self.state.load(Ordering::Acquire);
        self.readers_waiting(state)
    }

    #[cfg(test)]
    pub fn writers_waiting_count(&self) -> u16 {
        let state = self.state.load(Ordering::Acquire);
        self.writers_waiting(state)
    }
}

/// Object lock state with version support - optimized memory layout
#[derive(Debug)]
#[repr(align(64))] // Align to cache line boundary
pub struct ObjectLockState {
    // First cache line: Most frequently accessed data
    /// Atomic state for fast operations
    pub atomic_state: AtomicLockState,

    // Second cache line: Notification mechanisms
    /// Notification for readers (traditional)
    pub read_notify: Notify,
    /// Notification for writers (traditional)
    pub write_notify: Notify,
    /// Optimized notification system (optional)
    pub optimized_notify: OptimizedNotify,

    // Third cache line: Less frequently accessed data
    /// Current owner of exclusive lock (if any)
    pub current_owner: parking_lot::RwLock<Option<ExclusiveOwnerInfo>>,
    /// Shared owners - optimized for small number of readers
    pub shared_owners: parking_lot::RwLock<smallvec::SmallVec<[SharedOwnerEntry; 4]>>,
    /// Lock priority for conflict resolution
    pub priority: parking_lot::RwLock<LockPriority>,
}

#[derive(Clone, Debug)]
pub struct ExclusiveOwnerInfo {
    pub owner: Arc<str>,
    pub acquired_at: SystemTime,
    pub lock_timeout: Duration,
}

#[derive(Clone, Debug)]
pub struct SharedOwnerEntry {
    pub owner: Arc<str>,
    pub count: u32,
    pub acquired_at: SystemTime,
    pub lock_timeout: Duration,
}

impl Default for ObjectLockState {
    fn default() -> Self {
        Self::new()
    }
}

impl ObjectLockState {
    pub fn new() -> Self {
        Self {
            atomic_state: AtomicLockState::new(),
            read_notify: Notify::new(),
            write_notify: Notify::new(),
            optimized_notify: OptimizedNotify::new(),
            current_owner: parking_lot::RwLock::new(None),
            shared_owners: parking_lot::RwLock::new(smallvec::SmallVec::new()),
            priority: parking_lot::RwLock::new(LockPriority::Normal),
        }
    }

    /// Try fast path shared lock acquisition
    pub fn try_acquire_shared_fast(&self, owner: &Arc<str>, lock_timeout: Duration) -> bool {
        if !self.atomic_state.try_acquire_shared() {
            return false;
        }

        self.atomic_state.update_access_time();
        let mut shared = self.shared_owners.write();
        if let Some(entry) = shared.iter_mut().find(|entry| entry.owner.as_ref() == owner.as_ref()) {
            entry.count = entry.count.saturating_add(1);
            entry.acquired_at = SystemTime::now();
            entry.lock_timeout = lock_timeout;
        } else {
            shared.push(SharedOwnerEntry {
                owner: owner.clone(),
                count: 1,
                acquired_at: SystemTime::now(),
                lock_timeout,
            });
        }
        true
    }

    /// Try fast path exclusive lock acquisition
    pub fn try_acquire_exclusive_fast(&self, owner: &Arc<str>, lock_timeout: Duration) -> bool {
        if !self.atomic_state.try_acquire_exclusive() {
            return false;
        }

        self.atomic_state.update_access_time();
        let mut current = self.current_owner.write();
        *current = Some(ExclusiveOwnerInfo {
            owner: owner.clone(),
            acquired_at: SystemTime::now(),
            lock_timeout,
        });
        true
    }

    /// Release shared lock
    pub fn release_shared(&self, owner: &Arc<str>) -> bool {
        let mut shared = self.shared_owners.write();
        if let Some(pos) = shared.iter().position(|entry| entry.owner.as_ref() == owner.as_ref()) {
            let original_entry = shared[pos].clone();
            let removed_entry = if shared[pos].count > 1 {
                shared[pos].count -= 1;
                None
            } else {
                Some(shared.remove(pos))
            };
            if self.atomic_state.release_shared() {
                if shared.is_empty() {
                    drop(shared);
                    self.optimized_notify.notify_writer();
                }
                true
            } else {
                tracing::warn!(
                    "Atomic state inconsistency during shared lock release: owner={}, remaining_entries={}",
                    owner,
                    shared.len()
                );
                // Re-add owner entry to maintain consistency when release failed
                match removed_entry {
                    Some(entry) => {
                        shared.push(entry);
                    }
                    None => {
                        if let Some(existing) = shared.iter_mut().find(|existing| existing.owner.as_ref() == owner.as_ref()) {
                            existing.count = existing.count.saturating_add(1);
                        } else {
                            shared.push(original_entry);
                        }
                    }
                }
                false
            }
        } else {
            tracing::debug!(
                "Shared lock release failed - owner not found: owner={}, current_entries={:?}",
                owner,
                shared.iter().map(|s| s.owner.as_ref()).collect::<Vec<_>>()
            );
            false
        }
    }

    /// Release exclusive lock
    pub fn release_exclusive(&self, owner: &Arc<str>) -> bool {
        let mut current = self.current_owner.write();
        if current.as_ref().is_some_and(|info| info.owner.as_ref() == owner.as_ref()) {
            if self.atomic_state.release_exclusive() {
                *current = None;
                drop(current);
                // Notify waiters using optimized system - prefer writers over readers
                if self
                    .atomic_state
                    .writers_waiting(self.atomic_state.state.load(Ordering::Acquire))
                    > 0
                {
                    self.optimized_notify.notify_writer();
                } else {
                    self.optimized_notify.notify_readers();
                }
                true
            } else {
                // Atomic state inconsistency - current owner matches but atomic release failed
                tracing::warn!(
                    "Atomic state inconsistency during exclusive lock release: owner={}, atomic_state={:b}",
                    owner,
                    self.atomic_state.state.load(Ordering::Acquire)
                );
                false
            }
        } else {
            // Owner mismatch
            tracing::debug!(
                "Exclusive lock release failed - owner mismatch: expected_owner={}, actual_owner={:?}",
                owner,
                current.as_ref().map(|s| s.owner.as_ref())
            );
            false
        }
    }

    /// Check if object is locked
    pub fn is_locked(&self) -> bool {
        !self.atomic_state.is_free()
    }

    /// Get current lock mode
    pub fn current_mode(&self) -> Option<LockMode> {
        let state = self.atomic_state.state.load(Ordering::Acquire);
        if (state & WRITER_FLAG_MASK) != 0 {
            Some(LockMode::Exclusive)
        } else if self.atomic_state.readers_count(state) > 0 {
            Some(LockMode::Shared)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_atomic_lock_state() {
        let state = AtomicLockState::new();

        // Test shared lock
        assert!(state.try_acquire_shared());
        assert!(state.try_acquire_shared());
        assert!(!state.try_acquire_exclusive());

        assert!(state.release_shared());
        assert!(state.release_shared());
        assert!(!state.release_shared());

        // Test exclusive lock
        assert!(state.try_acquire_exclusive());
        assert!(!state.try_acquire_shared());
        assert!(!state.try_acquire_exclusive());

        assert!(state.release_exclusive());
        assert!(!state.release_exclusive());
    }

    // Regression for the waiter-preserving exclusive CAS.
    //
    // The acquisition CAS used to demand a fully-zero state word, which
    // includes the readers_waiting/writers_waiting counters. A free lock with
    // registered waiters was then acquirable by *no one* — including the
    // waiters themselves, each blocked by the others' registration — so
    // contended acquisition only succeeded in windows where every waiter
    // happened to be unregistered. Reverting to `expected = 0` must fail here.
    #[test]
    fn exclusive_acquisition_ignores_registered_waiters() {
        let state = AtomicLockState::new();

        // Waiters register while the lock is held, then the holder releases.
        assert!(state.try_acquire_exclusive());
        assert!(state.inc_writers_waiting());
        assert!(state.inc_readers_waiting());
        assert!(state.release_exclusive());

        // The lock is now free — only the waiting counters are set.
        assert!(
            state.try_acquire_exclusive(),
            "registered waiters must not block acquisition of a free lock"
        );
        // ...and the CAS must preserve those counters, not clobber them.
        assert_eq!(state.writers_waiting_count(), 1);
        assert_eq!(state.readers_waiting_count(), 1);

        assert!(state.release_exclusive());
        state.dec_writers_waiting();

        // Ownership bits still block, registered waiters or not.
        assert!(state.try_acquire_shared());
        assert!(!state.try_acquire_exclusive(), "an active reader must still block");
        assert!(state.release_shared());

        state.dec_readers_waiting();
        assert!(state.is_free());
    }

    #[test]
    fn test_object_lock_state() {
        let state = ObjectLockState::new();
        let owner1 = Arc::from("owner1");
        let owner2 = Arc::from("owner2");

        // Test shared locks
        let timeout = Duration::from_secs(30);

        assert!(state.try_acquire_shared_fast(&owner1, timeout));
        assert!(state.try_acquire_shared_fast(&owner2, timeout));
        assert!(!state.try_acquire_exclusive_fast(&owner1, timeout));

        assert!(state.release_shared(&owner1));
        assert!(state.release_shared(&owner2));

        // Test exclusive lock
        assert!(state.try_acquire_exclusive_fast(&owner1, timeout));
        assert!(!state.try_acquire_shared_fast(&owner2, timeout));
        assert!(state.release_exclusive(&owner1));
    }
}