mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat: add provider-key concurrent limit (#352)
* feat: add provider-key concurrent limit * Fix provider key concurrent limit checks --------- Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
@@ -263,6 +263,7 @@ pub struct StoredProviderCatalogKey {
|
||||
pub proxy: Option<serde_json::Value>,
|
||||
pub fingerprint: Option<serde_json::Value>,
|
||||
pub rpm_limit: Option<u32>,
|
||||
pub concurrent_limit: Option<i32>,
|
||||
pub learned_rpm_limit: Option<u32>,
|
||||
pub concurrent_429_count: Option<u32>,
|
||||
pub rpm_429_count: Option<u32>,
|
||||
@@ -338,6 +339,7 @@ impl StoredProviderCatalogKey {
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
rpm_limit: None,
|
||||
concurrent_limit: None,
|
||||
learned_rpm_limit: None,
|
||||
concurrent_429_count: None,
|
||||
rpm_429_count: None,
|
||||
@@ -410,6 +412,7 @@ impl StoredProviderCatalogKey {
|
||||
pub fn with_rate_limit_fields(
|
||||
mut self,
|
||||
rpm_limit: Option<u32>,
|
||||
concurrent_limit: Option<i32>,
|
||||
learned_rpm_limit: Option<u32>,
|
||||
concurrent_429_count: Option<u32>,
|
||||
rpm_429_count: Option<u32>,
|
||||
@@ -419,6 +422,7 @@ impl StoredProviderCatalogKey {
|
||||
success_count: Option<u32>,
|
||||
) -> Self {
|
||||
self.rpm_limit = rpm_limit;
|
||||
self.concurrent_limit = concurrent_limit;
|
||||
self.learned_rpm_limit = learned_rpm_limit;
|
||||
self.concurrent_429_count = concurrent_429_count;
|
||||
self.rpm_429_count = rpm_429_count;
|
||||
@@ -460,6 +464,53 @@ impl StoredProviderCatalogKey {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod transport_tests {
|
||||
use super::StoredProviderCatalogKey;
|
||||
|
||||
#[test]
|
||||
fn provider_catalog_key_defaults_concurrent_limit_to_none() {
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
|
||||
assert_eq!(key.concurrent_limit, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_catalog_key_rate_limit_builder_sets_concurrent_limit() {
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_rate_limit_fields(
|
||||
Some(120),
|
||||
Some(3),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(key.rpm_limit, Some(120));
|
||||
assert_eq!(key.concurrent_limit, Some(3));
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub enum ProviderCatalogKeyListOrder {
|
||||
#[default]
|
||||
|
||||
Reference in New Issue
Block a user