<commit_before>#[macro_use]
extern crate serde_derive;

extern crate futures;
extern crate serde;
extern crate serde_json;
extern crate textnonce;

use std::error::Error as StdError;
use std::os::unix::fs::{DirBuilderExt, OpenOptionsExt};

use serde::{Serialize, Deserialize};

/// A durable queue backed by the filesystem.
///
/// This queue stores items of type T as files in a spool directory. Files are serialized with
/// serde_json so T must derive serde's Serialize and Deserialize traits.
///
/// The queue's directory should only contain items of the same type. Any files in the spool
/// that fail to deserialize will be discarded.
#[derive(Clone, Debug)]
pub struct Queue<T> {
    path: String,
    seq: u32,
    _placeholder: std::marker::PhantomData<T>,
}

impl<T: Serialize + Deserialize> Queue<T> {
    /// Create a new `Queue<T>` using the given directory path for storage.
    pub fn new(path: &str) -> Result<Queue<T>, std::io::Error> {
        std::fs::DirBuilder::new()
            .recursive(true)
            .mode(0o700)
            .create(path)?;
        Ok(Queue::<T> {
               path: path.to_string(),
               seq: 0,
               _placeholder: std::marker::PhantomData,
           })
    }

    /// Push an item into the Queue.
    pub fn push(&mut self, item: T) -> Result<(), std::io::Error> {
        let mut item_path = std::path::PathBuf::from(&self.path);
        let item_name = format!("{:016x}-{}", self.seq, rand_string());
        item_path.push(item_name);
        let complete_path = item_path.to_str().unwrap().to_string();
        let incomplete_path = item_path
            .with_extension("inc")
            .to_str()
            .unwrap()
            .to_string();

        {
            let mut item_file = std::fs::OpenOptions::new()
                .write(true)
                .mode(0o600)
                .create_new(true)
                .open(&incomplete_path)?;
            serde_json::to_writer(&mut item_file, &item)
                .map_err(to_ioerror)?;
        }
        std::fs::rename(incomplete_path, complete_path)?;
        self.seq += 1;
        Ok(())
    }

    /// Pop an item off the queue.
    ///
    /// This method returns the first matching directory entry. Queue ordering cannot be guaranteed
    /// to be consistent across all operating systems and filesystems, as the serialized file will
    /// be chosen based on the filesystem's directory entry ordering.
    ///
    /// Popped items are not removed from the filesystem immediately; instead, they are marked for
    /// deletion. Use flush() to cause the items to be permanently removed from the underlying
    /// filesystem.
    pub fn pop(&self) -> Result<Option<T>, std::io::Error> {
        self.pop_filter(|_| true)
    }

    /// Pop an item off the queue matching the filter function.
    ///
    /// This method works the same as pop(), except the item must match the given function or it is
    /// not popped off the queue.
    pub fn pop_filter<F>(&self, filter: F) -> Result<Option<T>, std::io::Error>
        where F: Fn(&T) -> bool
    {
        let dirh = std::fs::read_dir(&self.path)?;
        for maybe_dirent in dirh {
            let item_path = match maybe_dirent {
                Ok(dirent) => {
                    let p = dirent.path();
                    if let Some(_) = p.extension() {
                        continue;
                    }
                    p
                }
                Err(e) => return Err(e),
            };
            let stage_path = item_path.with_extension("pop");
            {
                let item_file = std::fs::OpenOptions::new().read(true).open(&item_path)?;
                let item = serde_json::from_reader(item_file).map_err(to_ioerror)?;
                if filter(&item) {
                    std::fs::rename(item_path, stage_path)?;
                    return Ok(Some(item));
                }
            }
        }
        Ok(None)
    }

