From 234fd90580f7a0cf05d3a04506b8bf8b716a075c Mon Sep 17 00:00:00 2001 From: lifegpc Date: Tue, 10 May 2022 22:01:12 +0800 Subject: [PATCH] WebClient change all &mut self to &self --- src/cookies.rs | 4 + src/download.rs | 30 ++++--- src/pixiv_web.rs | 88 +++++++++++++++------ src/webclient.rs | 198 +++++++++++++++++++++++++++++++++++++---------- 4 files changed, 241 insertions(+), 79 deletions(-) diff --git a/src/cookies.rs b/src/cookies.rs index 3fbc075..de7c038 100644 --- a/src/cookies.rs +++ b/src/cookies.rs @@ -307,6 +307,10 @@ impl CookieJar { } } + pub fn clear(&mut self) { + self.cookies.clear(); + } + pub fn read(&mut self, file_name: &str) -> bool { self.cookies.clear(); let p = Path::new(file_name); diff --git a/src/download.rs b/src/download.rs index b08879d..dd7d6a2 100644 --- a/src/download.rs +++ b/src/download.rs @@ -21,19 +21,18 @@ use json::JsonValue; use spin_on::spin_on; use std::path::PathBuf; use std::sync::Arc; -use futures_util::lock::Mutex; impl Main { pub fn download(&mut self) -> i32 { - let pw = Arc::new(Mutex::new(PixivWebClient::new(&self))); - if !pw.try_lock().unwrap().init() { + let pw = Arc::new(PixivWebClient::new(&self)); + if !pw.init() { println!("{}", gettext("Failed to initialize pixiv web api client.")); return 1; } - if !pw.try_lock().unwrap().check_login() { + if !pw.check_login() { return 1; } - if !pw.try_lock().unwrap().logined() { + if !pw.logined() { println!( "{}", gettext("Warning: Web api client not logined, some future may not work.") @@ -52,7 +51,7 @@ impl Main { 0 } - pub async fn download_artwork_page(pw: Arc>, page: JsonValue, np: u16, progress_bars: Arc, datas: Arc, base: Arc, helper: Arc) -> i32 { + pub async fn download_artwork_page(pw: Arc, page: JsonValue, np: u16, progress_bars: Arc, datas: Arc, base: Arc, helper: Arc) -> i32 { let link = &page["urls"]["original"]; if !link.is_string() { println!("{}", gettext("Failed to get original picture's link.")); @@ -94,7 +93,6 @@ impl Main { } let r; { - let mut pw = pw.lock().await; r = pw.adownload_image(link).await; if r.is_none() { println!("{} {}", gettext("Failed to download image:"), link); @@ -120,19 +118,19 @@ impl Main { 0 } - pub fn download_artwork(&self, pw: Arc>, id: u64) -> i32 { + pub fn download_artwork(&self, pw: Arc, id: u64) -> i32 { let mut re = None; let pages; let mut ajax_ver = true; - let helper = Arc::new(pw.try_lock().unwrap().helper.clone()); + let helper = Arc::new(pw.helper.clone()); if helper.use_webpage() { - re = pw.try_lock().unwrap().get_artwork(id); + re = pw.get_artwork(id); if re.is_some() { ajax_ver = false; } } if re.is_none() { - re = pw.try_lock().unwrap().get_artwork_ajax(id); + re = pw.get_artwork_ajax(id); } if re.is_none() { return 1; @@ -150,7 +148,7 @@ impl Main { let pages = pages.unwrap(); let mut pages_data: Option = None; if pages > 1 { - pages_data = pw.try_lock().unwrap().get_illust_pages(id); + pages_data = pw.get_illust_pages(id); } if pages > 1 && pages_data.is_none() { println!("{}", gettext("Failed to get pages' data.")); @@ -180,7 +178,7 @@ impl Main { match illust_type { 0 => { } 2 => { - let ugoira_data = pw.try_lock().unwrap().get_ugoira(id); + let ugoira_data = pw.get_ugoira(id); if ugoira_data.is_none() { println!("{}", gettext("Failed to get ugoira's data.")); return 1; @@ -208,7 +206,7 @@ impl Main { true }; if dw { - let r = pw.try_lock().unwrap().download_image(src); + let r = pw.download_image(src); if r.is_none() { println!("{} {}", gettext("Failed to download ugoira:"), src); return 1; @@ -321,7 +319,7 @@ impl Main { } } } - let r = pw.try_lock().unwrap().download_image(link); + let r = pw.download_image(link); if r.is_none() { println!("{} {}", gettext("Failed to download image:"), link); return 1; @@ -391,7 +389,7 @@ impl Main { return 0; } } - let r = pw.try_lock().unwrap().download_image(link); + let r = pw.download_image(link); if r.is_none() { println!("{} {}", gettext("Failed to download image:"), link); return 1; diff --git a/src/pixiv_web.rs b/src/pixiv_web.rs index b71ae60..caa62cb 100644 --- a/src/pixiv_web.rs +++ b/src/pixiv_web.rs @@ -7,14 +7,21 @@ use json::JsonValue; use reqwest::IntoUrl; use reqwest::Response; use spin_on::spin_on; +use std::sync::Arc; +use std::sync::RwLock; +use std::sync::RwLockReadGuard; +use std::sync::RwLockWriteGuard; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; /// A client which use Pixiv's web API pub struct PixivWebClient { client: WebClient, pub helper: OptHelper, /// true if in is initialized - inited: bool, - data: Option, + inited: Arc, + data: RwLock>, } impl PixivWebClient { @@ -22,12 +29,46 @@ impl PixivWebClient { Self { client: WebClient::new(), helper: OptHelper::new(m.cmd.as_ref().unwrap().clone(), m.settings.as_ref().unwrap().clone()), - inited: false, - data: None, + inited: Arc::new(AtomicBool::new(false)), + data: RwLock::new(None), } } - pub fn init(&mut self) -> bool { + async fn aget_data_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, Option> { + loop { + match self.data.try_write() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + fn get_data_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, Option> { + spin_on(self.aget_data_as_mut()) + } + + async fn aget_data<'a>(&'a self) -> RwLockReadGuard<'a, Option> { + loop { + match self.data.try_read() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + fn get_data<'a>(&'a self) -> RwLockReadGuard<'a, Option> { + spin_on(self.aget_data()) + } + + pub fn is_inited(&self) -> bool { + self.inited.load(Ordering::Relaxed) + } + + pub fn init(&self) -> bool { let c = self.helper.cookies(); if c.is_some() { if !self.client.read_cookies(c.as_ref().unwrap()) { @@ -41,18 +82,18 @@ impl PixivWebClient { } else { self.client.set_header("Accept-Language", "ja"); } - self.client.verbose = self.helper.verbose(); + self.client.set_verbose(self.helper.verbose()); let retry = self.helper.retry(); if retry.is_some() { - self.client.retry = retry.unwrap(); + self.client.set_retry(retry.unwrap()); } - self.client.retry_interval = Some(self.helper.retry_interval()); - self.inited = true; + self.client.get_retry_interval_as_mut().replace(self.helper.retry_interval()); + self.inited.store(true, Ordering::Relaxed); true } - pub fn auto_init(&mut self) { - if !self.inited { + pub fn auto_init(&self) { + if !self.is_inited() { let r = self.init(); if !r { panic!("{}", gettext("Failed to initialize pixiv web api client.")); @@ -60,7 +101,7 @@ impl PixivWebClient { } } - pub fn check_login(&mut self) -> bool { + pub fn check_login(&self) -> bool { self.auto_init(); let r = self.client.get("https://www.pixiv.net/", None); if r.is_none() { @@ -87,11 +128,11 @@ impl PixivWebClient { if self.helper.verbose() { println!("{}\n{}", gettext("Main page's data:"), p.value.as_ref().unwrap().pretty(2).as_str()); } - self.data = Some(p.value.unwrap()); + self.get_data_as_mut().replace(p.value.unwrap()); true } - pub fn deal_json(&mut self, r: Response) -> Option { + pub fn deal_json(&self, r: Response) -> Option { let status = r.status(); let code = status.as_u16(); let is_status_err = code >= 400; @@ -140,7 +181,7 @@ impl PixivWebClient { Some(body.clone()) } - pub fn download_image(&mut self, url: U) -> Option { + pub fn download_image(&self, url: U) -> Option { self.auto_init(); let r = self.client.get(url, json::object!{"referer": "https://www.pixiv.net/"}); if r.is_none() { @@ -156,7 +197,7 @@ impl PixivWebClient { Some(r) } - pub async fn adownload_image(&mut self, url: U) -> Option { + pub async fn adownload_image(&self, url: U) -> Option { self.auto_init(); let r = self.client.aget(url, json::object!{"referer": "https://www.pixiv.net/"}).await; if r.is_none() { @@ -172,7 +213,7 @@ impl PixivWebClient { Some(r) } - pub fn get_artwork_ajax(&mut self, id: u64) -> Option { + pub fn get_artwork_ajax(&self, id: u64) -> Option { self.auto_init(); let r = self.client.get(format!("https://www.pixiv.net/ajax/illust/{}", id), None); if r.is_none() { @@ -186,7 +227,7 @@ impl PixivWebClient { v } - pub fn get_artwork(&mut self, id: u64) -> Option { + pub fn get_artwork(&self, id: u64) -> Option { self.auto_init(); let r = self.client.get(format!("https://www.pixiv.net/artworks/{}", id), None); if r.is_none() { @@ -216,7 +257,7 @@ impl PixivWebClient { Some(p.value.unwrap()) } - pub fn get_illust_pages(&mut self, id: u64) -> Option { + pub fn get_illust_pages(&self, id: u64) -> Option { self.auto_init(); let r = self.client.get(format!("https://www.pixiv.net/ajax/illust/{}/pages", id), None); if r.is_none() { @@ -230,7 +271,7 @@ impl PixivWebClient { v } - pub fn get_ugoira(&mut self, id: u64) -> Option { + pub fn get_ugoira(&self, id: u64) -> Option { self.auto_init(); let r = self.client.get(format!("https://www.pixiv.net/ajax/illust/{}/ugoira_meta", id), None); if r.is_none() { @@ -245,10 +286,11 @@ impl PixivWebClient { } pub fn logined(&self) -> bool { - if self.data.is_none() { + let data = self.get_data(); + if data.is_none() { return false; } - let value = self.data.as_ref().unwrap(); + let value = data.as_ref().unwrap(); let d = &value["userData"]; if d.is_object() { return true; @@ -259,7 +301,7 @@ impl PixivWebClient { impl Drop for PixivWebClient { fn drop(&mut self) { - if self.inited { + if self.is_inited() { let c = self.helper.cookies(); if c.is_some() { if !self.client.save_cookies(c.as_ref().unwrap()) { diff --git a/src/webclient.rs b/src/webclient.rs index b5a1db0..7636635 100644 --- a/src/webclient.rs +++ b/src/webclient.rs @@ -20,6 +20,12 @@ use std::fs::remove_file; use std::io::Write; use std::path::Path; use std::sync::Arc; +use std::sync::RwLock; +use std::sync::RwLockReadGuard; +use std::sync::RwLockWriteGuard; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use std::time::Duration; pub trait ToHeaders { @@ -59,11 +65,11 @@ impl ToHeaders for JsonValue { /// Generate `cookie` header for a url /// * `c` - Cookies /// * `url` - URL -pub fn gen_cookie_header(c: &mut CookieJar, url: U) -> String { - c.check_expired(); +pub fn gen_cookie_header(c: &WebClient, url: U) -> String { + c.get_cookies_as_mut().check_expired(); let mut s = String::from(""); let u = url.as_str(); - for a in c.iter() { + for a in c.get_cookies().iter() { if a.matched(u) { if s.len() > 0 { s += " "; @@ -78,28 +84,127 @@ pub fn gen_cookie_header(c: &mut CookieJar, url: U) -> String { pub struct WebClient { client: Client, /// HTTP Headers - pub headers: HashMap, - cookies: CookieJar, - pub verbose: bool, + headers: RwLock>, + cookies: RwLock, + verbose: Arc, /// Retry times, 0 means disable - pub retry: u64, + retry: Arc, /// Retry interval - pub retry_interval: Option>, + retry_interval: RwLock>>, } impl WebClient { pub fn new() -> Self { Self { client: Client::new(), - headers: HashMap::new(), - cookies: CookieJar::new(), - verbose: false, - retry: 3, - retry_interval: None, + headers: RwLock::new(HashMap::new()), + cookies: RwLock::new(CookieJar::new()), + verbose: Arc::new(AtomicBool::new(false)), + retry: Arc::new(AtomicU64::new(3)), + retry_interval: RwLock::new(None), } } - pub fn handle_set_cookie(&mut self, r: &Response) { + pub async fn aget_cookies_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, CookieJar> { + loop { + match self.cookies.try_write() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + pub fn get_cookies_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, CookieJar> { + spin_on(self.aget_cookies_as_mut()) + } + + pub async fn aget_cookies<'a>(&'a self) -> RwLockReadGuard<'a, CookieJar> { + loop { + match self.cookies.try_read() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + pub fn get_cookies<'a>(&'a self) -> RwLockReadGuard<'a, CookieJar> { + spin_on(self.aget_cookies()) + } + + pub async fn aget_headers_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, HashMap> { + loop { + match self.headers.try_write() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + pub fn get_headers_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, HashMap> { + spin_on(self.aget_headers_as_mut()) + } + + pub async fn aget_headers<'a>(&'a self) -> RwLockReadGuard<'a, HashMap> { + loop { + match self.headers.try_read() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + pub fn get_headers<'a>(&'a self) -> RwLockReadGuard<'a, HashMap> { + spin_on(self.aget_headers()) + } + + /// return retry times, 0 means disable + pub fn get_retry(&self) -> u64 { + self.retry.load(Ordering::Relaxed) + } + + pub async fn aget_retry_interval_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, Option>> { + loop { + match self.retry_interval.try_write() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + pub fn get_retry_interval_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, Option>> { + spin_on(self.aget_retry_interval_as_mut()) + } + + pub async fn aget_retry_interval<'a>(&'a self) -> RwLockReadGuard<'a, Option>> { + loop { + match self.retry_interval.try_read() { + Ok(f) => { return f; } + Err(_) => { + tokio::time::sleep(Duration::new(0, 1_000_000)).await; + } + } + } + } + + pub fn get_retry_interval<'a>(&'a self) -> RwLockReadGuard<'a, Option>> { + spin_on(self.aget_retry_interval()) + } + + pub fn get_verbose(&self) -> bool { + self.verbose.load(Ordering::Relaxed) + } + + pub fn handle_set_cookie(&self, r: &Response) { let u = r.url(); let h = r.headers(); let v = h.get_all("Set-Cookie"); @@ -110,7 +215,7 @@ impl WebClient { let c = Cookie::from_set_cookie(u.as_str(), val); match c { Some(c) => { - self.cookies.add(c); + self.get_cookies_as_mut().add(c); } None => { println!("{}", gettext("Failed to parse Set-Cookie header.")); @@ -124,20 +229,30 @@ impl WebClient { } } - pub fn read_cookies(&mut self, file_name: &str) -> bool { - let r = self.cookies.read(file_name); + pub fn read_cookies(&self, file_name: &str) -> bool { + let mut c = self.get_cookies_as_mut(); + let r = c.read(file_name); if !r { - self.cookies = CookieJar::new(); + c.clear(); } r } - pub fn save_cookies(&mut self, file_name: &str) -> bool { - self.cookies.save(file_name) + pub fn save_cookies(&self, file_name: &str) -> bool { + self.get_cookies_as_mut().save(file_name) } - pub fn set_header(&mut self, key: &str, value: &str) -> Option { - self.headers.insert(String::from(key), String::from(value)) + pub fn set_header(&self, key: &str, value: &str) -> Option { + self.get_headers_as_mut().insert(String::from(key), String::from(value)) + } + + /// Set retry times, 0 means disable + pub fn set_retry(&self, retry: u64) { + self.retry.store(retry, Ordering::Relaxed) + } + + pub fn set_verbose(&self, verbose: bool) { + self.verbose.store(verbose, Ordering::Relaxed) } /// Send GET requests with parameters /// * `param` - GET parameters. Should be a JSON object/array. If value in map is not a string, will dump it @@ -150,7 +265,7 @@ impl WebClient { /// client.get_with_param("https://test.com/a", json::array![["daa", "param1"]]); /// ``` /// It will GET `https://test.com/a?data=param1`, `https://test.com/a?daa=%7B%22ad%22%3A%22test%22%7D`, `https://test.com/a?daa=param1` - pub fn get_with_param(&mut self, url: U, param: JsonValue) -> Option { + pub fn get_with_param(&self, url: U, param: JsonValue) -> Option { let u = url.into_url(); if u.is_err() { println!("{} \"{}\"", gettext("Can not parse URL:"), u.unwrap_err()); @@ -211,17 +326,19 @@ impl WebClient { self.get(u.as_str(), None) } - pub fn get(&mut self, url: U, headers: H) -> Option { + pub fn get(&self, url: U, headers: H) -> Option { let mut count = 0u64; - while count <= self.retry { + let retry = self.get_retry(); + while count <= retry { let r = self._get(url.clone(), headers.clone()); if r.is_some() { return r; } count += 1; - if count <= self.retry { - if self.retry_interval.is_some() { - let t = self.retry_interval.as_ref().unwrap()[(count - 1).try_into().unwrap()]; + if count <= retry { + let ri = self.get_retry_interval(); + if ri.is_some() { + let t = ri.as_ref().unwrap()[(count - 1).try_into().unwrap()]; if !t.is_zero() { println!("{}", gettext("Retry after seconds.").replace("", format!("{}", t.as_secs_f64()).as_str()).as_str()); spin_on(tokio::time::sleep(t)); @@ -233,16 +350,17 @@ impl WebClient { None } - pub async fn aget(&mut self, url: U, headers: H) -> Option { + pub async fn aget(&self, url: U, headers: H) -> Option { let mut count = 0u64; - while count <= self.retry { + let retry = self.get_retry(); + while count <= retry { let r = self._aget2(url.clone(), headers.clone()).await; if r.is_some() { return r; } count += 1; - if count <= self.retry { - let t = self.retry_interval.as_ref().unwrap()[(count - 1).try_into().unwrap()]; + if count <= retry { + let t = self.get_retry_interval().as_ref().unwrap()[(count - 1).try_into().unwrap()]; if !t.is_zero() { println!("{}", gettext("Retry after seconds.").replace("", format!("{}", t.as_secs_f64()).as_str()).as_str()); tokio::time::sleep(t).await; @@ -254,7 +372,7 @@ impl WebClient { } /// Send GET requests - pub fn _get(&mut self, url: U, headers: H) -> Option { + pub fn _get(&self, url: U, headers: H) -> Option { let r = self._aget(url, headers); let r = r.send(); let r = spin_on(r); @@ -267,13 +385,13 @@ impl WebClient { } let r = r.unwrap(); self.handle_set_cookie(&r); - if self.verbose { + if self.get_verbose() { println!("{}", r.status()); } Some(r) } - pub async fn _aget2(&mut self, url: U, headers: H) -> Option { + pub async fn _aget2(&self, url: U, headers: H) -> Option { let r = self._aget(url, headers); let r = r.send().await; match r { @@ -285,19 +403,19 @@ impl WebClient { } let r = r.unwrap(); self.handle_set_cookie(&r); - if self.verbose { + if self.get_verbose() { println!("{}", r.status()); } Some(r) } - pub fn _aget(&mut self, url: U, headers: H) -> RequestBuilder { + pub fn _aget(&self, url: U, headers: H) -> RequestBuilder { let s = url.as_str(); - if self.verbose { + if self.get_verbose() { println!("GET {}", s); } let mut r = self.client.get(s); - for (k, v) in self.headers.iter() { + for (k, v) in self.get_headers().iter() { r = r.header(k, v); } let headers = headers.to_headers(); @@ -307,7 +425,7 @@ impl WebClient { r = r.header(k, v); } } - let c = gen_cookie_header(&mut self.cookies, s); + let c = gen_cookie_header(&self, s); if c.len() > 0 { r = r.header("Cookie", c.as_str()); }