Improve concurrency with streams

examples/client_pool/Cargo.toml

@@ -8,5 +8,7 @@ name = "client_pool"
 path = "client_pool.rs"
 
 [dependencies]
+futures = "0.3.17"
+tokio-stream = "0.1.7"
 lychee-lib = { path = "../../lychee-lib", version = "0.8.1" }
 tokio = { version = "1.14.0", features = ["full"] }

examples/client_pool/client_pool.rs

@@ -1,14 +1,14 @@
-use lychee_lib::{ClientBuilder, ClientPool, Input, Request, Result, Uri};
+use lychee_lib::{ClientBuilder, Input, Request, Result, Uri};
 use std::convert::TryFrom;
 use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
 
 const CONCURRENT_REQUESTS: usize = 4;
 
 #[tokio::main]
-#[allow(clippy::trivial_regex)]
 async fn main() -> Result<()> {
     // These channels are used to send requests and receive responses to and
-    // from the lychee client pool
+    // from lychee
     let (send_req, recv_req) = mpsc::channel(CONCURRENT_REQUESTS);
     let (send_resp, mut recv_resp) = mpsc::channel(CONCURRENT_REQUESTS);
 
@@ -18,7 +18,7 @@ async fn main() -> Result<()> {
         Input::Stdin,
     )];
 
