Skip to main content

nuclei/fs/
lock.rs

1//! Advisory locks on a whole file, without blocking the worker.
2//!
3//! `lock_file` waits in the kernel until the file is free, so it is a job of the
4//! blocking pool. A future dropped while the job waits must not leave the file
5//! locked with nobody to give it back, so the future and the job share one
6//! state word per call: WAITING, HELD or GONE. After the job got the lock it
7//! moves the word from WAITING to HELD, and if it finds GONE it unlocks the
8//! file itself. A future dropped before it returned its result swaps GONE in,
9//! and if the word was HELD it unlocks the file; then its job, if it has not
10//! started, is cancelled and never locks.
11//! `proofs/tla/flock/FileLock.tla` checks the protocol.
12
13use core::future::{Future, poll_fn};
14use core::pin::Pin;
15use core::sync::atomic::{AtomicU8, Ordering};
16use core::task::{Context, Poll};
17use std::fs::{File, TryLockError};
18use std::io;
19use std::sync::Arc;
20
21use futures_io::AsyncWrite;
22
23use crate::io::Handle;
24use crate::io::source::requests::Lease;
25use crate::rt::task::Join;
26
27/// The job has not got the lock yet.
28const WAITING: u8 = 0;
29/// The job got the lock and its future is still there to take it.
30const HELD: u8 = 1;
31/// The future was dropped before it took the result.
32const GONE: u8 = 2;
33
34/// A `lock_file` call: the future of the pool's job, and what it must give back.
35struct LockCall {
36    state: Arc<AtomicU8>,
37    lease: Lease,
38    job: Join<io::Result<()>>,
39    /// The result was taken: the lock is the caller's.
40    settled: bool,
41}
42
43impl Future for LockCall {
44    type Output = io::Result<()>;
45
46    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
47        let ended = core::task::ready!(Pin::new(&mut self.job).poll(cx));
48        self.settled = true;
49        Poll::Ready(ended.unwrap_or_else(|lost| Err(lost.into())))
50    }
51}
52
53impl Drop for LockCall {
54    fn drop(&mut self) {
55        if !self.settled && self.state.swap(GONE, Ordering::AcqRel) == HELD {
56            let _ = self.lease.file().unlock();
57        }
58    }
59}
60
61impl Handle<File> {
62    /// Locks the whole file, shared for any number of holders if `shared`,
63    /// else exclusively for one, as std's `File::lock_shared` and
64    /// `File::lock` do, waiting for the lock on the blocking pool so the
65    /// worker is free meanwhile. The lock is advisory, and it is released by
66    /// [`unlock_file`](Self::unlock_file) or when the file is closed.
67    ///
68    /// If the future is dropped while it waits, the lock is never left held
69    /// with nobody to release it: a job that had not started is cancelled,
70    /// and one that got the lock gives it back.
71    ///
72    /// std's own `File::lock`, reached through `Deref`, blocks the worker
73    /// thread that makes it.
74    ///
75    /// # Errors
76    ///
77    /// The system's own, as std's.
78    pub async fn lock_file(&self, shared: bool) -> io::Result<()> {
79        let lease = self.source().lease()?;
80        let state = Arc::new(AtomicU8::new(WAITING));
81        let (job_lease, job_state) = (lease.clone(), Arc::clone(&state));
82        let job = crate::runtime::spawn_blocking(move || {
83            let file = job_lease.file();
84            let locked = if shared {
85                file.lock_shared()
86            } else {
87                file.lock()
88            };
89            if locked.is_ok()
90                && job_state
91                    .compare_exchange(WAITING, HELD, Ordering::AcqRel, Ordering::Acquire)
92                    .is_err()
93            {
94                let _ = file.unlock();
95            }
96            locked
97        })
98        .join();
99        LockCall {
100            state,
101            lease,
102            job,
103            settled: false,
104        }
105        .await
106    }
107
108    /// Locks the whole file if it is free now, shared if `shared`, else
109    /// exclusively, as std's `File::try_lock_shared` and `File::try_lock` do.
110    /// The call does not wait, so it is made on the calling thread.
111    ///
112    /// # Errors
113    ///
114    /// `WouldBlock` when another holder has the lock, else the system's own
115    /// error, as std's.
116    pub async fn try_lock_file(&self, shared: bool) -> Result<(), TryLockError> {
117        let lease = self.source().lease().map_err(TryLockError::Error)?;
118        let file = lease.file();
119        if shared {
120            file.try_lock_shared()
121        } else {
122            file.try_lock()
123        }
124    }
125
126    /// Releases the lock this file holds, as std's `File::unlock` does, once
127    /// every write accepted before the call landed (as a flush waits for
128    /// them), so a holder that wrote under the lock lets go only after its
129    /// bytes are in the file.
130    ///
131    /// # Errors
132    ///
133    /// A write that failed (reported here as by a flush), or the system's own
134    /// error. The lock is released even when a write failed.
135    pub async fn unlock_file(&self) -> io::Result<()> {
136        let mut this = self;
137        let flushed = poll_fn(|cx| Pin::new(&mut this).poll_flush(cx)).await;
138        let unlocked = self
139            .source()
140            .lease()
141            .and_then(|lease| lease.file().unlock());
142        flushed.and(unlocked)
143    }
144}