From d642f467ecf2fc7f82fe49b1fa1abbb49ea6d307 Mon Sep 17 00:00:00 2001 From: monosans Date: Wed, 16 Jul 2025 16:17:15 +0300 Subject: [PATCH] refactor: update proxy handling in checker and scraper, simplify function signatures --- src/checker.rs | 36 ++++++++++++++++++++++-------------- src/main.rs | 35 +++++------------------------------ src/output.rs | 12 +++--------- src/scraper.rs | 13 +++++++++---- 4 files changed, 39 insertions(+), 57 deletions(-) diff --git a/src/checker.rs b/src/checker.rs index 2a8f540..b5442b9 100644 --- a/src/checker.rs +++ b/src/checker.rs @@ -1,6 +1,6 @@ -use std::{collections::HashSet, sync::Arc}; +use std::sync::Arc; -use color_eyre::eyre::WrapErr as _; +use color_eyre::eyre::{OptionExt as _, WrapErr as _}; #[cfg(feature = "tui")] use crate::event::{AppEvent, Event}; @@ -8,25 +8,31 @@ use crate::{config::Config, proxy::Proxy, utils::pretty_error}; pub async fn check_all( config: Arc, - proxies: Arc>>, + proxies: Vec, token: tokio_util::sync::CancellationToken, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, -) -> color_eyre::Result<()> { - let workers_count = - config.checking.max_concurrent_checks.min(proxies.lock().await.len()); - if workers_count == 0 { - return Ok(()); +) -> color_eyre::Result> { + if config.checking.check_url.is_empty() { + return Ok(proxies); } - let queue = Arc::new(tokio::sync::Mutex::new( - proxies.lock().await.drain().collect::>(), - )); + let workers_count = + config.checking.max_concurrent_checks.min(proxies.len()); + if workers_count == 0 { + return Ok(Vec::new()); + } + + #[cfg(not(feature = "tui"))] + tracing::info!("Started checking {} proxies", proxies.len()); + + let queue = Arc::new(tokio::sync::Mutex::new(proxies)); + let checked_proxies = Arc::new(tokio::sync::Mutex::new(Vec::new())); let mut join_set = tokio::task::JoinSet::>::new(); for _ in 0..workers_count { let queue = Arc::clone(&queue); let config = Arc::clone(&config); - let proxies = Arc::clone(&proxies); + let checked_proxies = Arc::clone(&checked_proxies); let token = token.clone(); #[cfg(feature = "tui")] let tx = tx.clone(); @@ -49,7 +55,7 @@ pub async fn check_all( drop(tx.send(Event::App(AppEvent::ProxyWorking( proxy.protocol.clone(), )))); - proxies.lock().await.insert(proxy); + checked_proxies.lock().await.push(proxy); } Err(e) if tracing::event_enabled!( @@ -86,5 +92,7 @@ pub async fn check_all( } } - Ok(()) + Ok(Arc::into_inner(checked_proxies) + .ok_or_eyre("failed to unwrap Arc")? + .into_inner()) } diff --git a/src/main.rs b/src/main.rs index a3edbf9..22cd546 100644 --- a/src/main.rs +++ b/src/main.rs @@ -49,7 +49,8 @@ mod scraper; #[cfg(feature = "tui")] mod tui; mod utils; -use std::{collections::HashSet, path::Path, sync::Arc}; + +use std::{path::Path, sync::Arc}; use color_eyre::eyre::WrapErr as _; use tracing_subscriber::{ @@ -152,30 +153,6 @@ async fn download_output_dependencies( Ok(()) } -async fn process_proxies( - config: Arc, - proxies: Arc>>, - token: tokio_util::sync::CancellationToken, - #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender< - event::Event, - >, -) -> color_eyre::Result<()> { - if config.checking.check_url.is_empty() { - return Ok(()); - } - #[cfg(not(feature = "tui"))] - tracing::info!("Started checking {} proxies", proxies.lock().await.len()); - - checker::check_all( - config, - proxies, - token, - #[cfg(feature = "tui")] - tx.clone(), - ) - .await -} - #[cfg(unix)] async fn watch_signals(token: tokio_util::sync::CancellationToken) { let token_clone = token.clone(); @@ -252,9 +229,8 @@ async fn main_task( ) -> color_eyre::Result<()> { let http_client = create_reqwest_client() .wrap_err("failed to create reqwest HTTP client")?; - let proxies = Arc::new(tokio::sync::Mutex::new(HashSet::new())); - tokio::try_join!( + let ((), mut proxies) = tokio::try_join!( download_output_dependencies( &config, http_client.clone(), @@ -265,16 +241,15 @@ async fn main_task( scraper::scrape_all( Arc::clone(&config), http_client, - Arc::clone(&proxies), token.clone(), #[cfg(feature = "tui")] tx.clone(), ), )?; - process_proxies( + proxies = checker::check_all( Arc::clone(&config), - Arc::clone(&proxies), + proxies, token, #[cfg(feature = "tui")] tx.clone(), diff --git a/src/output.rs b/src/output.rs index 6537311..f9286d5 100644 --- a/src/output.rs +++ b/src/output.rs @@ -1,11 +1,11 @@ use std::{ - collections::{HashMap, HashSet}, + collections::HashMap, io, iter, net::{IpAddr, Ipv4Addr}, sync::Arc, }; -use color_eyre::eyre::{OptionExt as _, WrapErr as _}; +use color_eyre::eyre::WrapErr as _; use crate::{ config::Config, @@ -60,14 +60,8 @@ fn group_proxies<'a>( #[expect(clippy::too_many_lines)] pub async fn save_proxies( config: Arc, - proxies: Arc>>, + mut proxies: Vec, ) -> color_eyre::Result<()> { - let mut proxies: Vec<_> = Arc::into_inner(proxies) - .ok_or_eyre("failed to unwrap Arc")? - .into_inner() - .into_iter() - .filter(|p| config.checking.check_url.is_empty() || p.is_checked()) - .collect(); if config.output.sort_by_speed { proxies.sort_by_key(sort_by_timeout); } else { diff --git a/src/scraper.rs b/src/scraper.rs index b820005..08888cf 100644 --- a/src/scraper.rs +++ b/src/scraper.rs @@ -41,7 +41,7 @@ async fn scrape_one( source: &str, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, ) -> color_eyre::Result<()> { - let text_result = fetch_text(&config, http_client.clone(), source).await; + let text_result = fetch_text(&config, http_client, source).await; #[cfg(feature = "tui")] drop(tx.send(Event::App(AppEvent::SourceScraped(proto.clone())))); @@ -119,10 +119,11 @@ async fn scrape_one( pub async fn scrape_all( config: Arc, http_client: reqwest::Client, - proxies: Arc>>, token: tokio_util::sync::CancellationToken, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, -) -> color_eyre::Result<()> { +) -> color_eyre::Result> { + let proxies = Arc::new(tokio::sync::Mutex::new(HashSet::new())); + let mut join_set = tokio::task::JoinSet::new(); for (proto, sources) in config.scraping.sources.clone() { #[cfg(feature = "tui")] @@ -161,5 +162,9 @@ pub async fn scrape_all( .wrap_err("proxy scraping task failed")?; } - Ok(()) + Ok(Arc::into_inner(proxies) + .ok_or_eyre("failed to unwrap Arc")? + .into_inner() + .into_iter() + .collect()) }