filespooler 1.2.4

Sequential, distributed, POSIX-style job queue processing
Documentation
/*! Handling the next-id sequence file

This code is used both for the queue and for the prepare step.
*/

/*
    Copyright (C) 2022 John Goerzen <jgoerzen@complete.org>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

use anyhow::Context;
use fd_lock;
use std::ffi::OsString;
use std::fs::{File, OpenOptions};
use std::io::{BufRead, BufReader, ErrorKind, Write};
use std::path::PathBuf;
use tracing::*;

/// The main file for recording sequence numbers.
///
/// Used on both the sender and recipient.
///
/// Contains one line, containing a decimal next-step counter terminated by \n.
#[derive(Debug)]
pub struct SeqFile<'a> {
    path: PathBuf,
    next: u64,
    #[allow(dead_code)]
    lock: fd_lock::RwLockWriteGuard<'a, File>,
}

/// Call this function to obtain a lock, which you can then pass to [`SeqFile::open`].
/// Due to how the lock works, it has to be held outside the SeqFile itself.
#[instrument(level = "debug")]
pub fn prepare_seqfile_lock(
    path: &OsString,
    create_if_missing: bool,
) -> Result<fd_lock::RwLock<File>, anyhow::Error> {
    let mut lockpath = path.clone();
    lockpath.push(".lock");
    debug!("Attempting to prepare lock at {:?}", lockpath);
    let lockfile = OpenOptions::new()
        .read(true)
        .write(true)
        .create(create_if_missing)
        .open(&lockpath)
        .context(format!(
            "Opening lock file at {:?}; if missing, this may be an append-only queue",
            lockpath
        ))?;

    let lock = fd_lock::RwLock::new(lockfile);
    Ok(lock)
}

impl<'a> SeqFile<'a> {
    /// Opens the seqfile.  Pass a lock that has been created for it by
    /// [`prepare_seqfile_lock`].
    ///
    /// If the file does not exist, it will be created and set to `1`.
    ///
    /// If it does exist but cannot be parsed, there will be an error.
    ///
    /// If it does exist and can be parsed, it will be treated appropriately.
    #[instrument(level = "debug", skip(lock))]
    pub fn open(
        path: &OsString,
        lock: &'a mut fd_lock::RwLock<File>,
    ) -> Result<Self, anyhow::Error> {
        let path = PathBuf::from(path);

        debug!("Attempting to acquire write lock");
        let retval = Self {
            path,
            lock: lock.try_write()?,
            next: 1,
        };

        debug!("Attempting to open file {:?}", retval.path);
        let file = match File::open(&retval.path) {
            Err(e) => {
                if e.kind() == ErrorKind::NotFound {
                    retval.write()?;
                    return Ok(retval);
                } else {
                    return Err(anyhow::Error::from(e));
                }
            }
            Ok(f) => f,
        };

        let next = SeqFile::read(file)?;
        Ok(Self { next, ..retval })
    }

    /// Sets the next counter, returning its previous value.
    pub fn set(&mut self, next: u64) -> Result<u64, anyhow::Error> {
        let retval = self.next;
        self.next = next;
        self.write()?;
        Ok(retval)
    }

    /// Increments the next counter, returning its previous value
    pub fn increment(&mut self) -> Result<u64, anyhow::Error> {
        self.set(self.next + 1)
    }

    /// Gets the current value of the counter
    pub fn get_next(&self) -> u64 {
        self.next
    }

    fn read(file: File) -> Result<u64, anyhow::Error> {
        let mut br = BufReader::new(file);
        let mut buf = String::new();
        br.read_line(&mut buf)?;
        let res: u64 = buf.trim().parse()?;
        debug!("Read next ID {} from seqfile", res);
        Ok(res)
    }

    fn write(&self) -> Result<(), anyhow::Error> {
        let mut tmppath = OsString::from(self.path.clone());
        tmppath.push(".temp");
        trace!("Writing {} to {:?}", self.next, tmppath);
        let mut file = File::create(PathBuf::from(&tmppath))?;
        writeln!(file, "{}", self.next)?;
        file.flush()?;
        file.sync_all()?;
        std::mem::drop(file);
        trace!("Renaming {:?} to {:?}", &tmppath, self.path);
        std::fs::rename(tmppath, &self.path)?;
        debug!("Sequence {} written to {:?}", self.next, self.path);

        Ok(())
    }
}

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

    use tempfile::tempdir;

    #[test]
    fn basic() {
        let dir = tempdir().unwrap();
        let path = OsString::from(dir.path().join("seqfile"));

        {
            let mut lock = prepare_seqfile_lock(&path, true).unwrap();
            let mut sf = SeqFile::open(&path, &mut lock).unwrap();

            assert_eq!(sf.get_next(), 1);

            sf.increment().unwrap();
            sf.increment().unwrap();
            assert_eq!(sf.get_next(), 3);

            // We shouldn't be able to access it while holding the lock.
            let mut lock2 = prepare_seqfile_lock(&path, false).unwrap();
            let _ = SeqFile::open(&path, &mut lock2).unwrap_err();
        }

        // OK, having dropped it, try again.
        {
            let mut lock = prepare_seqfile_lock(&path, false).unwrap();
            let sf = SeqFile::open(&path, &mut lock).unwrap();

            assert_eq!(sf.get_next(), 3);
        }

        // Truncate it and see what happens.
        File::create(&path).unwrap();
        {
            let mut lock = prepare_seqfile_lock(&path, false).unwrap();
            let _sf = SeqFile::open(&path, &mut lock).unwrap_err();
        }
    }
}