Improve concurrency with streams

examples/client_pool/Cargo.toml

@@ -10,3 +10,5 @@ path = "client_pool.rs"
 [dependencies]
 lychee-lib = { path = "../../lychee-lib", version = "0.7.2" }
 tokio = { version = "1.12.0", features = ["full"] }
+futures = "0.3.17"
+tokio-stream = "0.1.7"

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.7.2" }
 tokio = { version = "1.12.0", features = ["full"] }
 regex = "1.4.6"
 http = "0.2.5"
-reqwest = { version = "0.11.5", features = ["gzip"] }
+reqwest = { version = "0.11.5", features = ["gzip"] }
+tokio-stream = "0.1.7"

examples/collect_links/collect_links.rs

@@ -1,6 +1,7 @@
 use lychee_lib::{Collector, Input, Result};
 use reqwest::Url;
 use std::path::PathBuf;
+use tokio_stream::StreamExt;
 
 #[tokio::main]
 #[allow(clippy::trivial_regex)]
@@ -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

@@ -38,6 +38,8 @@ serde_json = "1.0.68"
 structopt = "0.3.21"
 tokio = { version = "1.12.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::{anyhow, 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()
+        .map_err(|e| anyhow!("Failed to create request client: {}", e))
+}