Concurrent Web Crawler

use regex::Regex;
use reqwest::Url;
use std::collections::BTreeSet;
use std::sync::LazyLock;

static HREF: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#"href="([^"]+)""#).unwrap());

/// Same-site http(s) links in `html`, resolved against the page's URL, without fragments.
pub fn extract_links(page: &Url, html: &str) -> BTreeSet<Url> {
    HREF.captures_iter(html)
        .filter_map(|captures| page.join(&captures[1]).ok())
        .filter(|url| matches!(url.scheme(), "http" | "https"))
        .filter(|url| url.host_str() == page.host_str() && url.port_or_known_default() == page.port_or_known_default())
        .map(|mut url| {
            url.set_fragment(None);
            url
        })
        .collect()
}

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

    fn paths(page: &str, html: &str) -> Vec<String> {
        extract_links(&Url::parse(page).unwrap(), html).iter().map(|u| u.path().to_string()).collect()
    }

    #[test]
    fn resolves_relative_links() {
        let html = r#"<a href="rust">R</a> <a href="../about">A</a> <a href="/shop">S</a>"#;
        assert_eq!(paths("http://site.test/blog/", html), ["/about", "/blog/rust", "/shop"]);
    }

    #[test]
    fn skips_other_sites_and_schemes() {
        let html = r#"<a href="https://other.test/x">x</a> <a href="mailto:me@site.test">mail</a> <a href="http://site.test:8080/">port</a>"#;
        assert!(paths("http://site.test/", html).is_empty());
    }

    #[test]
    fn drops_fragments_and_duplicates() {
        let html = r##"<a href="/about#team">t</a> <a href="/about">a</a> <a href="#top">top</a>"##;
        assert_eq!(paths("http://site.test/", html), ["/", "/about"]);
    }
}

use axum::Router;
use axum::response::{Html, Redirect};
use axum::routing::get;
use reqwest::{Client, StatusCode, header};
use std::time::Duration;

async fn slow() -> Html<&'static str> {
    tokio::time::sleep(Duration::from_secs(2)).await;
    Html("<p>finally</p>")
}