    /// Flush removes all pending item files marked for deletion.
    pub fn flush(&self) -> Result<(), std::io::Error> {
        let dirh = std::fs::read_dir(&self.path)?;
        for maybe_dirent in dirh {
            match maybe_dirent {
                Ok(dirent) => {
                    let p = dirent.path();
                    match p.extension() {
                        Some(e) => {
                            if e != "pop" {
                                continue;
                            }
                        }
                        None => continue,
                    }
                    std::fs::remove_file(p)?;
                }
                Err(e) => return Err(e),
            }
        }<commit_msg>Add drain method.<commit_after>#[macro_use]
extern crate serde_derive;

extern crate futures;
extern crate serde;
extern crate serde_json;
extern crate textnonce;

use std::error::Error as StdError;
use std::os::unix::fs::{DirBuilderExt, OpenOptionsExt};

use serde::{Serialize, Deserialize};

/// A durable queue backed by the filesystem.
///
/// This queue stores items of type T as files in a spool directory. Files are serialized with
/// serde_json so T must derive serde's Serialize and Deserialize traits.
///
/// The queue's directory should only contain items of the same type. Any files in the spool
/// that fail to deserialize will be discarded.
#[derive(Clone, Debug)]
pub struct Queue<T> {
    path: String,
    seq: u32,
    _placeholder: std::marker::PhantomData<T>,
}

impl<T: Serialize + Deserialize> Queue<T> {
    /// Create a new `Queue<T>` using the given directory path for storage.
    pub fn new(path: &str) -> Result<Queue<T>, std::io::Error> {
        std::fs::DirBuilder::new()
            .recursive(true)
            .mode(0o700)
            .create(path)?;
        Ok(Queue::<T> {
               path: path.to_string(),
               seq: 0,
               _placeholder: std::marker::PhantomData,
           })
    }

    /// Push an item into the Queue.
    pub fn push(&mut self, item: T) -> Result<(), std::io::Error> {
        let mut item_path = std::path::PathBuf::from(&self.path);
        let item_name = format!("{:016x}-{}", self.seq, rand_string());
        item_path.push(item_name);
        let complete_path = item_path.to_str().unwrap().to_string();
        let incomplete_path = item_path
            .with_extension("inc")
            .to_str()
            .unwrap()
            .to_string();

        {
            let mut item_file = std::fs::OpenOptions::new()
                .write(true)
                .mode(0o600)
                .create_new(true)
                .open(&incomplete_path)?;
            serde_json::to_writer(&mut item_file, &item)
                .map_err(to_ioerror)?;
        }
        std::fs::rename(incomplete_path, complete_path)?;
        self.seq += 1;
        Ok(())
    }

    /// Pop an item off the queue.
    ///
    /// This method returns the first matching directory entry. Queue ordering cannot be guaranteed
    /// to be consistent across all operating systems and filesystems, as the serialized file will
    /// be chosen based on the filesystem's directory entry ordering.
    ///
    /// Popped items are not removed from the filesystem immediately; instead, they are marked for
    /// deletion. Use flush() to cause the items to be permanently removed from the underlying
    /// filesystem.
    pub fn pop(&self) -> Result<Option<T>, std::io::Error> {
        self.pop_filter(|_| true)
    }

    /// Pop an item off the queue matching the filter function.
    ///
    /// This method works the same as pop(), except the item must match the given function or it is
    /// not popped off the queue.
    pub fn pop_filter<F>(&self, filter: F) -> Result<Option<T>, std::io::Error>
        where F: Fn(&T) -> bool
    {
        match self.drain_filter(filter, Some(1)) {
            Ok(ref mut result) => {
                if result.len() > 0 {
                    return Ok(Some(result.remove(0)));
                }
            }
            Err(e) => return Err(e),
        }
        Ok(None)
    }

    pub fn drain(&self) -> Result<Vec<T>, std::io::Error> {
        self.drain_filter(|_| true, None)
    }

    fn drain_filter<F>(&self, filter: F, limit: Option<usize>) -> Result<Vec<T>, std::io::Error>
        where F: Fn(&T) -> bool
    {
        let mut result = vec![];
        let dirh = std::fs::read_dir(&self.path)?;
        for maybe_dirent in dirh {
            let item_path = match maybe_dirent {
                Ok(dirent) => {
                    let p = dirent.path();
                    if let Some(_) = p.extension() {
                        continue;
                    }
                    p
                }
                Err(e) => return Err(e),
            };
            let stage_path = item_path.with_extension("pop");
            {
                let item_file = std::fs::OpenOptions::new().read(true).open(&item_path)?;
                let item = serde_json::from_reader(item_file).map_err(to_ioerror)?;
                if filter(&item) {
                    std::fs::rename(item_path, stage_path)?;
                    result.push(item);
                }
                if let Some(l) = limit {
                    if result.len() == l {
                        return Ok(result);
                    }
                }
            }
        }
        Ok(result)
    }

    /// Flush removes all pending item files marked for deletion.
    pub fn flush(&self) -> Result<(), std::io::Error> {
        let dirh = std::fs::read_dir(&self.path)?;
        for maybe_dirent in dirh {
            match maybe_dirent {
                Ok(dirent) => {
                    let p = dirent.path();
                    match p.extension() {
                        Some(e) => {
                            if e != "pop" {
                                continue;
                            }
                        }
                        None => continue,
                    }
                    std::fs::remove_file(p)?;
                }
                Err(e) => return Err(e),
            }
        }
{'fraction_non_alphanumeric': 0.12481528913407546, 'fraction_numerical': 0.0032509112402718944, 'mean_word_length': 1.6348300025953801, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 69, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4761595', 'n_tokens_mistral': 2917, 'n_tokens_neox': 2814, 'n_words': 809}