15
\$\begingroup\$

I am learning Rust. I have written a web crawler that would scrape all the pages from my own blog (which is running on Ghost) and would generate a static version of it. Because of this, I'm not interested in handling robots.txt or having rate limiting.

extern crate select;

use std::io::Read;
use select::document::Document;
use select::predicate::Name;
use select::predicate::Predicate;

use std::collections::HashSet;
use reqwest::Url;
use std::path::Path;
use std::time::Instant;
use std::fs;
use rayon::prelude::*;
use std::sync::{Arc, Mutex};

fn get_links_from_html(html: &String) -> HashSet<String> {
    Document::from(html.as_str())
        .find(Name("a").or(Name("link")))
        .filter_map(|n| n.attr("href"))
        .filter(has_extension)
        .filter_map(normalize_url)
        .collect::<HashSet<String>>()
}

fn normalize_url(url: &str) -> Option<String> {
    let new_url = Url::parse(url);
    match new_url {
        Ok(new_url) => {
            if new_url.has_host() && new_url.host_str().unwrap() == "rolisz.ro" {
                Some(url.to_string())
            } else {
                None
            }
        },
        Err(_e) => {
            // Relative urls are not parsed by Reqwest
            if url.starts_with('/') {
                Some(format!("https://rolisz.ro{}", url))
            } else {
                None
            }
        }
    }
}

fn fetch_url(client: &reqwest::blocking::Client, url: &str) -> String {
    let mut res = client.get(url).send().unwrap();
    println!("Status for {}: {}", url, res.status());

    let mut body  = String::new();
    res.read_to_string(&mut body).unwrap();
    return body
}

fn has_extension(url: &&str) -> bool {
    Path::new(&url).extension().is_none()
}

fn write_file(path: &str, content: &String) {
    let dir = fs::create_dir_all(format!("static{}", path)).unwrap();
    fs::write(format!("static{}/index.html", path), content);
}

fn main() {
    let now = Instant::now();

    let client = reqwest::blocking::Client::new();
    let origin_url = "https://rolisz.ro/";

    let body= fetch_url(&client, origin_url);

    write_file("", &body);
    let mut visited = Arc::new(Mutex::new(HashSet::new()));
    visited.lock().unwrap().insert(origin_url.to_string());
    let found_urls = get_links_from_html(&body);
    let mut new_urls = found_urls.difference(&visited.lock().unwrap()).map(|x| x.to_string()).collect::<HashSet<String>>();

    while new_urls.len() > 0 {
        let mut found_urls = Arc::new(Mutex::new(HashSet::new()));
        new_urls.par_iter().for_each(|url| {
            let body = fetch_url(&client, url);
            write_file(&url[origin_url.len()-1..], &body);

            let links = get_links_from_html(&body);
            println!("Visited: {} found {} links", url, links.len());
            found_urls.lock().unwrap().extend(links);
            visited.lock().unwrap().insert(url.to_string());
        });
        new_urls = found_urls.lock().unwrap()
            .difference(&visited.lock().unwrap()).map(|x| x.to_string())
            .collect::<HashSet<String>>();
        println!("New urls: {}", new_urls.len())
    }
    println!("URLs: {:#?}", found_urls);
    println!("{}", now.elapsed().as_secs());

}

I'm most familiar with Python, so I'm looking mostly for feedback on how to write more idiomatic Rust. One particular thing that stands out is the repetition of the .lock().unwrap() for the HashSets wrapped in Mutex and Arc. What's the most elegant way to handle that in Rust?

Any other feedback is also welcome.

\$\endgroup\$

2 Answers 2

23
\$\begingroup\$

Cargo Fmt

There's a very common tool accessible through cargo which can format all of the code in your project according to Rust's official style guide. Many major open source Rust libraries use this tool (and even enforce it through CI on pull requests), which you can access through cargo fmt. You can also customize its output using a .rustfmt config file. See the project's repo under the official rust-lang organization here.

Time to try_ again?

