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

fix: 修复手动代理节点请求不计数及延迟/心跳不显示的问题 (#368)

Browse Source 
* fix: 修复手动代理节点请求不计数及延迟/心跳不显示的问题

问题描述:
- 手动添加的代理节点请求数始终为0，不会递增
- 手动代理节点的延迟和最后心跳时间不显示

根因:
Rust 重写版中缺失了 Python 版的手动代理节点请求计数逻辑。
隧道节点通过心跳上报计数，但手动节点没有心跳机制，
需要在 usage recording 路径中递增计数并更新延迟信息。

修复方案:
1. 在 ExecutionPlan 中提取 proxy 信息，注入到 request_metadata
2. 在 request_metadata 白名单中添加 proxy 字段
3. 新增 INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL，
   递增 total_requests/failed_requests，
   同时更新 avg_latency_ms 和 last_heartbeat_at
4. 在 ProxyNodeWriteRepository trait 新增 increment_manual_node_requests 方法
5. 在 write_event_record 中解析 request_metadata 中的 proxy 信息，
   对非隧道模式的手动节点调用递增方法
6. 为 GatewayDataState 实现 ManualProxyNodeCounter trait

影响范围:
- 仅影响手动代理节点的统计数据
- 隧道节点不受影响（继续通过心跳计数）
- 不影响请求转发逻辑

* style: cargo fmt

---------

Co-authored-by: root <root@ser406777952330.local>
This commit is contained in:
jiuwovo-ai
2026-05-03 14:50:36 +08:00
committed by GitHub
parent c4ea042eb4
commit 3e2eca4fd0
9 changed files with 195 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
apps/aether-gateway/src/data/state/integrations.rs
View File

@@ -248,6 +248,26 @@ impl UsageRuntimeAccess for GatewayDataState {
}
}
#[async_trait]
impl aether_usage_runtime::ManualProxyNodeCounter for GatewayDataState {
async fn increment_manual_proxy_node_requests(
&self,
node_id: &str,
total_delta: i64,
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), DataLayerError> {
match &self.proxy_node_writer {
Some(repository) => {
repository
.increment_manual_node_requests(node_id, total_delta, failed_delta, latency_ms)
.await
}
None => Ok(()),
}
}
}
#[async_trait]
impl UsageRecordWriter for GatewayDataState {
async fn upsert_usage_record(

26
crates/aether-data/src/repository/proxy_nodes/memory.rs
View File

@@ -552,6 +552,32 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
node.updated_at_unix_secs = Self::now_unix_secs();
Ok(Some(node.clone()))
}
async fn increment_manual_node_requests(
&self,
node_id: &str,
total_delta: i64,
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), DataLayerError> {
let mut nodes = self.nodes.write().expect("proxy node repository lock");
let Some(node) = nodes.get_mut(node_id) else {
return Ok(());
};
if !node.is_manual {
return Ok(());
}
if total_delta > 0 {
node.total_requests += total_delta;
}
if failed_delta > 0 {
node.failed_requests += failed_delta;
}
if let Some(ms) = latency_ms {
node.avg_latency_ms = Some(ms as f64);
}
Ok(())
}
}
#[cfg(test)]

30
crates/aether-data/src/repository/proxy_nodes/sql.rs
View File

@@ -357,6 +357,18 @@ WHERE is_manual = FALSE
AND tunnel_connected = TRUE
"#;
const INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL: &str = r#"
UPDATE proxy_nodes
SET
total_requests = total_requests + GREATEST($1::bigint, 0),
failed_requests = failed_requests + GREATEST($2::bigint, 0),
avg_latency_ms = COALESCE($3, avg_latency_ms),
last_heartbeat_at = NOW(),
updated_at = NOW()
WHERE id = $4
AND is_manual = TRUE
"#;
#[derive(Debug, Clone)]
pub struct SqlxProxyNodeRepository {
pool: PgPool,
@@ -986,6 +998,24 @@ VALUES (
self.find_proxy_node(&mutation.node_id).await
}
async fn increment_manual_node_requests(
&self,
node_id: &str,
total_delta: i64,
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), DataLayerError> {
sqlx::query(INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL)
.bind(total_delta)
.bind(failed_delta)
.bind(latency_ms)
.bind(node_id)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(())
}
}
#[cfg(test)]

