refactor(key-pool): change error handler signature

This commit is contained in:
pyrite 2025-08-10 18:53:57 +02:00
parent 8a8b34506a
commit 44c5df9a7f
Signed by: pyrite
GPG key ID: 7F1BA9170CD35D15
4 changed files with 27 additions and 21 deletions

View file

@ -245,7 +245,11 @@ where
error_hooks: HashMap<
u16,
Box<
dyn for<'a> Fn(&'a S, &'a S::Key) -> BoxFuture<'a, Result<bool, S::Error>>
dyn for<'a> Fn(
&'a S,
&'a S::Key,
&'a ApiRequest,
) -> BoxFuture<'a, Result<bool, S::Error>>
+ Send
+ Sync,
>,
@ -287,27 +291,29 @@ where
self
}
pub fn error_hook<F>(mut self, code: u16, handler: F) -> Self
pub fn error_hook<F>(mut self, error: ApiError, handler: F) -> Self
where
F: for<'a> Fn(&'a S, &'a S::Key) -> BoxFuture<'a, Result<bool, S::Error>>
F: for<'a> Fn(&'a S, &'a S::Key, &'a ApiRequest) -> BoxFuture<'a, Result<bool, S::Error>>
+ 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<S> KeyPool<S>
where
S: KeyPoolStorage + Send + Sync + 'static,
{
pub fn torn_api<I>(&self, selector: I) -> KeyPoolExecutor<S>
pub fn torn_api<I>(&self, selector: I) -> KeyPoolExecutor<'_, S>
where
I: IntoSelector<S::Key, S::Domain>,
{
@ -503,7 +509,7 @@ where
&self,
selector: I,
distance: Duration,
) -> ThrottledKeyPoolExecutor<S>
) -> ThrottledKeyPoolExecutor<'_, S>
where
I: IntoSelector<S::Key, S::Domain>,
{