mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix(proxy): record manual proxy node traffic
- add manual proxy traffic mutations to the data layer - record sync and stream outcomes for manual proxy requests - prevent tunnel proxy nodes from receiving manual traffic updates
This commit is contained in:
@@ -9,8 +9,8 @@ use uuid::Uuid;
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -437,6 +437,26 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
|
||||
Ok(Some(node.clone()))
|
||||
}
|
||||
|
||||
async fn record_traffic(
|
||||
&self,
|
||||
mutation: &ProxyNodeTrafficMutation,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let mut nodes = self.nodes.write().expect("proxy node repository lock");
|
||||
let Some(node) = nodes.get_mut(&mutation.node_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if !node.is_manual {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
node.total_requests += mutation.total_requests_delta.max(0);
|
||||
node.failed_requests += mutation.failed_requests_delta.max(0);
|
||||
node.dns_failures += mutation.dns_failures_delta.max(0);
|
||||
node.stream_errors += mutation.stream_errors_delta.max(0);
|
||||
node.updated_at_unix_secs = Self::now_unix_secs();
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn update_tunnel_status(
|
||||
&self,
|
||||
mutation: &ProxyNodeTunnelStatusMutation,
|
||||
|
||||
@@ -9,7 +9,7 @@ pub use types::{
|
||||
reconcile_remote_config_after_heartbeat, remote_config_scheduling_state,
|
||||
remote_config_upgrade_target, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
|
||||
ProxyNodeManualUpdateMutation, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
|
||||
ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
|
||||
StoredProxyNode, StoredProxyNodeEvent, PROXY_NODE_SCHEDULING_STATE_CORDONED,
|
||||
PROXY_NODE_SCHEDULING_STATE_DRAINING,
|
||||
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
PROXY_NODE_SCHEDULING_STATE_CORDONED, PROXY_NODE_SCHEDULING_STATE_DRAINING,
|
||||
};
|
||||
|
||||
@@ -6,8 +6,8 @@ use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
use super::types::{
|
||||
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
|
||||
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
||||
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
use crate::{
|
||||
error::{postgres_error, SqlxResultExt},
|
||||
@@ -320,6 +320,18 @@ WHERE id = $1
|
||||
AND is_manual = TRUE
|
||||
"#;
|
||||
|
||||
const RECORD_PROXY_NODE_TRAFFIC_SQL: &str = r#"
|
||||
UPDATE proxy_nodes
|
||||
SET
|
||||
total_requests = total_requests + GREATEST($2, 0),
|
||||
failed_requests = failed_requests + GREATEST($3, 0),
|
||||
dns_failures = dns_failures + GREATEST($4, 0),
|
||||
stream_errors = stream_errors + GREATEST($5, 0),
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND is_manual = TRUE
|
||||
"#;
|
||||
|
||||
const UNREGISTER_PROXY_NODE_SQL: &str = r#"
|
||||
UPDATE proxy_nodes
|
||||
SET
|
||||
@@ -820,6 +832,23 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
|
||||
Ok(Some(updated))
|
||||
}
|
||||
|
||||
async fn record_traffic(
|
||||
&self,
|
||||
mutation: &ProxyNodeTrafficMutation,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(RECORD_PROXY_NODE_TRAFFIC_SQL)
|
||||
.bind(&mutation.node_id)
|
||||
.bind(mutation.total_requests_delta)
|
||||
.bind(mutation.failed_requests_delta)
|
||||
.bind(mutation.dns_failures_delta)
|
||||
.bind(mutation.stream_errors_delta)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
async fn update_tunnel_status(
|
||||
&self,
|
||||
mutation: &ProxyNodeTunnelStatusMutation,
|
||||
|
||||
@@ -162,6 +162,15 @@ pub struct ProxyNodeHeartbeatMutation {
|
||||
pub proxy_version: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProxyNodeTrafficMutation {
|
||||
pub node_id: String,
|
||||
pub total_requests_delta: i64,
|
||||
pub failed_requests_delta: i64,
|
||||
pub dns_failures_delta: i64,
|
||||
pub stream_errors_delta: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProxyNodeRegistrationMutation {
|
||||
pub name: String,
|
||||
@@ -392,6 +401,11 @@ pub trait ProxyNodeWriteRepository: Send + Sync {
|
||||
mutation: &ProxyNodeHeartbeatMutation,
|
||||
) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
|
||||
|
||||
async fn record_traffic(
|
||||
&self,
|
||||
mutation: &ProxyNodeTrafficMutation,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn update_tunnel_status(
|
||||
&self,
|
||||
mutation: &ProxyNodeTunnelStatusMutation,
|
||||
|
||||
Reference in New Issue
Block a user