i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 02:47:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: update gateway pressure observability

Browse Source
This commit is contained in:
elky
2026-06-30 17:01:39 +08:00
parent 974def5fef
commit f179ee72f9
69 changed files with 14843 additions and 1609 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-testkit/src/bin/gateway_openai_chat_c_compare.rs
+1
View File
@@ -386,6 +386,7 @@ fn load_config(
warmup_url: Some(gateway_health_url.to_string()),
method: Method::POST,
headers,
header_sets: Vec::new(),
body: Some(
serde_json::json!({
"model": "gpt-5",
crates/aether-testkit/src/bin/gateway_pressure_probe.rs
+1990 -15
View File
File diff suppressed because it is too large Load Diff
crates/aether-testkit/src/bin/gateway_pressure_seed.rs
+756
View File
@@ -0,0 +1,756 @@
use std::env;
use std::fs;
use std::path::PathBuf;
use aether_data::repository::auth::CreateStandaloneApiKeyRecord;
use aether_data::repository::wallet::WalletLookupKey;
use aether_data::{
DataBackends, DataLayerConfig, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig,
};
use aether_data_contracts::repository::global_models::{
CreateAdminGlobalModelRecord, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use serde_json::json;
use sha2::Digest;
const DEFAULT_POSTGRES_URL: &str = "postgresql://postgres:[email protected]:5432/aether";
const DEFAULT_OUTPUT_ENV_PATH: &str = "/tmp/aether_local_env.sh";
const DEFAULT_OUTPUT_KEY_PATH: &str = "/tmp/aether_fullchain_api_key";
const DEFAULT_OUTPUT_KEY_LIST_PATH: &str = "/tmp/aether_fullchain_api_keys";
const DEFAULT_PROVIDER_ID: &str = "provider-local-pressure-openai";
const DEFAULT_ENDPOINT_ID: &str = "endpoint-local-pressure-openai-chat";
const DEFAULT_PROVIDER_KEY_ID: &str = "provider-key-local-pressure-openai";
const DEFAULT_GLOBAL_MODEL_ID: &str = "gm-local-pressure-gpt-5-mini";
const DEFAULT_MODEL_ID: &str = "model-local-pressure-gpt-5-mini";
const DEFAULT_API_KEY_ID: &str = "api-key-local-pressure";
const DEFAULT_OPERATOR_ID: &str = "pressure-local";
const DEFAULT_MODEL: &str = "gpt-5-mini";
const DEFAULT_MOCK_UPSTREAM_BASE_URL: &str = "http://127.0.0.1:18181/v1";
const DEFAULT_GATEWAY_BASE_URL: &str = "http://127.0.0.1:8084";
const DEFAULT_API_KEY: &str = "sk-aether-local-pressure";
#[derive(Debug, Clone)]
struct Config {
database_url: String,
output_env_path: PathBuf,
output_key_path: PathBuf,
output_key_list_path: PathBuf,
provider_id: String,
endpoint_id: String,
provider_key_id: String,
global_model_id: String,
model_id: String,
api_key_id: String,
operator_id: String,
model: String,
mock_upstream_base_url: String,
gateway_base_url: String,
api_key: String,
api_key_count: usize,
provider_api_key: String,
postgres_min_connections: u32,
postgres_max_connections: u32,
}
impl Config {
fn from_env_and_args() -> Result<Self, String> {
let mut config = Self {
database_url: env_value("DATABASE_URL")
.or_else(|| env_value("AETHER_DATABASE_URL"))
.or_else(|| env_value("AETHER_GATEWAY_DATA_POSTGRES_URL"))
.unwrap_or_else(|| DEFAULT_POSTGRES_URL.to_string()),
output_env_path: PathBuf::from(
env_value("OUTPUT_ENV_PATH").unwrap_or_else(|| DEFAULT_OUTPUT_ENV_PATH.to_string()),
),
output_key_path: PathBuf::from(
env_value("OUTPUT_KEY_PATH").unwrap_or_else(|| DEFAULT_OUTPUT_KEY_PATH.to_string()),
),
output_key_list_path: PathBuf::from(
env_value("OUTPUT_KEY_LIST_PATH")
.or_else(|| env_value("PRESSURE_API_KEY_LIST_FILE"))
.unwrap_or_else(|| DEFAULT_OUTPUT_KEY_LIST_PATH.to_string()),
),
provider_id: env_value("PRESSURE_PROVIDER_ID")
.unwrap_or_else(|| DEFAULT_PROVIDER_ID.to_string()),
endpoint_id: env_value("PRESSURE_ENDPOINT_ID")
.unwrap_or_else(|| DEFAULT_ENDPOINT_ID.to_string()),
provider_key_id: env_value("PRESSURE_PROVIDER_KEY_ID")
.unwrap_or_else(|| DEFAULT_PROVIDER_KEY_ID.to_string()),
global_model_id: env_value("PRESSURE_GLOBAL_MODEL_ID")
.unwrap_or_else(|| DEFAULT_GLOBAL_MODEL_ID.to_string()),
model_id: env_value("PRESSURE_MODEL_ID")
.unwrap_or_else(|| DEFAULT_MODEL_ID.to_string()),
api_key_id: env_value("PRESSURE_API_KEY_ID")
.unwrap_or_else(|| DEFAULT_API_KEY_ID.to_string()),
operator_id: env_value("PRESSURE_OPERATOR_ID")
.unwrap_or_else(|| DEFAULT_OPERATOR_ID.to_string()),
model: env_value("PRESSURE_MODEL").unwrap_or_else(|| DEFAULT_MODEL.to_string()),
mock_upstream_base_url: env_value("PRESSURE_MOCK_UPSTREAM_BASE_URL")
.unwrap_or_else(|| DEFAULT_MOCK_UPSTREAM_BASE_URL.to_string()),
gateway_base_url: env_value("GATEWAY_BASE_URL")
.unwrap_or_else(|| DEFAULT_GATEWAY_BASE_URL.to_string()),
api_key: env_value("AETHER_API_KEY").unwrap_or_else(|| DEFAULT_API_KEY.to_string()),
api_key_count: env_value("PRESSURE_API_KEY_COUNT")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(1),
provider_api_key: env_value("PRESSURE_PROVIDER_API_KEY")
.unwrap_or_else(|| "dummy-local-pressure-provider-key".to_string()),
postgres_min_connections: 1,
postgres_max_connections: 8,
};
let args = env::args().skip(1).collect::<Vec<_>>();
let mut index = 0;
while index < args.len() {
let arg = &args[index];
match arg.as_str() {
"--database-url" => config.database_url = arg_value(&args, &mut index, arg)?,
"--output-env" => {
config.output_env_path = PathBuf::from(arg_value(&args, &mut index, arg)?)
}
"--output-key" => {
config.output_key_path = PathBuf::from(arg_value(&args, &mut index, arg)?)
}
"--output-key-list" => {
config.output_key_list_path = PathBuf::from(arg_value(&args, &mut index, arg)?)
}
"--provider-id" => config.provider_id = arg_value(&args, &mut index, arg)?,
"--endpoint-id" => config.endpoint_id = arg_value(&args, &mut index, arg)?,
"--provider-key-id" => config.provider_key_id = arg_value(&args, &mut index, arg)?,
"--global-model-id" => config.global_model_id = arg_value(&args, &mut index, arg)?,
"--model-id" => config.model_id = arg_value(&args, &mut index, arg)?,
"--api-key-id" => config.api_key_id = arg_value(&args, &mut index, arg)?,
"--operator-id" => config.operator_id = arg_value(&args, &mut index, arg)?,
"--model" => config.model = arg_value(&args, &mut index, arg)?,
"--mock-upstream-base-url" => {
config.mock_upstream_base_url = arg_value(&args, &mut index, arg)?
}
"--gateway-base-url" => {
config.gateway_base_url = arg_value(&args, &mut index, arg)?
}
"--api-key" => config.api_key = arg_value(&args, &mut index, arg)?,
"--api-key-count" => {
config.api_key_count = parse_usize(&arg_value(&args, &mut index, arg)?, arg)?
}
"--provider-api-key" => {
config.provider_api_key = arg_value(&args, &mut index, arg)?
}
"--postgres-min-connections" => {
config.postgres_min_connections =
parse_u32(&arg_value(&args, &mut index, arg)?, arg)?
}
"--postgres-max-connections" => {
config.postgres_max_connections =
parse_u32(&arg_value(&args, &mut index, arg)?, arg)?
}
"--help" | "-h" => {
print_help();
std::process::exit(0);
}
_ => return Err(format!("unknown argument: {arg}")),
}
index += 1;
}
config.validate()?;
Ok(config)
}
fn validate(&self) -> Result<(), String> {
for (name, value) in [
("database_url", &self.database_url),
("provider_id", &self.provider_id),
("endpoint_id", &self.endpoint_id),
("provider_key_id", &self.provider_key_id),
("global_model_id", &self.global_model_id),
("model_id", &self.model_id),
("api_key_id", &self.api_key_id),
("operator_id", &self.operator_id),
("model", &self.model),
("mock_upstream_base_url", &self.mock_upstream_base_url),
("gateway_base_url", &self.gateway_base_url),
("api_key", &self.api_key),
] {
if value.trim().is_empty() {
return Err(format!("{name} cannot be empty"));
}
}
if self.postgres_min_connections > self.postgres_max_connections {
return Err("postgres min connections cannot exceed max connections".to_string());
}
if self.api_key_count == 0 {
return Err("api_key_count must be positive".to_string());
}
Ok(())
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = Config::from_env_and_args().map_err(|err| format!("invalid config: {err}"))?;
let backends = DataBackends::from_config(DataLayerConfig::from_database(SqlDatabaseConfig {
driver: DatabaseDriver::Postgres,
url: config.database_url.clone(),
pool: SqlPoolConfig {
min_connections: config.postgres_min_connections,
max_connections: config.postgres_max_connections,
acquire_timeout_ms: 5_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 300_000,
statement_cache_capacity: 128,
require_ssl: false,
},
}))?;
seed_provider_catalog(&backends, &config).await?;
seed_models(&backends, &config).await?;
let operator_user_id = seed_operator_user(&backends, &config).await?;
seed_api_keys(&backends, &config, &operator_user_id).await?;
verify_candidate_selection(&backends, &config).await?;
write_outputs(&config)?;
println!("gateway pressure seed complete");
println!("provider_id={}", config.provider_id);
println!("endpoint_id={}", config.endpoint_id);
println!("provider_key_id={}", config.provider_key_id);
println!("model={}", config.model);
println!("api_key_id={}", config.api_key_id);
println!("api_key_count={}", config.api_key_count);
println!("env written to {}", config.output_env_path.display());
println!("api key written to {}", config.output_key_path.display());
println!(
"api key list written to {}",
config.output_key_list_path.display()
);
Ok(())
}
async fn seed_provider_catalog(
backends: &DataBackends,
config: &Config,
) -> Result<(), Box<dyn std::error::Error>> {
let reader = backends
.read()
.provider_catalog()
.ok_or("provider catalog reader unavailable")?;
let writer = backends
.write()
.provider_catalog()
.ok_or("provider catalog writer unavailable")?;
let provider = StoredProviderCatalogProvider::new(
config.provider_id.clone(),
"Local pressure OpenAI mock".to_string(),
Some("http://127.0.0.1:18181".to_string()),
"openai".to_string(),
)?
.with_transport_fields(
true,
false,
false,
None,
Some(0),
None,
Some(120.0),
Some(30.0),
None,
)
.with_routing_fields(0)
.with_description(Some(
"Local OpenAI-compatible mock provider for gateway pressure tests".to_string(),
));
if reader
.list_providers_by_ids(std::slice::from_ref(&config.provider_id))
.await?
.is_empty()
{
writer.create_provider(&provider, None).await?;
} else {
writer.update_provider(&provider).await?;
}
let endpoint = StoredProviderCatalogEndpoint::new(
config.endpoint_id.clone(),
config.provider_id.clone(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat_completions".to_string()),
true,
)?
.with_transport_fields(
config.mock_upstream_base_url.clone(),
None,
None,
Some(0),
None,
None,
None,
None,
)?
.with_health_score(1.0);
if reader
.list_endpoints_by_ids(std::slice::from_ref(&config.endpoint_id))
.await?
.is_empty()
{
writer.create_endpoint(&endpoint).await?;
} else {
writer.update_endpoint(&endpoint).await?;
}
let provider_key = StoredProviderCatalogKey::new(
config.provider_key_id.clone(),
config.provider_id.clone(),
"Local pressure mock key".to_string(),
"api_key".to_string(),
Some(json!({"streaming": true})),
true,
)?
.with_transport_fields(
Some(json!(["openai:chat"])),
Some(config.provider_api_key.clone()),
None,
None,
None,
Some(json!([config.model.clone()])),
None,
None,
None,
)?
.with_rate_limit_fields(None, None, None, None, None, None, None, None, None)
.with_health_fields(
Some(json!({"openai:chat": {"status": "healthy"}})),
Some(json!({"openai:chat": {"state": "closed"}})),
);
if reader
.list_keys_by_ids(std::slice::from_ref(&config.provider_key_id))
.await?
.is_empty()
{
writer.create_key(&provider_key).await?;
} else {
writer.update_key(&provider_key).await?;
}
Ok(())
}
async fn seed_models(
backends: &DataBackends,
config: &Config,
) -> Result<(), Box<dyn std::error::Error>> {
let reader = backends
.read()
.global_models()
.ok_or("global model reader unavailable")?;
let writer = backends
.write()
.global_models()
.ok_or("global model writer unavailable")?;
let capabilities = Some(json!({
"streaming": true,
"chat": true
}));
let global_config = Some(json!({
"model_mappings": [config.model],
"pressure_seed": true
}));
if reader
.get_admin_global_model_by_id(&config.global_model_id)
.await?
.is_some()
{
writer
.update_admin_global_model(&UpdateAdminGlobalModelRecord::new(
config.global_model_id.clone(),
config.model.clone(),
true,
Some(0.0),
None,
capabilities.clone(),
global_config.clone(),
)?)
.await?;
} else {
writer
.create_admin_global_model(&CreateAdminGlobalModelRecord::new(
config.global_model_id.clone(),
config.model.clone(),
config.model.clone(),
true,
Some(0.0),
None,
capabilities.clone(),
global_config.clone(),
)?)
.await?;
}
let provider_model_mappings = Some(json!([
{
"name": config.model,
"priority": 0,
"api_formats": ["openai:chat"],
"endpoint_ids": [config.endpoint_id]
}
]));
let provider_model = UpsertAdminProviderModelRecord::new(
config.model_id.clone(),
config.provider_id.clone(),
config.global_model_id.clone(),
config.model.clone(),
provider_model_mappings,
Some(0.0),
None,
Some(false),
Some(false),
Some(true),
Some(false),
Some(false),
true,
true,
Some(json!({"pressure_seed": true})),
)?;
if reader
.list_admin_provider_models_by_global_model_id(&config.global_model_id)
.await?
.iter()
.any(|model| model.id == config.model_id)
{
writer.update_admin_provider_model(&provider_model).await?;
} else {
writer.create_admin_provider_model(&provider_model).await?;
}
Ok(())
}
async fn seed_api_keys(
backends: &DataBackends,
config: &Config,
operator_user_id: &str,
) -> Result<(), Box<dyn std::error::Error>> {
for index in 0..config.api_key_count {
seed_api_key(backends, config, operator_user_id, index).await?;
}
Ok(())
}
async fn seed_api_key(
backends: &DataBackends,
config: &Config,
operator_user_id: &str,
key_index: usize,
) -> Result<(), Box<dyn std::error::Error>> {
let auth_reader = backends
.read()
.auth_api_keys()
.ok_or("auth api key reader unavailable")?;
let auth_writer = backends
.write()
.auth_api_keys()
.ok_or("auth api key writer unavailable")?;
let wallet_reader = backends
.read()
.wallets()
.ok_or("wallet reader unavailable")?;
let api_key_id = pressure_api_key_id(config, key_index);
let api_key_value = pressure_api_key_value(config, key_index);
let existing = auth_reader
.find_export_standalone_api_key_by_id(&api_key_id)
.await?;
if existing.is_none() {
auth_writer
.create_standalone_api_key(CreateStandaloneApiKeyRecord {
user_id: operator_user_id.to_string(),
api_key_id: api_key_id.clone(),
key_hash: sha256_hex(&api_key_value),
key_encrypted: Some(api_key_value),
name: Some(format!("Local pressure API key {}", key_index + 1)),
allowed_providers: Some(vec![config.provider_id.clone()]),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_models: Some(vec![config.model.clone()]),
ip_rules: None,
rate_limit: Some(0),
concurrent_limit: None,
force_capabilities: None,
is_active: true,
expires_at_unix_secs: None,
auto_delete_on_expiry: false,
total_requests: 0,
total_tokens: 0,
total_cost_usd: 0.0,
})
.await?;
} else {
auth_writer
.update_standalone_api_key_basic(
aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord {
api_key_id: api_key_id.clone(),
name: Some(format!("Local pressure API key {}", key_index + 1)),
rate_limit_present: true,
rate_limit: Some(0),
concurrent_limit_present: true,
concurrent_limit: None,
allowed_providers: Some(Some(vec![config.provider_id.clone()])),
allowed_api_formats: Some(Some(vec!["openai:chat".to_string()])),
allowed_models: Some(Some(vec![config.model.clone()])),
ip_rules: Some(None),
expires_at_present: true,
expires_at_unix_secs: None,
auto_delete_on_expiry_present: true,
auto_delete_on_expiry: false,
},
)
.await?;
auth_writer
.set_standalone_api_key_active(&api_key_id, true)
.await?;
}
if wallet_reader
.find(WalletLookupKey::ApiKeyId(&api_key_id))
.await?
.is_none()
{
wallet_reader
.initialize_auth_api_key_wallet(&api_key_id, 0.0, true)
.await?;
} else {
wallet_reader
.update_auth_api_key_wallet_limit_mode(&api_key_id, "unlimited")
.await?;
}
Ok(())
}
async fn seed_operator_user(
backends: &DataBackends,
config: &Config,
) -> Result<String, Box<dyn std::error::Error>> {
let user_reader = backends.read().users().ok_or("user reader unavailable")?;
let wallet_reader = backends
.read()
.wallets()
.ok_or("wallet reader unavailable")?;
let username = format!("{}-user", config.operator_id);
let user = match user_reader.find_user_auth_by_username(&username).await? {
Some(user) => user_reader
.update_local_auth_user_admin_fields(
&user.id,
Some("admin".to_string()),
true,
Some(vec![config.provider_id.clone()]),
true,
Some(vec!["openai:chat".to_string()]),
true,
Some(vec![config.model.clone()]),
true,
None,
Some(true),
)
.await?
.unwrap_or(user),
None => user_reader
.create_local_auth_user_with_settings(
Some(format!("{}@local.pressure", config.operator_id)),
true,
username,
"local-pressure-password-disabled".to_string(),
"admin".to_string(),
Some(vec![config.provider_id.clone()]),
Some(vec!["openai:chat".to_string()]),
Some(vec![config.model.clone()]),
None,
)
.await?
.ok_or("failed to create pressure operator user")?,
};
if wallet_reader
.find(WalletLookupKey::UserId(&user.id))
.await?
.is_none()
{
wallet_reader
.initialize_auth_user_wallet(&user.id, 0.0, true)
.await?;
} else {
wallet_reader
.update_auth_user_wallet_limit_mode(&user.id, "unlimited")
.await?;
}
Ok(user.id)
}
async fn verify_candidate_selection(
backends: &DataBackends,
config: &Config,
) -> Result<(), Box<dyn std::error::Error>> {
let reader = backends
.read()
.minimal_candidate_selection()
.ok_or("candidate selection reader unavailable")?;
let rows = reader
.list_for_exact_api_format_and_requested_model("openai:chat", &config.model)
.await?;
let has_pressure_row = rows.iter().any(|row| {
row.provider_id == config.provider_id
&& row.endpoint_id == config.endpoint_id
&& row.key_id == config.provider_key_id
&& row.global_model_id == config.global_model_id
&& row.model_id == config.model_id
&& row.provider_is_active
&& row.endpoint_is_active
&& row.key_is_active
&& row.model_is_active
&& row.model_is_available
});
if !has_pressure_row {
return Err(format!(
"seeded candidate not visible for model {} and openai:chat",
config.model
)
.into());
}
Ok(())
}
fn write_outputs(config: &Config) -> Result<(), Box<dyn std::error::Error>> {
if let Some(parent) = config.output_env_path.parent() {
fs::create_dir_all(parent)?;
}
if let Some(parent) = config.output_key_path.parent() {
fs::create_dir_all(parent)?;
}
if let Some(parent) = config.output_key_list_path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(&config.output_key_path, format!("{}\n", config.api_key))?;
let key_list = (0..config.api_key_count)
.map(|index| pressure_api_key_value(config, index))
.collect::<Vec<_>>()
.join("\n");
fs::write(&config.output_key_list_path, format!("{key_list}\n"))?;
let env_content = format!(
concat!(
"export AETHER_API_KEY_FILE={key_path}\n",
"export AETHER_API_KEY_LIST_FILE={key_list_path}\n",
"export AETHER_API_KEY=$(cat {key_path})\n",
"export GATEWAY_BASE_URL={gateway_base_url}\n",
"export TARGET_URL={gateway_base_url}/v1/chat/completions\n",
"export METRICS_URL={gateway_base_url}/_gateway/metrics\n",
"export PRESSURE_MODEL={model}\n",
"export PRESSURE_MOCK_UPSTREAM_BASE_URL={mock_upstream_base_url}\n",
"export PRESSURE_MOCK_UPSTREAM_METRICS_URL=http://127.0.0.1:18181/metrics\n"
),
key_path = shell_escape(&config.output_key_path.display().to_string()),
key_list_path = shell_escape(&config.output_key_list_path.display().to_string()),
gateway_base_url = shell_escape(config.gateway_base_url.trim_end_matches('/')),
model = shell_escape(&config.model),
mock_upstream_base_url = shell_escape(&config.mock_upstream_base_url),
);
fs::write(&config.output_env_path, env_content)?;
Ok(())
}
fn sha256_hex(value: &str) -> String {
let mut hasher = sha2::Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn pressure_api_key_id(config: &Config, index: usize) -> String {
if index == 0 {
config.api_key_id.clone()
} else {
format!("{}-{}", config.api_key_id, index + 1)
}
}
fn pressure_api_key_value(config: &Config, index: usize) -> String {
if index == 0 {
config.api_key.clone()
} else {
format!("{}-{}", config.api_key, index + 1)
}
}
fn shell_escape(value: &str) -> String {
if value
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '/' | '.' | ':' | '_' | '-'))
{
return value.to_string();
}
format!("'{}'", value.replace('\'', "'\\''"))
}
fn env_value(name: &str) -> Option<String> {
env::var(name)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn arg_value(args: &[String], index: &mut usize, name: &str) -> Result<String, String> {
*index += 1;
args.get(*index)
.filter(|value| !value.starts_with("--"))
.cloned()
.ok_or_else(|| format!("{name} requires a value"))
}
fn parse_u32(value: &str, name: &str) -> Result<u32, String> {
value
.parse::<u32>()
.map_err(|_| format!("{name} must be an unsigned integer"))
}
fn parse_usize(value: &str, name: &str) -> Result<usize, String> {
value
.parse::<usize>()
.map_err(|_| format!("{name} must be an unsigned integer"))
}
fn print_help() {
println!(
"Usage: cargo run -p aether-testkit --bin gateway_pressure_seed -- [options]\n\
\n\
Options:\n\
--database-url URL\n\
--output-env PATH\n\
--output-key PATH\n\
--output-key-list PATH\n\
--gateway-base-url URL\n\
--mock-upstream-base-url URL\n\
--model NAME\n\
--api-key VALUE\n\
--api-key-count N\n\
--provider-api-key VALUE\n\
--provider-id ID\n\
--endpoint-id ID\n\
--provider-key-id ID\n\
--global-model-id ID\n\
--model-id ID\n\
--api-key-id ID\n"
);
}
crates/aether-testkit/src/load.rs
+125 -8
View File
@@ -11,6 +11,8 @@ use tokio::sync::Mutex;
use crate::runtime::{BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot};
const MAX_ERROR_SAMPLES: usize = 32;
const MAX_STATUS_SAMPLES: usize = 32;
const MAX_STATUS_SAMPLE_BODY_CHARS: usize = 512;
const FIRST_BODY_BACKGROUND_DRAIN_CHUNKS_ENV: &str =
"AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS";
const FIRST_BODY_BACKGROUND_DRAIN_MS_ENV: &str = "AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_MS";
@@ -34,6 +36,7 @@ pub struct HttpLoadProbeConfig {
pub warmup_url: Option<String>,
pub method: Method,
pub headers: BTreeMap<String, String>,
pub header_sets: Vec<BTreeMap<String, String>>,
pub body: Option<Vec<u8>>,
pub total_requests: usize,
pub concurrency: usize,
@@ -56,6 +59,7 @@ impl Default for HttpLoadProbeConfig {
warmup_url: None,
method: Method::GET,
headers: BTreeMap::new(),
header_sets: Vec::new(),
body: None,
total_requests: 100,
concurrency: 10,
@@ -93,6 +97,11 @@ impl HttpLoadProbeConfig {
if self.client_shards == 0 {
return Err("load probe client_shards must be positive".to_string());
}
for (index, headers) in self.header_sets.iter().enumerate() {
if headers.is_empty() {
return Err(format!("load probe header_sets[{index}] cannot be empty"));
}
}
if self.http1_only && self.http2_prior_knowledge {
return Err(
"load probe cannot enable both http1_only and http2_prior_knowledge".to_string(),
@@ -114,6 +123,15 @@ pub struct HttpLoadProbeErrorSample {
pub source: Option<String>,
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
pub struct HttpLoadProbeStatusSample {
pub request_index: usize,
pub url: String,
pub status: u16,
pub elapsed_ms: u64,
pub body: String,
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
pub struct HttpLoadProbeResult {
pub url: String,
@@ -159,6 +177,8 @@ pub struct HttpLoadProbeResult {
pub error_counts: BTreeMap<String, usize>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub error_samples: Vec<HttpLoadProbeErrorSample>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub non_success_status_samples: Vec<HttpLoadProbeStatusSample>,
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
@@ -207,6 +227,8 @@ pub struct MultiUrlHttpLoadProbeResult {
pub error_counts: BTreeMap<String, usize>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub error_samples: Vec<HttpLoadProbeErrorSample>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub non_success_status_samples: Vec<HttpLoadProbeStatusSample>,
}
pub async fn run_http_load_probe(
@@ -253,6 +275,7 @@ pub async fn run_http_load_probe(
status_counts: result.status_counts,
error_counts: result.error_counts,
error_samples: result.error_samples,
non_success_status_samples: result.non_success_status_samples,
})
}
@@ -274,7 +297,7 @@ async fn run_http_load_probe_against_urls(
let effective_client_shards = effective_probe_client_shards(config);
let clients = Arc::new(build_probe_clients(config, effective_client_shards)?);
let total_requests = config.total_requests;
let request_headers = build_headers(&config.headers)?;
let request_headers = build_request_header_sets(config)?;
let request_body = config.body.clone().map(Arc::new);
let response_mode = config.response_mode;
let first_body_hold = config.first_body_hold;
@@ -290,6 +313,7 @@ async fn run_http_load_probe_against_urls(
let status_counts = Arc::new(Mutex::new(BTreeMap::<u16, usize>::new()));
let error_counts = Arc::new(Mutex::new(BTreeMap::<String, usize>::new()));
let error_samples = Arc::new(Mutex::new(Vec::<HttpLoadProbeErrorSample>::new()));
let non_success_status_samples = Arc::new(Mutex::new(Vec::<HttpLoadProbeStatusSample>::new()));
let target_request_counts = Arc::new(Mutex::new(BTreeMap::<String, usize>::new()));
let failed_requests = Arc::new(AtomicUsize::new(0));
let completed_requests = Arc::new(AtomicUsize::new(0));
@@ -304,12 +328,13 @@ async fn run_http_load_probe_against_urls(
let status_counts = Arc::clone(&status_counts);
let error_counts = Arc::clone(&error_counts);
let error_samples = Arc::clone(&error_samples);
let non_success_status_samples = Arc::clone(&non_success_status_samples);
let target_request_counts = Arc::clone(&target_request_counts);
let failed_requests = Arc::clone(&failed_requests);
let completed_requests = Arc::clone(&completed_requests);
let method = config.method.clone();
let urls = urls.to_vec();
let request_headers = request_headers.clone();
let request_headers = Arc::clone(&request_headers);
let request_body = request_body.clone();
let start_delay = worker_start_delay(start_ramp, worker_index, config.concurrency);
@@ -326,7 +351,8 @@ async fn run_http_load_probe_against_urls(
let started_at = Instant::now();
let url = urls[current % urls.len()].clone();
let mut request = client.request(method.clone(), &url);
for (name, value) in request_headers.iter() {
let headers = &request_headers[current % request_headers.len()];
for (name, value) in headers.iter() {
request = request.header(name, value);
}
if let Some(body) = request_body.as_ref() {
@@ -357,6 +383,17 @@ async fn run_http_load_probe_against_urls(
.await;
}
Ok(observation) => {
if !(200..300).contains(&status) {
record_non_success_status_sample(
&non_success_status_samples,
current,
&url,
status,
started_at.elapsed().as_millis() as u64,
observation.body_sample.as_deref().unwrap_or_default(),
)
.await;
}
let mut counts = status_counts.lock().await;
*counts.entry(status).or_insert(0) += 1;
drop(counts);
@@ -404,6 +441,7 @@ async fn run_http_load_probe_against_urls(
let status_counts = status_counts.lock().await.clone();
let error_counts = error_counts.lock().await.clone();
let error_samples = error_samples.lock().await.clone();
let non_success_status_samples = non_success_status_samples.lock().await.clone();
let target_request_counts = target_request_counts.lock().await.clone();
let mut latencies = latencies_ms.lock().await.clone();
let mut header_latencies = header_latencies_ms.lock().await.clone();
@@ -460,6 +498,7 @@ async fn run_http_load_probe_against_urls(
status_counts,
error_counts,
error_samples,
non_success_status_samples,
})
}
@@ -613,6 +652,26 @@ async fn record_load_error(
}
}
async fn record_non_success_status_sample(
non_success_status_samples: &Arc<Mutex<Vec<HttpLoadProbeStatusSample>>>,
request_index: usize,
url: &str,
status: u16,
elapsed_ms: u64,
body: &str,
) {
let mut samples = non_success_status_samples.lock().await;
if samples.len() < MAX_STATUS_SAMPLES {
samples.push(HttpLoadProbeStatusSample {
request_index,
url: url.to_string(),
status,
elapsed_ms,
body: compact_error_text(body, MAX_STATUS_SAMPLE_BODY_CHARS),
});
}
}
fn classify_reqwest_error(phase: &str, err: &reqwest::Error) -> ClassifiedLoadError {
let kind = if err.is_timeout() && err.is_connect() {
"connect_timeout"
@@ -689,17 +748,19 @@ async fn observe_response_body(
)
})?;
let first_body_latency_ms = started_at.elapsed().as_millis() as u64;
drop(first);
let body_sample = compact_bytes_sample(&first, MAX_STATUS_SAMPLE_BODY_CHARS);
if !first_body_hold.is_zero() {
tokio::time::sleep(first_body_hold).await;
}
drain_first_body_response_tail(response).await?;
Ok(BodyObservation {
first_body_latency_ms: Some(first_body_latency_ms),
body_sample: Some(body_sample),
})
}
HttpLoadProbeResponseMode::FullBody => {
let mut first_body_latency_ms = None;
let mut body_sample = Vec::new();
while let Some(chunk) = response
.chunk()
.await
@@ -708,7 +769,7 @@ async fn observe_response_body(
if first_body_latency_ms.is_none() {
first_body_latency_ms = Some(started_at.elapsed().as_millis() as u64);
}
drop(chunk);
append_body_sample(&mut body_sample, &chunk, MAX_STATUS_SAMPLE_BODY_CHARS);
}
if first_body_latency_ms.is_none() {
return Err(ClassifiedLoadError::static_body(
@@ -718,11 +779,24 @@ async fn observe_response_body(
}
Ok(BodyObservation {
first_body_latency_ms,
body_sample: Some(String::from_utf8_lossy(&body_sample).into_owned()),
})
}
}
}
fn compact_bytes_sample(bytes: &[u8], max_chars: usize) -> String {
compact_error_text(String::from_utf8_lossy(bytes), max_chars)
}
fn append_body_sample(target: &mut Vec<u8>, chunk: &[u8], max_chars: usize) {
if target.len() >= max_chars {
return;
}
let remaining = max_chars.saturating_sub(target.len());
target.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
}
async fn drain_first_body_response_tail(
mut response: reqwest::Response,
) -> Result<(), ClassifiedLoadError> {
@@ -773,9 +847,26 @@ fn env_u64(key: &str, default_value: u64) -> u64 {
.unwrap_or(default_value)
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct BodyObservation {
first_body_latency_ms: Option<u64>,
body_sample: Option<String>,
}
fn build_request_header_sets(config: &HttpLoadProbeConfig) -> Result<Arc<Vec<HeaderMap>>, String> {
let raw_sets = if config.header_sets.is_empty() {
vec![config.headers.clone()]
} else {
config.header_sets.clone()
};
let mut sets = Vec::with_capacity(raw_sets.len().max(1));
for headers in raw_sets {
sets.push(build_headers(&headers)?);
}
if sets.is_empty() {
sets.push(HeaderMap::new());
}
Ok(Arc::new(sets))
}
fn build_headers(headers: &BTreeMap<String, String>) -> Result<HeaderMap, String> {
@@ -815,8 +906,8 @@ fn percentile(latencies: &[u64], percentile: u8) -> u64 {
#[cfg(test)]
mod tests {
use super::{
build_headers, summarize_latencies, worker_start_delay, HttpLoadProbeConfig,
HttpLoadProbeResponseMode,
build_headers, build_request_header_sets, summarize_latencies, worker_start_delay,
HttpLoadProbeConfig, HttpLoadProbeResponseMode,
};
use reqwest::Method;
use std::collections::BTreeMap;
@@ -886,6 +977,7 @@ mod tests {
assert_eq!(config.method, Method::GET);
assert!(config.warmup_url.is_none());
assert!(config.headers.is_empty());
assert!(config.header_sets.is_empty());
assert!(config.body.is_none());
assert_eq!(config.total_requests, 100);
assert_eq!(config.concurrency, 10);
@@ -916,6 +1008,31 @@ mod tests {
assert!(build_headers(&invalid).is_err());
}
#[test]
fn validates_header_sets() {
let mut config = HttpLoadProbeConfig {
url: "http://127.0.0.1/".to_string(),
header_sets: vec![BTreeMap::new()],
..HttpLoadProbeConfig::default()
};
assert!(config.validate().is_err());
config.header_sets = vec![BTreeMap::from([(
"authorization".to_string(),
"Bearer test".to_string(),
)])];
assert!(config.validate().is_ok());
let sets = build_request_header_sets(&config).expect("header set should build");
assert_eq!(sets.len(), 1);
assert_eq!(
sets[0]
.get("authorization")
.and_then(|value| value.to_str().ok()),
Some("Bearer test")
);
}
#[test]
fn spreads_worker_start_delay_across_ramp() {
assert_eq!(
crates/aether-testkit/src/metrics.rs
+88 -7
View File
@@ -85,22 +85,69 @@ fn parse_prometheus_line(line: &str) -> Option<PrometheusSample> {
fn parse_labels(raw: &str) -> BTreeMap<String, String> {
let mut labels = BTreeMap::new();
for pair in raw.split(',').filter(|pair| !pair.is_empty()) {
for pair in split_label_pairs(raw) {
if let Some((key, value)) = pair.split_once('=') {
labels.insert(
key.trim().to_string(),
value
.trim()
.trim_matches('"')
.replace("\\\"", "\"")
.replace("\\n", "\n")
.replace("\\\\", "\\"),
unescape_label_value(value.trim().trim_matches('"')),
);
}
}
labels
}
fn split_label_pairs(raw: &str) -> Vec<&str> {
let mut pairs = Vec::new();
let mut start = 0;
let mut in_string = false;
let mut escaped = false;
for (index, ch) in raw.char_indices() {
if escaped {
escaped = false;
continue;
}
match ch {
'\\' if in_string => escaped = true,
'"' => in_string = !in_string,
',' if !in_string => {
let pair = raw[start..index].trim();
if !pair.is_empty() {
pairs.push(pair);
}
start = index + ch.len_utf8();
}
_ => {}
}
}
let pair = raw[start..].trim();
if !pair.is_empty() {
pairs.push(pair);
}
pairs
}
fn unescape_label_value(value: &str) -> String {
let mut output = String::with_capacity(value.len());
let mut chars = value.chars();
while let Some(ch) = chars.next() {
if ch != '\\' {
output.push(ch);
continue;
}
match chars.next() {
Some('n') => output.push('\n'),
Some('\\') => output.push('\\'),
Some('"') => output.push('"'),
Some(next) => {
output.push('\\');
output.push(next);
}
None => output.push('\\'),
}
}
output
}
#[cfg(test)]
mod tests {
use super::{find_metric_value_u64, parse_prometheus_samples};
@@ -133,4 +180,38 @@ aether_gateway_concurrency_rejected_total{gate="gateway_requests"} 12
Some(12)
);
}
#[test]
fn parses_quoted_label_values_containing_commas() {
let samples = parse_prometheus_samples(
r#"
metric_with_sql{rank="1",query_prefix="SELECT id, name, created_at FROM request_candidates",state="active"} 2
"#,
);
assert_eq!(samples.len(), 1);
assert_eq!(samples[0].labels.get("rank").map(String::as_str), Some("1"));
assert_eq!(
samples[0].labels.get("query_prefix").map(String::as_str),
Some("SELECT id, name, created_at FROM request_candidates")
);
assert_eq!(
samples[0].labels.get("state").map(String::as_str),
Some("active")
);
}
#[test]
fn unescapes_quoted_label_values() {
let samples = parse_prometheus_samples(
r#"
metric_with_escape{message="bad\"line\nx\\y"} 1
"#,
);
assert_eq!(
samples[0].labels.get("message").map(String::as_str),
Some("bad\"line\nx\\y")
);
}
}