async fn start_site() -> String {
    let app = Router::new()
        .route("/", get(|| async { Html(r#"<a href="/about">About</a>"#) }))
        .route("/old", get(|| async { Redirect::permanent("/") }))
        .route("/broken", get(|| async { StatusCode::INTERNAL_SERVER_ERROR }))
        .route("/slow", get(slow))
        .route("/logo.png", get(|| async { ([(header::CONTENT_TYPE, "image/png")], vec![0u8; 1024]) }));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
    format!("http://{addr}")
}

/// Status (or error) and, for HTML pages, the body.
async fn fetch(client: &Client, url: &str) -> Result<(StatusCode, String, Option<String>), String> {
    let response = client.get(url).send().await.map_err(|e| if e.is_timeout() { String::from("timed out") } else { e.to_string() })?;
    let status = response.status();
    let final_path = response.url().path().to_string();
    let is_html = response.headers().get(header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).is_some_and(|v| v.starts_with("text/html"));
    let body = if status.is_success() && is_html { Some(response.text().await.map_err(|e| e.to_string())?) } else { None };
    Ok((status, final_path, body))
}

#[tokio::main]
async fn main() {
    let site = start_site().await;
    let client = Client::builder().timeout(Duration::from_millis(500)).user_agent("linkcheck/0.1").build().unwrap();

    for path in ["/", "/old", "/missing", "/broken", "/slow", "/logo.png"] {
        match fetch(&client, &format!("{site}{path}")).await {
            Ok((status, final_path, body)) => {
                let body = body.map_or(String::from("(not read)"), |b| format!("{} bytes of HTML", b.len()));
                println!("{path:<10} {status} -> {final_path}, {body}");
            }
            Err(error) => println!("{path:<10} error: {error}"),
        }
    }
}

use axum::Router;
use axum::response::{Html, Redirect};
use axum::routing::get;
use regex::Regex;
use reqwest::{Client, StatusCode, Url, header};
use std::collections::{BTreeSet, HashMap, HashSet};
use std::sync::{Arc, LazyLock};
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;
use tokio::task::JoinSet;

// ---------- links.rs (step 1) ----------

static HREF: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#"href="([^"]+)""#).unwrap());

fn extract_links(page: &Url, html: &str) -> BTreeSet<Url> {
    HREF.captures_iter(html)
        .filter_map(|captures| page.join(&captures[1]).ok())
        .filter(|url| matches!(url.scheme(), "http" | "https"))
        .filter(|url| url.host_str() == page.host_str() && url.port_or_known_default() == page.port_or_known_default())
        .map(|mut url| {
            url.set_fragment(None);
            url
        })
        .collect()
}

// ---------- crawl.rs ----------

enum Outcome {
    Ok(StatusCode),
    Broken(String),
}

struct Page {
    url: Url,
    outcome: Outcome,
    links: BTreeSet<Url>,
}

async fn fetch(client: &Client, url: Url) -> Page {
    let response = match client.get(url.clone()).send().await {
        Ok(response) => response,
        Err(error) => {
            let reason = if error.is_timeout() { String::from("timed out") } else { String::from("connection failed") };
            return Page { url, outcome: Outcome::Broken(reason), links: BTreeSet::new() };
        }
    };
    let status = response.status();
    if !status.is_success() {
        return Page { url, outcome: Outcome::Broken(status.to_string()), links: BTreeSet::new() };
    }
    let final_url = response.url().clone(); // after redirects
    let is_html = response.headers().get(header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).is_some_and(|v| v.starts_with("text/html"));
    let links = if is_html {
        let html = response.text().await.unwrap_or_default();
        extract_links(&final_url, &html)
    } else {
        BTreeSet::new()
    };
    Page { url, outcome: Outcome::Ok(status), links }
}

struct Report {
    pages: Vec<(usize, Page)>,          // (depth, page)
    linked_from: HashMap<Url, Url>,     // first page that linked to each URL
}

async fn crawl(client: Client, start: Url, max_depth: usize, max_in_flight: usize) -> Report {
    let permits = Arc::new(Semaphore::new(max_in_flight));
    let mut visited: HashSet<Url> = HashSet::from([start.clone()]);
    let mut linked_from = HashMap::new();
    let mut pages = Vec::new();
    let mut frontier = vec![start];

    for depth in 0..=max_depth {
        let mut tasks = JoinSet::new();
        for url in frontier.drain(..) {
            let client = client.clone();
            let permits = Arc::clone(&permits);
            tasks.spawn(async move {
                let _permit = permits.acquire_owned().await.expect("semaphore closed");
                fetch(&client, url).await
            });
        }

        let mut level: Vec<Page> = Vec::new();
        while let Some(result) = tasks.join_next().await {
            level.push(result.expect("a crawl task panicked"));
        }
        level.sort_by(|a, b| a.url.cmp(&b.url)); // completion order is random; make the report stable

        for page in &level {
            if depth == max_depth {
                break;
            }
            for link in &page.links {
                if visited.insert(link.clone()) {
                    linked_from.insert(link.clone(), page.url.clone());
                    frontier.push(link.clone());
                }
            }
        }
        pages.extend(level.into_iter().map(|page| (depth, page)));
    }
    Report { pages, linked_from }
}

// ---------- a local website to crawl ----------

async fn start_site() -> Url {
    async fn page(links: &[&str]) -> Html<String> {
        Html(links.iter().map(|l| format!(r#"<a href="{l}">{l}</a>"#)).collect::<Vec<_>>().join("\n"))
    }
    let app = Router::new()
        .route("/", get(|| page(&["/about#team", "/blog/", "/shop", "https://elsewhere.test/", "mailto:hi@site.test"])))
        .route("/about", get(|| page(&["/", "/team", "/logo.png"])))
        .route("/blog/", get(|| page(&["rust", "go", "slow", "/old-blog"])))
        .route("/blog/rust", get(|| page(&["/blog/", "async"])))
        .route("/blog/go", get(|| page(&["../about", "rust"])))
        .route("/blog/async", get(|| page(&["/unreachable"]))) // depth 3: beyond the limit
        .route("/blog/slow", get(|| async {
            tokio::time::sleep(Duration::from_secs(2)).await;
            Html("<p>too late</p>")
        }))
        .route("/old-blog", get(|| async { Redirect::permanent("/blog/") }))
        .route("/logo.png", get(|| async { ([(header::CONTENT_TYPE, "image/png")], vec![0u8; 512]) }))
        .route("/shop", get(|| async { StatusCode::INTERNAL_SERVER_ERROR }));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
    Url::parse(&format!("http://{addr}/")).unwrap()
}

// ---------- main.rs ----------

#[tokio::main]
async fn main() {
    let start = start_site().await;
    let client = Client::builder().timeout(Duration::from_millis(500)).user_agent("linkcheck/0.1").build().unwrap();

    let began = Instant::now();
    let report = crawl(client, start, 2, 4).await;
    let elapsed = began.elapsed();

    println!("crawled {} pages", report.pages.len());
    for (depth, page) in &report.pages {
        let status = match &page.outcome {
            Outcome::Ok(status) => status.as_u16().to_string(),
            Outcome::Broken(_) => String::from("---"),
        };
        println!("  depth {depth}  {status:>3}  {}", page.url.path());
    }

    println!("broken links:");
    for (_, page) in &report.pages {
        if let Outcome::Broken(reason) = &page.outcome {
            let from = report.linked_from.get(&page.url).map_or("", |u| u.path());
            println!("  {:<12} {:<26} (linked from {from})", page.url.path(), reason);
        }
    }
    // The slow page costs one 500 ms timeout; everything else ran alongside it.
    println!("finished in under 1.5 s: {}", elapsed < Duration::from_millis(1500));
}