The closure that starts here contains several .unwrap() because some locks you do might not yield anything.

        new_urls.par_iter().for_each(|url| {
            // ...
            found_urls.lock().unwrap().extend(&links);
            visited.lock().unwrap().insert(url.to_string());

Instead of .unwrap()ing whenever something goes wrong, which could poison all of your locked Mutexes, consider try_for_each.

With try_for_each, your closure has to return a Result<T, E> or an Option<T> instead of () (nothing, unit). This allows you to use Rust's special ? operator, which is like a shorter version of .unwrap() that's actually a bit nicer because instead of crashing your program, it returns the error to be handled somewhere else.

            found_urls.lock().ok()?.extend(links);
            visited.lock().ok()?.insert(url.to_string());
            Some(())

Note that in this case, we do have to use .ok() because the PoisonError the Mutex returns also contains a reference to the Mutex, which isn't thread safe. A better practice here might be to use a custom Error enum. (more on that later)

This practice can be propagated throughout the code base.

Compiler Warnings

The Rust compiler is your friend! When I compile your code on my machine, I get several warnings.

warning: unused variable: `dir`
  --> src\main.rs:60:9
   |
60 |     let dir = fs::create_dir_all(format!("static{}", path)).unwrap();
   |         ^^^ help: consider prefixing with an underscore: `_dir`
   |
   = note: `#[warn(unused_variables)]` on by default

warning: variable does not need to be mutable
  --> src\main.rs:73:9
   |
73 |     let mut visited = Arc::new(Mutex::new(HashSet::new()));
   |         ----^^^^^^^
   |         |
   |         help: remove this `mut`
   |
   = note: `#[warn(unused_mut)]` on by default

warning: variable does not need to be mutable
  --> src\main.rs:82:13
   |
82 |         let mut found_urls = Arc::new(Mutex::new(HashSet::new()));
   |             ----^^^^^^^^^^
   |             |
   |             help: remove this `mut`

warning: unused `std::result::Result` that must be used
  --> src\main.rs:61:5
   |
61 |     fs::write(format!("static{}/index.html", path), content);
   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   |
   = note: `#[warn(unused_must_use)]` on by default
   = note: this `Result` may be an `Err` variant, which should be handled

These can all be handled in straightforward ways.

fn write_file(path: &str, content: &String) {
   let dir = fs::create_dir_all(format!("static{}", path)).unwrap();
   fs::write(format!("static{}/index.html", path), content);
}

Here, for example, you create a dir variable that's never used, and fs::write could fail but you never handle the error.

fn write_file(path: &str, content: &String) -> Result<(), ()> {

Filling in unit () types will get the compiler to tell us what types should actually go there.

PS C:\Users\cedhu\reqtwest> cargo check
    Checking reqtwest v0.1.0 (C:\Users\cedhu\reqtwest)
error[E0308]: mismatched types
  --> src\main.rs:60:5
   |
59 | fn write_file(path: &str, content: &String) -> Result<(), ()> {
   |                                                -------------- expected `std::result::Result<(), ()>` because of return type
60 |     fs::write(format!("static{}/index.html", path), content)
   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `()`, found struct `std::io::Error`
   |
   = note: expected enum `std::result::Result<_, ()>`
              found enum `std::result::Result<_, std::io::Error>`

error: aborting due to previous error

Now, thanks mostly due to the Rust compiler, we have

fn write_file(path: &str, content: &String) -> Result<(), std::io::Error> {
    fs::create_dir_all(format!("static{}", path))?;
    fs::write(format!("static{}/index.html", path), content)
}

And now the Rust compiler will continue by warning us about places where we don't handle the error that write_file now returns. Handling errors seems like a lot of typing, but it helps you make much more stable programs, and Rust's Result types really do force you to think about edge cases you would otherwise simply ignore.

write_file(&url[origin_url.len() - 1..], &body)
    .unwrap_or_else(|e| panic!("Couldn't write {:?} to file: {}", url, e));

Custom Error Type

new_urls.par_iter().try_for_each(|url| {
    let body = fetch_url(&client, url);
    write_file(&url[origin_url.len() - 1..], &body)
        .unwrap_or_else(|e| panic!("Couldn't write {:?} to file: {}", url, e));

    let links = get_links_from_html(&body);
    println!("Visited: {} found {} links", url, links.len());
    found_urls.lock().ok()?.extend(links);
    visited.lock().ok()?.insert(url.to_string());
    Some(())
}).unwrap();

Our par_iter looks like this now, but it's a bit ugly and inconsistent. For the file writing, we just panic if that creates an error, but for the .lock()s we return None through ? to create an early return. Furthermore, that None won't contain any information that could be useful for debugging. We do have information to return from the write_file, so we could stop using Option<()> and maybe start using Result<(), String>, and get some information from the failed .lock()s and put that in a String, but Strings are really big and messy, so let's try to avoid turning anything into a String until the last possible moment. To do that, we can use an enum to represent all of the possible failures our program might make.

enum Error {
    // TODO: get useful information from the mutex
    // before dropping it and store it in this error.
    Lock,

    Write {
        url: String,
        e: IoErr
    }
}

In this example, Error::Lock doesn't store any useful information quite yet, but Error::Write does. Let's start by doing two things that will make it very easy for us to change our write_file function to use our custom error type.

type Result<T> = std::result::Result<T, Error>;

This little bit of code will allow us to just write Result<()> and have that turn into Result<(), Error>, since all of the errors used in our program will use be our Error type, we don't want to have to keep typing Error all of the time.

use std::io::Error as IoErr;

// ...

impl<S: AsRef<str>> From<(S, IoErr)> for Error {
    fn from((url, e): (S, IoErr)) -> Self {
        Error::Write {
            url: url.as_ref().to_string(), 
            e
        }
    }
}

This code will let us use ? to create Errors from std::io::Errors like fs::write_file returns. This will make it very easy for us to make a write_file function that returns our custom error type.

fn write_file(path: &str, content: &String) -> Result<()> {
    let dir = format!("static{}", path);
    fs::create_dir_all(&dir).map_err(|e| (&dir, e))?;

    let index = format!("static{}/index.html", path);
    fs::write(&index, content).map_err(|e| (&index, e))?;

    Ok(())
}

The errors will now contain lots of very useful information that will make debugging a breeze.

Before we can clean up our par_iter call, now we need to take care of turning .lock() errors into our custom error type. For now, this will suffice.

impl<T> From<PoisonError<T>> for Error {
    fn from(_: PoisonError<T>) -> Self {
        //TODO: get useful information from the Mutex and store it in the Lock
        Error::Lock
    }
}

Now our par_iter call can be just:

new_urls.par_iter().try_for_each::<_, Result<()>>(|url| {
    let body = fetch_url(&client, url);
    write_file(&url[origin_url.len() - 1..], &body)?;

    let links = get_links_from_html(&body);
    println!("Visited: {} found {} links", url, links.len());
    found_urls.lock()?.extend(links);
    visited.lock()?.insert(url.to_string());
    Ok(())
}).unwrap();

There are no more .unwrap()s or even .ok()s everywhere (except for at the very end), and we can easily store information for debugging without having to panic! and risk poisoning Mutexs.

Before we go ahead and start getting useful information to put in our Lock errors, let's try to clean up fn main a bit, since it uses .unwrap() a lot as well.

fn main() -> Result<()> {

Let's have main return a result with one of our fancy custom errors.

write_file("", &body)?;
// ...
visited.lock()?.insert(origin_url.to_string());
// ...
let mut new_urls = found_urls
    .difference(&*visited.lock()?)
    // ...

Notice that we need &*visited.lock()? now to coerce our MutexGuard<HashMap<_>> into a HashMap<_>, because ? is already converting it from a Result<MutexGuard> to a MutexGuard, so we need to do two derefs, one of which will have to be explicit since only one can be done implicitly.

new_urls.par_iter().try_for_each::<_, Result<()>>(|url| {
    // ...
})?;

Notice that the par_iter call can now end simply with ? instead of .unwrap(), because our fn main returns a Result<_>.

new_urls = found_urls
    .lock()
    .unwrap()
    .difference(&visited.lock().unwrap())

becomes now just (two unwraps gone!)

new_urls = found_urls
    .lock()?
    .difference(&*visited.lock()?)

Question every use of .unwrap()

.unwrap() is a scaaaary thing. Like maggots sprinkled in your code, each of them represents a limitation on the domain of your program. Another situation in which your crawler collapses into the fetal position and cries for its mother. Another reminder that you really can't control much anything at all and you've only been lying to yourself this entire time.

Take, for example, this code.

if new_url.has_host() && new_url.host_str().unwrap() == "rolisz.ro" {
    Some(url.to_string())
} else {
    None
}

Grab the URL you found, make sure it has a host, crash the program if it doesn't ... prosaic, indeed

hold up! crash the program!? ... I don't like the sound of that! Granted, in this case, you're almost guaranteed not to crash because you just checked on the other side of the && that you wouldn't, but then... why write code that at the first glance reads like it would crash? In this case, you either want to return the url if it has "rolisz.ro" as the host string, or you want to return None in any other situation. Let Rust express what you really mean.

if let Some("rolisz.ro") = new_url.host_str() {
    Some(url.to_string())
} else {
    None
}

There are some people who use Rust who like Haskell. They might tell you to write

Ok(new_url) => new_url
    .host_str()
    .filter(|&h| h == "rolisz.ro")
    .map(|_| url.to_string()),

Before choosing this solution please keep in mind that these Haskell people are mean to other people who don't have PhDs.

Also keep in mind that there are other places where .unwrap() and .expect() are used in your application. Adding to your error enum may be necessary, but in other cases like the one above it may be as simple as rethinking the problem you were trying to solve and finding a new way to express it using the tools Rust gives you. Only add to your error enum when you're sure you've found an edge case that's probably outside of the scope of your program.

Multiple Producer, Single Consumer

.unwrap() is pretty evil, and .lock() is really just more of the same. Every .lock() you find in your codebase is an exclamation that reminds you that your program isn't truly asynchronous. Every .lock() is somewhere where your program has to wait on another thread. That's slow. Why wait when you could, like, do things?. In order to write fast programs, you want to avoid .lock() as much as you can.

This program uses .lock() a lot. That's because it's using async badly. You slapped par_iter on there, but ... each of your threads has to keep begging for the mutex to get anything done. What if, instead of waiting for exclusive control of the mutex, your threads could just ship off what they have to another thread and get right back to work?

mpsc is a channel you can use to send data across threads.

\$\endgroup\$
5
\$\begingroup\$

Clippy

Clippy helps automatically scan for common mistakes in code. Many of these are from its automatic review. Aside from the suggestions in other answers, consider:

Don't explicitly use return

in fetch_url, and just leave it as the one word body.

Don't use &String: use &str

on lines 16 and 59. You should never (unless you're doing some weird stuff with capacity) pass a &String. This is because anything you'd want to do with it can be done with a simple slice instead, i.e. &str. If you take a String as a parameter, then you require your users to allocate memory on the heap, instead of using a slice that they already have.

Don't bind a variable to an empty type

on line 60. Aside from being unused entirely, its type is (), meaning that it contains no data anyways.

Don't use extern crate

That was the old pre-2018 syntax, and you don't even use that syntax for other libraries anyways.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.