diff --git a/Cargo.lock b/Cargo.lock index dfa6268..b9a1f51 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1579,6 +1579,7 @@ version = "0.7.2" dependencies = [ "cached", "check-if-email-exists", + "deadpool", "doc-comment", "fast_chemail", "glob", diff --git a/examples/builder/builder.rs b/examples/builder/builder.rs index 1b3f53b..08f189e 100644 --- a/examples/builder/builder.rs +++ b/examples/builder/builder.rs @@ -40,7 +40,7 @@ async fn main() -> Result<()> { .build() .client()?; - let response = client.check("https://example.org").await?; + let response = client.check("http://example.org").await?; dbg!(&response); assert!(response.status().is_success()); Ok(()) diff --git a/examples/client_pool/client_pool.rs b/examples/client_pool/client_pool.rs index 064fe83..ef4bc8d 100644 --- a/examples/client_pool/client_pool.rs +++ b/examples/client_pool/client_pool.rs @@ -1,4 +1,4 @@ -use lychee_lib::{ClientBuilder, Input, Request, Result, Uri}; +use lychee_lib::{ClientBuilder, ClientPool, Input, Request, Result, Uri}; use std::convert::TryFrom; use tokio::sync::mpsc; @@ -9,7 +9,7 @@ const CONCURRENT_REQUESTS: usize = 4; async fn main() -> Result<()> { // These channels are used to send requests and receive responses to and // from the lychee client pool - let (send_req, mut recv_req) = mpsc::channel(CONCURRENT_REQUESTS); + let (send_req, recv_req) = mpsc::channel(CONCURRENT_REQUESTS); let (send_resp, mut recv_resp) = mpsc::channel(CONCURRENT_REQUESTS); // Add as many requests as you like @@ -29,17 +29,13 @@ 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 tokio::spawn(async move { - while let Some(req) = recv_req.recv().await { - // Client::check() may fail only because Request::try_from() may fail - // here request is already Request, so it never fails - let resp = client.check(req).await.unwrap(); - send_resp - .send(resp) - .await - .expect("Cannot send response to channel"); - } + clients.listen().await; }); // Finally, listen to incoming responses from lychee diff --git a/lychee-bin/src/main.rs b/lychee-bin/src/main.rs index 8c49c58..cb66b6b 100644 --- a/lychee-bin/src/main.rs +++ b/lychee-bin/src/main.rs @@ -69,7 +69,7 @@ use std::{collections::HashSet, fs, str::FromStr}; use anyhow::{anyhow, Context, Result}; use headers::HeaderMapExt; use indicatif::{ProgressBar, ProgressStyle}; -use lychee_lib::{ClientBuilder, Collector, Input, Request, Response}; +use lychee_lib::{ClientBuilder, ClientPool, Collector, Input, Request, Response}; use openssl_sys as _; // required for vendored-openssl feature use regex::RegexSet; use ring as _; // required for apple silicon @@ -205,7 +205,7 @@ async fn run(cfg: &Config, inputs: Vec) -> Result { .require_https(cfg.require_https) .build() .client() - .map_err(|e| anyhow!(e))?; + .map_err(|e| anyhow!("Failed to create request client: {}", e))?; let links = Collector::new(cfg.base.clone(), cfg.skip_missing, max_concurrency) .collect_links(&inputs) @@ -228,7 +228,7 @@ async fn run(cfg: &Config, inputs: Vec) -> Result { Some(bar) }; - let (send_req, mut recv_req) = mpsc::channel(max_concurrency); + let (send_req, recv_req) = mpsc::channel(max_concurrency); let (send_resp, mut recv_resp) = mpsc::channel(max_concurrency); let mut stats = ResponseStats::new(); @@ -245,15 +245,9 @@ async fn run(cfg: &Config, inputs: Vec) -> Result { // Start receiving requests tokio::spawn(async move { - while let Some(req) = recv_req.recv().await { - // `Client::check()` may fail only because `Request::try_from()` may - // fail. Here `req` is already a valid `Request`, so it never fails. - let resp = client.check(req).await.unwrap(); - send_resp - .send(resp) - .await - .expect("Cannot send response to channel"); - } + let clients = vec![client; max_concurrency]; + let mut clients = ClientPool::new(send_resp, recv_req, clients); + clients.listen().await; }); while let Some(response) = recv_resp.recv().await { diff --git a/lychee-lib/Cargo.toml b/lychee-lib/Cargo.toml index 67b5043..1ae84ab 100644 --- a/lychee-lib/Cargo.toml +++ b/lychee-lib/Cargo.toml @@ -18,6 +18,7 @@ version = "0.7.2" [dependencies] check-if-email-exists = "0.8.25" +deadpool = "0.7.0" fast_chemail = "0.9.6" glob = "0.3.0" html5ever = "0.25.1" diff --git a/lychee-lib/src/client_pool.rs b/lychee-lib/src/client_pool.rs new file mode 100644 index 0000000..a2c8a31 --- /dev/null +++ b/lychee-lib/src/client_pool.rs @@ -0,0 +1,49 @@ +use client::Client; +use deadpool::unmanaged::Pool; +use tokio::sync::mpsc; + +use crate::{client, types}; + +#[allow(missing_debug_implementations)] +/// Manages a channel for incoming requests +/// and a pool of lychee clients to handle them +/// +/// Note: Although `reqwest` has its own pool, +/// it only works for connections to the same host, so +/// a single client can still be blocked until a request is done. +pub struct ClientPool { + tx: mpsc::Sender, + rx: mpsc::Receiver, + pool: deadpool::unmanaged::Pool, +} + +impl ClientPool { + #[must_use] + /// Creates a new client pool + pub fn new( + tx: mpsc::Sender, + rx: mpsc::Receiver, + clients: Vec, + ) -> Self { + let pool = Pool::from(clients); + ClientPool { tx, rx, pool } + } + + #[allow(clippy::missing_panics_doc)] + /// Start listening for incoming requests and send each of them + /// asynchronously to a client from the pool + pub async fn listen(&mut self) { + while let Some(req) = self.rx.recv().await { + let client = self.pool.get().await; + let tx = self.tx.clone(); + tokio::spawn(async move { + // Client::check() may fail only because Request::try_from() may fail + // here request is already Request, so it never fails + let resp = client.check(req).await.unwrap(); + tx.send(resp) + .await + .expect("Cannot send response to channel"); + }); + } + } +} diff --git a/lychee-lib/src/lib.rs b/lychee-lib/src/lib.rs index e4d1de0..58c0a6a 100644 --- a/lychee-lib/src/lib.rs +++ b/lychee-lib/src/lib.rs @@ -47,6 +47,7 @@ doc_comment::doctest!("../../README.md"); mod client; +mod client_pool; /// A pool of clients, to handle concurrent checks pub mod collector; mod helpers; @@ -72,7 +73,8 @@ use ring as _; // required for apple silicon #[doc(inline)] pub use crate::{ - client::{check, Client, ClientBuilder}, + client::{check, ClientBuilder}, + client_pool::ClientPool, collector::Collector, filter::{Excludes, Filter, Includes}, types::{Base, ErrorKind, Input, Request, Response, ResponseBody, Result, Status, Uri},