Merge pull request #370 from yao177/aether-rust-pioneer

fix(proxy): record manual proxy node traffic

crates/aether-data/src/repository/proxy_nodes/memory.rs

@@ -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,

crates/aether-data/src/repository/proxy_nodes/mod.rs

@@ -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,
 };

crates/aether-data/src/repository/proxy_nodes/sql.rs

@@ -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
@@ -832,6 +844,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,

crates/aether-data/src/repository/proxy_nodes/types.rs

@@ -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,