rust-kernel

//! POSIX record locks: fcntl(F_GETLK / F_SETLK / F_SETLKW), stage 16.
//!
//! A program can lock a byte range of a file: a *read* (shared) lock, which
//! many processes may hold at once, or a *write* (exclusive) lock, which
//! nobody else may overlap. They're "advisory": read() and write() don't
//! look at them, only other fcntl calls do. Databases (SQLite) and mail
//! programs use them so two processes don't change one file at once.
//!
//! The rules, as in Linux:
//!   - locks belong to a *process* (not to an fd): locking a range you
//!     already hold replaces your old lock there (splitting it if needed);
//!     your own locks never conflict with each other;
//!   - closing *any* fd of the file drops all your locks on it, and so does
//!     exiting;
//!   - F_SETLK fails with EAGAIN on a conflict, F_SETLKW waits (a signal
//!     interrupts it), unless waiting would deadlock: EDEADLK.
//!
//! Files are told apart by a key: the path hash stat calls st_ino, or for
//! files without a path (pipes, sockets) the open file's address -- the
//! same number fstat reports as st_ino (`key_of`).
//! Ranges are [start, end) with end = u64::MAX for "to the end of the file
//! and beyond" (l_len = 0).

use crate::errno::Errno;
use crate::thread::WaitQueue;
use alloc::{collections::BTreeMap, vec::Vec};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Lock {
    pub pid: u64,
    pub start: u64,
    pub end: u64,
    pub write: bool,
}

impl Lock {
    fn overlaps(&self, start: u64, end: u64) -> bool {
        self.start < end && start < self.end
    }
}

/// More than this many locks in total: ENOLCK (keeps a runaway program
/// from eating the kernel heap).
const MAX_LOCKS: usize = 1024;

struct Table {
    /// file key -> its locks.
    files: BTreeMap<u64, Vec<Lock>>,
    /// pid -> what it waits for in F_SETLKW (for deadlock detection).
    waiting: BTreeMap<u64, (u64, Lock)>,
}

/// Locked with interrupts off (it's looked at in a wait condition).
static TABLE: Mutex<Table> = Mutex::new(Table { files: BTreeMap::new(), waiting: BTreeMap::new() });
/// F_SETLKW sleepers: woken whenever a lock goes away.
static UNLOCKED: WaitQueue = WaitQueue::new();

fn conflict_in(t: &Table, key: u64, want: &Lock) -> Option<Lock> {
    t.files.get(&key)?.iter().find(|l| l.pid != want.pid && l.overlaps(want.start, want.end) && (l.write || want.write)).copied()
}

/// The key of an open file (see the top).
pub fn key_of(file: &dyn crate::fs::file::File) -> u64 {
    match file.path() {
        Some(path) => crate::fs::file::ino_hash(path.as_bytes()),
        None => file as *const dyn crate::fs::file::File as *const u8 as u64,
    }
}

/// F_GETLK: the first lock that would stop `want`, if any.
pub fn conflict(key: u64, want: Lock) -> Option<Lock> {
    without_interrupts(|| conflict_in(&TABLE.lock(), key, &want))
}

/// Replace `pid`'s locks on [start, end) of file `key` with `new` (None =
/// unlock that range). Doesn't check for conflicts.
fn replace(t: &mut Table, key: u64, pid: u64, start: u64, end: u64, new: Option<bool>) -> Result<(), Errno> {
    let total: usize = t.files.values().map(|v| v.len()).sum();
    let locks = t.files.entry(key).or_default();
    let mut kept = Vec::with_capacity(locks.len() + 2);
    for l in locks.drain(..) {
        if l.pid != pid || !l.overlaps(start, end) {
            kept.push(l);
            continue;
        }
        // Keep the parts of our old lock outside the new range.
        if l.start < start {
            kept.push(Lock { end: start, ..l });
        }
        if l.end > end {
            kept.push(Lock { start: end, ..l });
        }
    }
    if let Some(write) = new {
        if total + 2 > MAX_LOCKS {
            *locks = kept;
            return Err(Errno::ENOLCK);
        }
        kept.push(Lock { pid, start, end, write });
    }
    *locks = kept;
    if locks.is_empty() {
        t.files.remove(&key);
    }
    Ok(())
}

/// F_SETLK: take (Some(write)) or drop (None) a lock. EAGAIN on a conflict.
pub fn set(key: u64, pid: u64, start: u64, end: u64, new: Option<bool>) -> Result<(), Errno> {
    without_interrupts(|| {
        let mut t = TABLE.lock();
        if let Some(write) = new {
            if conflict_in(&t, key, &Lock { pid, start, end, write }).is_some() {
                return Err(Errno::EAGAIN);
            }
        }
        replace(&mut t, key, pid, start, end, new)
    })?;
    UNLOCKED.wake_all(); // (something may have been unlocked)
    Ok(())
}

/// Would `pid` waiting for `want` close a circle of waiters (A waits for
/// B who waits for ... A)?
fn would_deadlock(t: &Table, key: u64, want: &Lock) -> bool {
    let mut next = conflict_in(t, key, want);
    for _ in 0..64 {
        let Some(holder) = next else { return false };
        if holder.pid == want.pid {
            return true;
        }
        next = t.waiting.get(&holder.pid).and_then(|(k, w)| conflict_in(t, *k, w));
    }
    false
}

/// F_SETLKW: like `set`, but wait while somebody else holds a conflicting
/// lock. A signal ends the wait with ERESTARTSYS (EINTR for the program).
pub fn set_wait(key: u64, pid: u64, start: u64, end: u64, write: bool) -> Result<(), Errno> {
    let want = Lock { pid, start, end, write };
    let result = UNLOCKED.wait_until_intr(|| {
        let mut t = TABLE.lock();
        if conflict_in(&t, key, &want).is_none() {
            t.waiting.remove(&pid);
            return Some(replace(&mut t, key, pid, start, end, Some(write)));
        }
        if would_deadlock(&t, key, &want) {
            t.waiting.remove(&pid);
            return Some(Err(Errno::EDEADLK));
        }
        t.waiting.insert(pid, (key, want));
        None
    });
    without_interrupts(|| TABLE.lock().waiting.remove(&pid));
    result?
}

/// `pid` closed an fd of file `key`: all its locks on that file go.
pub fn file_closed(key: u64, pid: u64) {
    let removed = without_interrupts(|| {
        let mut t = TABLE.lock();
        let Some(locks) = t.files.get_mut(&key) else { return false };
        let before = locks.len();
        locks.retain(|l| l.pid != pid);
        let removed = locks.len() != before;
        if locks.is_empty() {
            t.files.remove(&key);
        }
        removed
    });
    if removed {
        UNLOCKED.wake_all();
    }
}

/// Process `pid` exited: drop all its locks.
pub fn process_exited(pid: u64) {
    without_interrupts(|| {
        let mut t = TABLE.lock();
        t.waiting.remove(&pid);
        for locks in t.files.values_mut() {
            locks.retain(|l| l.pid != pid);
        }
        t.files.retain(|_, v| !v.is_empty());
    });
    UNLOCKED.wake_all();
}

/// How many locks exist (tests: nothing left behind).
pub fn count() -> usize {
    without_interrupts(|| TABLE.lock().files.values().map(|v| v.len()).sum())
}