From 44c5df9a7f0f2be4f94681bd2bc49d02b36ea2c5 Mon Sep 17 00:00:00 2001 From: pyrite Date: Sun, 10 Aug 2025 18:53:57 +0200 Subject: [PATCH] refactor(key-pool): change error handler signature --- Cargo.lock | 2 +- torn-key-pool/Cargo.toml | 2 +- torn-key-pool/src/lib.rs | 30 ++++++++++++++++++------------ torn-key-pool/src/postgres.rs | 14 +++++++------- 4 files changed, 27 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dbc87a5..42deb4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2329,7 +2329,7 @@ dependencies = [ [[package]] name = "torn-key-pool" -version = "1.1.3" +version = "1.2.0" dependencies = [ "chrono", "futures", diff --git a/torn-key-pool/Cargo.toml b/torn-key-pool/Cargo.toml index 2bf3edd..9d12f0f 100644 --- a/torn-key-pool/Cargo.toml +++ b/torn-key-pool/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "torn-key-pool" -version = "1.1.3" +version = "1.2.0" edition = "2021" authors = ["Pyrit [2111649]"] license = { workspace = true } diff --git a/torn-key-pool/src/lib.rs b/torn-key-pool/src/lib.rs index 4f7d8cb..3d84bfc 100644 --- a/torn-key-pool/src/lib.rs +++ b/torn-key-pool/src/lib.rs @@ -245,7 +245,11 @@ where error_hooks: HashMap< u16, Box< - dyn for<'a> Fn(&'a S, &'a S::Key) -> BoxFuture<'a, Result> + dyn for<'a> Fn( + &'a S, + &'a S::Key, + &'a ApiRequest, + ) -> BoxFuture<'a, Result> + Send + Sync, >, @@ -287,27 +291,29 @@ where self } - pub fn error_hook(mut self, code: u16, handler: F) -> Self + pub fn error_hook(mut self, error: ApiError, handler: F) -> Self where - F: for<'a> Fn(&'a S, &'a S::Key) -> BoxFuture<'a, Result> + F: for<'a> Fn(&'a S, &'a S::Key, &'a ApiRequest) -> BoxFuture<'a, Result> + Send + Sync + 'static, { - self.options.error_hooks.insert(code, Box::new(handler)); + self.options + .error_hooks + .insert(error.code(), Box::new(handler)); self } pub fn use_default_hooks(self) -> Self { - self.error_hook(2, |storage, key| { + self.error_hook(ApiError::IncorrectKey, |storage, key, _| { async move { storage.remove_key(KeySelector::Id(key.id())).await?; Ok(true) } .boxed() }) - .error_hook(5, |storage, key| { + .error_hook(ApiError::TooManyRequest, |storage, key, _| { async move { storage .timeout_key(KeySelector::Id(key.id()), Duration::from_secs(60)) @@ -316,14 +322,14 @@ where } .boxed() }) - .error_hook(10, |storage, key| { + .error_hook(ApiError::KeyOwnerInFederalJail, |storage, key, _| { async move { storage.remove_key(KeySelector::Id(key.id())).await?; Ok(true) } .boxed() }) - .error_hook(13, |storage, key| { + .error_hook(ApiError::TemporaryInactivity, |storage, key, _| { async move { storage .timeout_key(KeySelector::Id(key.id()), Duration::from_secs(24 * 3_600)) @@ -332,7 +338,7 @@ where } .boxed() }) - .error_hook(18, |storage, key| { + .error_hook(ApiError::Paused, |storage, key, _| { async move { storage .timeout_key(KeySelector::Id(key.id()), Duration::from_secs(24 * 3_600)) @@ -391,7 +397,7 @@ where if let Some(err) = decode_error(&bytes)? { if let Some(handler) = self.options.error_hooks.get(&err.code()) { - let retry = (*handler)(&self.storage, key).await?; + let retry = (*handler)(&self.storage, key, request).await?; if retry { return Ok(RequestResult::Retry); @@ -492,7 +498,7 @@ impl KeyPool where S: KeyPoolStorage + Send + Sync + 'static, { - pub fn torn_api(&self, selector: I) -> KeyPoolExecutor + pub fn torn_api(&self, selector: I) -> KeyPoolExecutor<'_, S> where I: IntoSelector, { @@ -503,7 +509,7 @@ where &self, selector: I, distance: Duration, - ) -> ThrottledKeyPoolExecutor + ) -> ThrottledKeyPoolExecutor<'_, S> where I: IntoSelector, { diff --git a/torn-key-pool/src/postgres.rs b/torn-key-pool/src/postgres.rs index aa142e4..1ac1e88 100644 --- a/torn-key-pool/src/postgres.rs +++ b/torn-key-pool/src/postgres.rs @@ -155,7 +155,7 @@ where pub async fn initialise(&self) -> Result<(), PgKeyPoolError> { if let Some(schema) = self.schema.as_ref() { - sqlx::query(&format!("create schema if not exists {}", schema)) + sqlx::query(&format!("create schema if not exists {schema}")) .execute(&self.pool) .await?; } @@ -306,7 +306,7 @@ where fn recurse( storage: &PgKeyPoolStorage, selector: KeySelector, D>, - ) -> BoxFuture, PgKeyPoolError>> + ) -> BoxFuture<'_, Result, PgKeyPoolError>> where D: PgKeyDomain, { @@ -445,7 +445,7 @@ where storage: &PgKeyPoolStorage, selector: KeySelector, D>, number: i64, - ) -> BoxFuture>, PgKeyPoolError>> + ) -> BoxFuture<'_, Result>, PgKeyPoolError>> where D: PgKeyDomain, { @@ -686,7 +686,7 @@ pub(crate) mod test { let (storage, _) = setup(pool).await; if let Err(e) = storage.initialise().await { - panic!("Initialising key storage failed: {:?}", e); + panic!("Initialising key storage failed: {e:?}"); } } @@ -815,7 +815,7 @@ pub(crate) mod test { let (storage, _) = setup(pool).await; if let Err(e) = storage.acquire_key(Domain::All).await { - panic!("Acquiring key failed: {:?}", e); + panic!("Acquiring key failed: {e:?}"); } } @@ -843,7 +843,7 @@ pub(crate) mod test { let (storage, _) = setup(pool).await; match storage.acquire_many_keys(Domain::All, 30).await { - Err(e) => panic!("Acquiring key failed: {:?}", e), + Err(e) => panic!("Acquiring key failed: {e:?}"), Ok(keys) => assert_eq!(keys.len(), 30), } } @@ -888,7 +888,7 @@ pub(crate) mod test { for i in 0..24 { storage - .store_key(1, format!("{}", i), vec![Domain::All]) + .store_key(1, format!("{i}"), vec![Domain::All]) .await .unwrap(); }