-    // Send requests to pool
+    // Queue requests
     tokio::spawn(async move {
         for request in requests {
             println!("Sending {}", request);
@@ -29,13 +29,18 @@ async fn main() -> Result<()> {
     // Create a default lychee client
     let client = ClientBuilder::default().client()?;
 
-    // Create a pool with four lychee clients
-    let clients = vec![client; CONCURRENT_REQUESTS];
-    let mut clients = ClientPool::new(send_resp, recv_req, clients);
-
-    // Handle requests in a client pool
+    // Start receiving requests
+    // Requests get streamed into the client and run concurrently
     tokio::spawn(async move {
-        clients.listen().await;
+        futures::StreamExt::for_each_concurrent(
+            ReceiverStream::new(recv_req),
+            CONCURRENT_REQUESTS,
+            |req| async {
+                let resp = client.check(req).await.unwrap();
+                send_resp.send(resp).await.unwrap();
+            },
+        )
+        .await;
     });
 
     // Finally, listen to incoming responses from lychee

examples/collect_links/Cargo.toml

@@ -12,4 +12,5 @@ lychee-lib = { path = "../../lychee-lib", version = "0.8.1" }
 tokio = { version = "1.14.0", features = ["full"] }
 regex = "1.4.6"
 http = "0.2.5"
+tokio-stream = "0.1.7"
 reqwest = { version = "0.11.7", features = ["gzip"] }

examples/collect_links/collect_links.rs

@@ -1,12 +1,13 @@
 use lychee_lib::{Collector, Input, Result};
 use reqwest::Url;
 use std::path::PathBuf;
+use tokio_stream::StreamExt;
 
 #[tokio::main]
 #[allow(clippy::trivial_regex)]
 async fn main() -> Result<()> {
     // Collect all links from the following inputs
-    let inputs: &[Input] = &[
+    let inputs = vec![
         Input::RemoteUrl(Box::new(
             Url::parse("https://github.com/lycheeverse/lychee").unwrap(),
         )),
@@ -21,6 +22,8 @@ async fn main() -> Result<()> {
     .collect_links(
         inputs, // base url or directory
     )
+    .await
+    .collect::<Result<Vec<_>>>()
     .await?;
 
     dbg!(links);

lychee-bin/Cargo.toml

@@ -39,6 +39,8 @@ structopt = "0.3.25"
 tabled = "0.3.0"
 tokio = { version = "1.14.0", features = ["full"] }
 toml = "0.5.8"
+futures = "0.3.17"
+tokio-stream = "0.1.7"
 once_cell = "1.8.0"
 
 [dev-dependencies]

lychee-bin/src/client.rs

@@ -0,0 +1,52 @@
+use crate::options::Config;
+use crate::parse::{parse_basic_auth, parse_headers, parse_statuscodes, parse_timeout};
+use anyhow::{Context, Result};
+use headers::HeaderMapExt;
+use lychee_lib::{Client, ClientBuilder};
+use regex::RegexSet;
+use std::iter::FromIterator;
+use std::{collections::HashSet, str::FromStr};
+
+/// Creates a client according to the command-line config
+pub(crate) fn create(cfg: &Config) -> Result<Client> {
+    let mut headers = parse_headers(&cfg.headers)?;
+    if let Some(auth) = &cfg.basic_auth {
+        let auth_header = parse_basic_auth(auth)?;
+        headers.typed_insert(auth_header);
+    }
+
+    let accepted = cfg.accept.clone().and_then(|a| parse_statuscodes(&a).ok());
+    let timeout = parse_timeout(cfg.timeout);
+    let method: reqwest::Method = reqwest::Method::from_str(&cfg.method.to_uppercase())?;
+    let include = RegexSet::new(&cfg.include)?;
+    let exclude = RegexSet::new(&cfg.exclude)?;
+
+    // Offline mode overrides the scheme
+    let schemes = if cfg.offline {
+        vec!["file".to_string()]
+    } else {
+        cfg.scheme.clone()
+    };
+
+    ClientBuilder::builder()
+        .includes(include)
+        .excludes(exclude)
+        .exclude_all_private(cfg.exclude_all_private)
+        .exclude_private_ips(cfg.exclude_private)
+        .exclude_link_local_ips(cfg.exclude_link_local)
+        .exclude_loopback_ips(cfg.exclude_loopback)
+        .exclude_mail(cfg.exclude_mail)
+        .max_redirects(cfg.max_redirects)
+        .user_agent(cfg.user_agent.clone())
+        .allow_insecure(cfg.insecure)
+        .custom_headers(headers)
+        .method(method)
+        .timeout(timeout)
+        .github_token(cfg.github_token.clone())
+        .schemes(HashSet::from_iter(schemes))
+        .accepted(accepted)
+        .require_https(cfg.require_https)
+        .build()
+        .client()
+        .context("Failed to create request client")
+}

lychee-bin/src/commands/check.rs

@@ -0,0 +1,116 @@
+use indicatif::ProgressBar;
+use indicatif::ProgressStyle;
+use lychee_lib::Result;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::StreamExt;
+
+use crate::{
+    options::Config,
+    stats::{color_response, ResponseStats},
+    ExitCode,
+};
+use lychee_lib::{Client, Request, Response};
+
+pub(crate) async fn check<S>(
+    client: Client,
+    requests: S,
+    cfg: &Config,
+) -> Result<(ResponseStats, ExitCode)>
+where
+    S: futures::Stream<Item = Result<Request>>,
+{
+    let (send_req, recv_req) = mpsc::channel(cfg.max_concurrency);
+    let (send_resp, mut recv_resp) = mpsc::channel(cfg.max_concurrency);
+    let max_concurrency = cfg.max_concurrency;
+    let mut stats = ResponseStats::new();
+
+    // Start receiving requests
+    tokio::spawn(async move {
+        futures::StreamExt::for_each_concurrent(
+            ReceiverStream::new(recv_req),
+            max_concurrency,
+            |request: Result<Request>| async {
+                let request: Request = request.expect("cannot read request");
+                let response = client.check(request).await.expect("cannot check request");
+                send_resp
+                    .send(response)
+                    .await
+                    .expect("cannot send response to queue");
+            },
+        )
+        .await;
+    });
+
+    let pb = if cfg.no_progress {
+        None
+    } else {
+        let bar = ProgressBar::new_spinner().with_style(ProgressStyle::default_bar().template(
+            "{spinner:.red.bright} {pos}/{len:.dim} [{elapsed_precise}] {bar:25} {wide_msg}",
+        ));
+        bar.set_length(0);
+        bar.set_message("Extracting links");
+        bar.enable_steady_tick(100);
+        Some(bar)
+    };
+
+    let bar = pb.clone();
+    let show_results_task = tokio::spawn({
+        let verbose = cfg.verbose;
+        async move {
+            while let Some(response) = recv_resp.recv().await {
+                show_progress(&pb, &response, verbose);
+                stats.add(response);
+            }
+            (pb, stats)
+        }
+    });
+
+    tokio::pin!(requests);
+
+    while let Some(request) = requests.next().await {
+        let request = request?;
+        if let Some(pb) = &bar {
+            pb.inc_length(1);
+            pb.set_message(&request.to_string());
+        };
+        send_req
+            .send(Ok(request))
+            .await
+            .expect("Cannot send request");
+    }
+    // required for the receiver task to end, which closes send_resp, which allows
+    // the show_results_task to finish
+    drop(send_req);
+
+    let (pb, stats) = show_results_task.await?;
+
+    // Note that print statements may interfere with the progress bar, so this
+    // must go before printing the stats
+    if let Some(pb) = &pb {
+        pb.finish_and_clear();
+    }
+
+    let code = if stats.is_success() {
+        ExitCode::Success
+    } else {
+        ExitCode::LinkCheckFailure
+    };
+    Ok((stats, code))
+}
+
+fn show_progress(progress_bar: &Option<ProgressBar>, response: &Response, verbose: bool) {
+    let out = color_response(&response.1);
+    if let Some(pb) = progress_bar {
+        pb.inc(1);
+        pb.set_message(&out);
+        if verbose {
+            pb.println(out);
+        }
+    } else {
+        if (response.status().is_success() || response.status().is_excluded()) && !verbose {
+            return;
+        }
+        println!("{}", out);
+    }
+}

lychee-bin/src/commands/dump.rs

@@ -0,0 +1,52 @@
+use std::io::{self, StdoutLock, Write};
+
+use lychee_lib::Result;
+use lychee_lib::{Client, Request};
+use tokio_stream::StreamExt;
+
+use crate::ExitCode;
+
+/// Dump all detected links to stdout without checking them
+pub(crate) async fn dump<'a, S>(client: Client, requests: S, verbose: bool) -> Result<ExitCode>
+where
+    S: futures::Stream<Item = Result<Request>>,
+{
+    // Lock stdout for better performance
+    let stdout = io::stdout();
+    let mut handle = stdout.lock();
+
+    tokio::pin!(requests);
+
+    while let Some(request) = requests.next().await {
+        let request = request?;
+
+        if client.filtered(&request.uri) {
+            continue;
+        }
+
+        // Avoid panic on broken pipe.
+        // See https://github.com/rust-lang/rust/issues/46016
+        // This can occur when piping the output of lychee
+        // to another program like `grep`.
+        if let Err(e) = write(&mut handle, &request, verbose) {
+            if e.kind() != io::ErrorKind::BrokenPipe {
+                eprintln!("{}", e);
+                return Ok(ExitCode::UnexpectedFailure);
+            }
+        }
+    }
+
+    Ok(ExitCode::Success)
+}
+
+/// Dump request to stdout
+/// Only print source in verbose mode. This way the normal link output
+/// can be fed into another tool without data mangling.
+fn write(handle: &mut StdoutLock<'_>, request: &Request, verbose: bool) -> io::Result<()> {
+    let output = if verbose {
+        request.to_string()
+    } else {
+        request.uri.to_string()
+    };
+    writeln!(*handle, "{}", output)
+}

lychee-bin/src/commands/mod.rs

@@ -0,0 +1,5 @@
+pub(crate) mod check;
+pub(crate) mod dump;
+
+pub(crate) use check::check;
+pub(crate) use dump::dump;