8
crates/aether-data/src/repository/proxy_nodes/types.rs
View File

@@ -411,6 +411,14 @@ pub trait ProxyNodeWriteRepository: Send + Sync {
&self,
mutation: &ProxyNodeRemoteConfigMutation,
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
async fn increment_manual_node_requests(
&self,
node_id: &str,
total_delta: i64,
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), crate::DataLayerError>;
}
#[cfg(test)]

4
crates/aether-usage-runtime/src/lib.rs
View File

@@ -44,8 +44,8 @@ pub use settlement::{settle_usage_if_needed, UsageSettlementWriter};
pub use standardized_usage::StandardizedUsage;
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
pub use worker::{
build_usage_queue_worker, write_event_record, UsageDataEventRecorder, UsageEventRecorder,
UsageQueueWorker, UsageRecordWriter,
build_usage_queue_worker, write_event_record, ManualProxyNodeCounter, UsageDataEventRecorder,
UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
};
pub use write::{
build_lifecycle_usage_seed, build_pending_usage_record, build_pending_usage_record_from_seed,

2
crates/aether-usage-runtime/src/request_metadata.rs
View File

@@ -94,6 +94,7 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_number(source, target, "cache_creation_price_per_1m");
copy_number(source, target, "cache_read_price_per_1m");
copy_number(source, target, "price_per_request");
copy_non_null_value(source, target, "proxy");
}
fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map<String, Value>) {
@@ -125,6 +126,7 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_number(&mut source, target, "cache_creation_price_per_1m");
remove_number(&mut source, target, "cache_read_price_per_1m");
remove_number(&mut source, target, "price_per_request");
remove_non_null_value(&mut source, target, "proxy");
}
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {

7
crates/aether-usage-runtime/src/runtime.rs
View File

@@ -53,7 +53,12 @@ impl Default for UsageBodyCapturePolicy {
#[async_trait]
pub trait UsageRuntimeAccess:
UsageRecordWriter + UsageSettlementWriter + UsageBillingEventEnricher + Send + Sync
UsageRecordWriter
+ UsageSettlementWriter
+ UsageBillingEventEnricher
+ crate::worker::ManualProxyNodeCounter
+ Send
+ Sync
{
fn has_usage_writer(&self) -> bool;
fn has_usage_worker_runner(&self) -> bool;

83
crates/aether-usage-runtime/src/worker.rs
View File

@@ -18,6 +18,17 @@ pub trait UsageEventRecorder: Send + Sync {
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError>;
}
#[async_trait]
pub trait ManualProxyNodeCounter: Send + Sync {
async fn increment_manual_proxy_node_requests(
&self,
node_id: &str,
total_delta: i64,
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), DataLayerError>;
}
#[async_trait]
pub trait UsageRecordWriter: Send + Sync {
async fn upsert_usage_record(
@@ -39,7 +50,7 @@ impl<T> UsageDataEventRecorder<T> {
#[async_trait]
impl<T> UsageEventRecorder for UsageDataEventRecorder<T>
where
T: UsageRecordWriter + UsageSettlementWriter + Send + Sync,
T: UsageRecordWriter + UsageSettlementWriter + ManualProxyNodeCounter + Send + Sync,
{
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
write_event_record(self.data.as_ref(), event).await
@@ -198,22 +209,73 @@ pub fn build_usage_queue_worker<T>(
config: UsageRuntimeConfig,
) -> Result<UsageQueueWorker, DataLayerError>
where
T: UsageRecordWriter + UsageSettlementWriter + Send + Sync + 'static,
T: UsageRecordWriter + UsageSettlementWriter + ManualProxyNodeCounter + Send + Sync + 'static,
{
UsageQueueWorker::new(runner, Arc::new(UsageDataEventRecorder::new(data)), config)
}
pub async fn write_event_record<T>(data: &T, event: &UsageEvent) -> Result<(), DataLayerError>
where
T: UsageRecordWriter + UsageSettlementWriter + Send + Sync,
T: UsageRecordWriter + UsageSettlementWriter + ManualProxyNodeCounter + Send + Sync,
{
let record = build_upsert_usage_record_from_event(event)?;
if let Some(stored) = data.upsert_usage_record(record).await? {
settle_usage_if_needed(data, &stored).await?;
}
increment_manual_proxy_node_from_event(data, event).await;
Ok(())
}
async fn increment_manual_proxy_node_from_event<T>(data: &T, event: &UsageEvent)
where
T: ManualProxyNodeCounter + Send + Sync,
{
let is_terminal = matches!(
event.event_type,
crate::UsageEventType::Completed | crate::UsageEventType::Failed
);
if !is_terminal {
return;
}
let Some(node_id) = extract_manual_proxy_node_id(event) else {
return;
};
let failed = matches!(event.event_type, crate::UsageEventType::Failed);
let failed_delta = if failed { 1i64 } else { 0i64 };
let latency_ms = event.data.response_time_ms.map(|v| v as i64);
if let Err(err) = data
.increment_manual_proxy_node_requests(&node_id, 1, failed_delta, latency_ms)
.await
{
warn!(
event_name = "manual_proxy_node_increment_failed",
log_type = "ops",
node_id = %node_id,
error = ?err,
"failed to increment manual proxy node request count"
);
}
}
fn extract_manual_proxy_node_id(event: &UsageEvent) -> Option<String> {
let metadata = event.data.request_metadata.as_ref()?;
let proxy = metadata.get("proxy")?.as_object()?;
let mode = proxy
.get("mode")
.and_then(serde_json::Value::as_str)
.unwrap_or("")
.trim();
if mode == "tunnel" {
return None;
}
proxy
.get("node_id")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|v| !v.is_empty())
.map(String::from)
}
fn consumer_name() -> String {
let host = std::env::var("HOSTNAME")
.ok()
@@ -233,7 +295,7 @@ mod tests {
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use async_trait::async_trait;
use super::{write_event_record, UsageRecordWriter};
use super::{write_event_record, ManualProxyNodeCounter, UsageRecordWriter};
use crate::{UsageEvent, UsageEventData, UsageEventType, UsageSettlementWriter};
#[derive(Default)]
@@ -317,6 +379,19 @@ mod tests {
}
}
#[async_trait]
impl ManualProxyNodeCounter for TestUsageStore {
async fn increment_manual_proxy_node_requests(
&self,
_node_id: &str,
_total_delta: i64,
_failed_delta: i64,
_latency_ms: Option<i64>,
) -> Result<(), aether_data_contracts::DataLayerError> {
Ok(())
}
}
fn sample_event() -> UsageEvent {
UsageEvent::new(
UsageEventType::Completed,

24
crates/aether-usage-runtime/src/write.rs
View File

@@ -1509,14 +1509,34 @@ fn build_runtime_request_metadata_seed(
let provider_request_has_inline_body =
context_has_inline_body(context, "provider_request_body")
|| plan_has_inline_json_body_for_usage(plan);
build_runtime_request_metadata_seed_from_parts(
let mut metadata = build_runtime_request_metadata_seed_from_parts(
context,
request_has_inline_body,
request_body_ref.as_deref(),
provider_request_has_inline_body,
provider_request_body_ref.as_deref(),
plan.body.body_bytes_b64.as_deref(),
)
);
if let Some(proxy) = plan.proxy.as_ref() {
if let Some(node_id) = proxy
.node_id
.as_deref()
.map(str::trim)
.filter(|v| !v.is_empty())
{
let mode = proxy.mode.as_deref().unwrap_or("").trim();
let mut proxy_obj = serde_json::Map::new();
proxy_obj.insert("node_id".to_string(), Value::String(node_id.to_string()));
if !mode.is_empty() {
proxy_obj.insert("mode".to_string(), Value::String(mode.to_string()));
}
let obj = metadata.get_or_insert_with(|| Value::Object(serde_json::Map::new()));
if let Value::Object(map) = obj {
map.insert("proxy".to_string(), Value::Object(proxy_obj));
}
}
}
metadata
}
fn build_runtime_request_metadata_seed_from_parts(