16
\$\begingroup\$

This module has a ensure_downloaded function which takes a slice of Urls. It downloads all of the urls to a local downloads folder (if they aren't already downloaded) and returns.

use curl::easy::Easy;
use failure::Error;
use pbr::{ProgressBar, Units};
use std::fs::{rename, DirBuilder, File};
use std::io::Write;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use url::Url;

/// Return the local path where the given url will be stored.
pub fn downloaded_path(url: &Url) -> PathBuf {
    PathBuf::from("downloads")
        .join(url.host_str().unwrap())
        .join(&url.path()[1..])
}

/// Ensure that all the urls are downloaded into a local downloads directory.
pub fn ensure_downloaded(urls: &[&Url]) -> Result<(), Error> {
    // There is a single progress bar, the individual urls all contribute to it
    let mut progress = ProgressBar::new(0);
    progress.set_units(Units::Bytes);
    progress.set_max_refresh_rate(Some(Duration::from_millis(100)));
    let progress = Arc::new(Mutex::new(progress));

    // The following explicitly typed so that the closures have the correct error return type
    // and ? will work.
    let handles: Result<Vec<thread::JoinHandle<Result<(), Error>>>, Error> = urls.into_iter()
        // skip files which already exist
        .map(|url| (url, downloaded_path(url)))
        .filter(|(_, path)| !path.exists())
        .map(|(url, path)| {
            // ensure directory exists to contain the file
            DirBuilder::new()
                .recursive(true)
                .create(path.parent().unwrap())?;

            let mut easy = Easy::new();
            easy.url(url.as_str())?;
            easy.progress(true)?;

            let progress = progress.clone();

            // Create the file with a .tmp suffix
            // so if the download gets interrupted, we won't think it finished
            let mut temp_path = path.as_os_str().to_owned();
            temp_path.push(".tmp");
            let mut file = File::create(temp_path.clone())?;

            Ok(thread::spawn(move || {
                let mut transfer = easy.transfer();

                // curl doesn't accept actual errors of any sort returned here
                // so we just panic if we can't write for some reason
                transfer.write_function(|data| Ok(file.write(data).unwrap()))?;

                // We keep track of the last know total/completed
                // so we can adjust the progress bar
                let mut last_total = 0;
                let mut last_completed = 0;

                transfer.progress_function(move |total_download, completed_download, _, _| {
                    if let Ok(mut progress) = progress.lock() {
                        let total_download = total_download as u64;
                        let completed_download = completed_download as u64;

                        progress.total += total_download - last_total;
                        progress.add(completed_download - last_completed);

                        last_total = total_download;
                        last_completed = completed_download;
                    }

                    true
                })?;

                transfer.perform()?;

                // Only after completion, we rename the file to the target

                rename(temp_path, path)?;

                Ok(())
            }))
        })
        .collect();

    for handle in handles? {
        match handle.join() {
            Ok(result) => result?,
            Err(error) => return Err(format_err!("Thread failed {:?}", error)),
        }
    }

    Ok(())
}
\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy