generate_files.rs

extern crate scoped_pool;

use std::{env, process, fs};
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::error::Error;

use scoped_pool::Pool;

static LOREM_IPSUM: &'static [u8] =
b"Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod
tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam,
quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo
consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse
cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non
proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
";

fn main() {
    let mut args = env::args();
    args.next();
    if let Some(destination) = args.next() {
        if ! Path::new(&destination).is_dir() {
            println!("The directory '{}' doesn't exist. Create it manually.", destination);
            process::exit(1);
        }
        if let Some(count) = args.next() {
            if let Ok(count_u) = u64::from_str_radix(&count, 10) {
                generate_files(&destination, count_u);
            }
            else { usage(); }
        }
        else { usage(); }
    }
    else { usage(); }
}

fn usage() {
    println!("Usage: {} /folder/to/dump/in 10000000", env::args().next().unwrap());
    process::exit(1);
}

fn generate_files(destination: &str, count: u64) {
    let pool = Pool::new(16);
    let batch_divisor = 4096;
    let batches = count / batch_divisor;
    let last_batch_total = count % batch_divisor;
    for batch in 0..batches {
        pool.scoped(|scope| {
            for step in 0..batch_divisor {
                let seed = batch * batch_divisor + step as u64;
                let path_str = format!("{}/{:04x}/{:04x}/{:04x}/{:04x}",
                                       destination,
                                       seed.rotate_right(48)%65536,
                                       seed.rotate_right(32)%65536,
                                       seed.rotate_right(16)%65536,
                                       seed%65536);
                scope.execute(move || {
                    create_file(path_str);
                });
            }
        });
    }
    pool.scoped(|scope| {
        for seed in (batch_divisor * batches)..(batch_divisor * batches + last_batch_total) {
            let path_str = format!("{}/{:04x}/{:04x}/{:04x}/{:04x}",
                                   destination,
                                   seed.rotate_right(48)%65536,
                                   seed.rotate_right(32)%65536,
                                   seed.rotate_right(16)%65536,
                                   seed%65536);
            scope.execute(move || {
                create_file(path_str);
            });
        }
    });
}

fn create_file(path_str: String) {
    let path = Path::new(&path_str);
    let parent_dir = path.parent().unwrap();
    if ! parent_dir.exists() {
        fs::create_dir_all(parent_dir).unwrap();
    }
    match File::create(&path) {
        Err(why) => {
            println!("failed creating {}: {}",
                     path_str,
                     why.description());
        },
        Ok(mut file) => {
            match file.write_all(LOREM_IPSUM) {
                Err(why) => {
                    println!("failed writing {}: {}",
                             path_str,
                             why.description());
                },
                Ok(_) => println!("wrote {}", path_str)
            }
            file.flush().unwrap();
        }
    };
}