Recursion Support (closes #55, #78) #165
+306
−61
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,4 @@ | ||
| use anyhow::{anyhow, bail, Context, Result}; | ||
| use headers::authorization::Basic; | ||
| use headers::{Authorization, HeaderMap, HeaderMapExt, HeaderName}; | ||
| use indicatif::{ProgressBar, ProgressStyle}; | ||
|
|
@@ -8,15 +8,18 @@ use stats::color_response; | |
| use std::{collections::HashSet, time::Duration}; | ||
| use std::{fs, str::FromStr}; | ||
| use structopt::StructOpt; | ||
| use tokio::sync::mpsc::{self, Sender}; | ||
|
|
||
| mod options; | ||
| mod stats; | ||
|
|
||
| use crate::options::{Config, LycheeOptions}; | ||
| use crate::stats::ResponseStats; | ||
|
|
||
| use lychee::{ | ||
| collector::{self, Input}, | ||
| Cache, Request, | ||
| }; | ||
| use lychee::{ClientBuilder, ClientPool, Response}; | ||
|
|
||
| /// A C-like enum that can be cast to `i32` and used as process exit code. | ||
|
|
@@ -86,6 +89,20 @@ fn fmt(stats: &ResponseStats, format: &Format) -> Result<String> { | |
| }) | ||
| } | ||
|
|
||
| // Get the set of input domains | ||
| // This is needed for supporting recursion | ||
| fn input_domains(inputs: Vec<Input>) -> HashSet<String> { | ||
| let mut domains = HashSet::new(); | ||
| for input in inputs { | ||
| if let Input::RemoteUrl(url) = input { | ||
| if let Some(domain) = url.domain() { | ||
| domains.insert(domain.to_string()); | ||
| } | ||
| } | ||
| } | ||
| domains | ||
| } | ||
|
|
||
| async fn run(cfg: &Config, inputs: Vec<Input>) -> Result<i32> { | ||
| let mut headers = parse_headers(&cfg.headers)?; | ||
| if let Some(auth) = &cfg.basic_auth { | ||
|
|
@@ -118,21 +135,26 @@ async fn run(cfg: &Config, inputs: Vec<Input>) -> Result<i32> { | |
| .accepted(accepted) | ||
| .build()?; | ||
|
|
||
| // Create link cache to keep track of seen links | ||
| let mut cache = Cache::new(); | ||
|
|
||
| let links = collector::collect_links( | ||
| &inputs, | ||
| cfg.base_url.clone(), | ||
| cfg.skip_missing, | ||
| max_concurrency, | ||
| ) | ||
| .await?; | ||
| let mut total_requests = links.len(); | ||
|
|
||
| let pb = match cfg.no_progress { | ||
| true => None, | ||
| false => { | ||
| let bar = ProgressBar::new(links.len() as u64) | ||
| .with_style(ProgressStyle::default_bar().template( | ||
| "{spinner:.red.bright} {pos}/{len:.dim} [{elapsed_precise}] {bar:25.magenta.bright/white} {wide_msg}", | ||
| ) | ||
| .progress_chars("██")); | ||
| bar.enable_steady_tick(100); | ||
| Some(bar) | ||
| } | ||
|
|
@@ -144,12 +166,13 @@ async fn run(cfg: &Config, inputs: Vec<Input>) -> Result<i32> { | |
| let mut stats = ResponseStats::new(); | ||
|
|
||
| let bar = pb.clone(); | ||
| let sr = send_req.clone(); | ||
| tokio::spawn(async move { | ||
| for link in links { | ||
| if let Some(pb) = &bar { | ||
| pb.set_message(&link.to_string()); | ||
| }; | ||
| sr.send(link).await.unwrap(); | ||
| } | ||
| }); | ||
|
|
||
|
|
@@ -160,9 +183,33 @@ async fn run(cfg: &Config, inputs: Vec<Input>) -> Result<i32> { | |
| clients.listen().await; | ||
| }); | ||
|
|
||
| let input_domains: HashSet<String> = input_domains(inputs); | ||
|
|
||
| // We keep track of the total number of requests | ||
| // and exit the loop once we reach it. | ||
| // Otherwise the sender would never be dropped and | ||
| // we'd be stuck indefinitely. | ||
| let mut curr = 0; | ||
|
|
||
| while curr < total_requests { | ||
| curr += 1; | ||
|
There was a problem hiding this comment. @lebensterben, I'm not too happy about the above three lines. Is there a better way to handle the request queue compared to counting the number of total requests like I did? There was a problem hiding this comment. There was a problem hiding this comment. tokio-stream crate has a wrapper type for Reciever of the channel. There was a problem hiding this comment. Ah that's great. Sounds like what we need. |
||
| let response = recv_resp.recv().await.context("Receive channel closed")?; | ||
|
|
||
| show_progress(&pb, &response, cfg.verbose); | ||
| stats.add(response.clone()); | ||
|
|
||
| if cfg.recursive { | ||
| let count = recurse( | ||
| response, | ||
| &mut cache, | ||
| &input_domains, | ||
| &cfg, | ||
| &pb, | ||
| send_req.clone(), | ||
| ) | ||
| .await?; | ||
| total_requests += count; | ||
| } | ||
| } | ||
|
|
||
| // Note that print statements may interfere with the progress bar, so this | ||
|
|
@@ -184,6 +231,68 @@ async fn run(cfg: &Config, inputs: Vec<Input>) -> Result<i32> { | |
| } | ||
| } | ||
|
|
||
| async fn recurse( | ||
| response: Response, | ||
| cache: &mut Cache, | ||
| input_domains: &HashSet<String>, | ||
| cfg: &Config, | ||
| pb: &Option<ProgressBar>, | ||
| send_req: Sender<Request>, | ||
| ) -> Result<usize> { | ||
| let recursion_level = response.recursion_level + 1; | ||
|
|
||
| if let Some(depth) = cfg.depth { | ||
| if recursion_level > depth { | ||
| // Maximum recursion depth reached; stop link checking. | ||
| return Ok(0); | ||
| } | ||
| } | ||
|
|
||
| if !response.status.is_success() { | ||
| return Ok(0); | ||
| } | ||
| if cache.contains(response.uri.as_str()) { | ||
|
There was a problem hiding this comment. Nitpick: I'm a bit skeptical of the name |
||
| return Ok(0); | ||
| } | ||
| cache.insert(response.uri.to_string()); | ||
|
There was a problem hiding this comment. Since we're allowing infinite recursion, this data structure can really take up the entire memory with sufficient number of links. Depending on the expected maximum number of links checked in recursive mode, we might consider a bloom filter here:
Super-rough calculations below. (Note that I'm just calculating the size of this one structure, there's probably going to be more memory overheads from various bits of the program). Assumptions:
Calculations:
So, HashSet with 1 million unique links will take up ~64+24+1 = ~89 MB (~85 MiB) For comparison, a Bloom Filter with ~1e-7 false positive probability is expected to take up:
So, to be honest, HashSet should be absolutely fine for the most users. But we should make a remark in the README that memory usage when using recursive mode with tens millions of links might cause RAM usage issues. There was a problem hiding this comment. Hi, would it be merged soon? and Docker image updated with the change to make GH action use the latest? |
||
|
|
||
| if let lychee::Uri::Website(url) = response.uri { | ||
| let input = collector::Input::RemoteUrl(url.clone()); | ||
|
|
||
| match url.domain() { | ||
| None => bail!("Cannot find domain in url: {}", url), | ||
| Some(domain) => { | ||
| if !input_domains.contains(domain) { | ||
| return Ok(0); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| let links = collector::collect_links( | ||
| &[input], | ||
| cfg.base_url.clone(), | ||
| cfg.skip_missing, | ||
| cfg.max_concurrency, | ||
| ) | ||
| .await?; | ||
| let count = links.len(); | ||
|
|
||
| let bar = pb.clone(); | ||
| tokio::spawn(async move { | ||
| for mut link in links { | ||
| link.recursion_level = recursion_level; | ||
| if let Some(pb) = &bar { | ||
| pb.inc_length(1); | ||
| pb.set_message(&link.to_string()); | ||
| }; | ||
| send_req.send(link).await.unwrap(); | ||
| } | ||
| }); | ||
| return Ok(count); | ||
| }; | ||
| Ok(0) | ||
| } | ||
|
|
||
| fn read_header(input: &str) -> Result<(String, String)> { | ||
| let elements: Vec<_> = input.split('=').collect(); | ||
| if elements.len() != 2 { | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I believe we should mention that depth less than zero also means infinite.