diff --git a/Cargo.lock b/Cargo.lock index a43ea2a9..4f478a6e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -212,6 +212,7 @@ dependencies = [ "utoipa", "utoipa-swagger-ui", "uuid", + "x25519-dalek", ] [[package]] diff --git a/agent/src/client.rs b/agent/src/client.rs index 7a8e9d46..85d12a73 100644 --- a/agent/src/client.rs +++ b/agent/src/client.rs @@ -117,6 +117,12 @@ pub struct ResourceGroupPeer { pub wg_tunnel_ip: String, } +#[derive(Debug, Clone, serde::Deserialize)] +pub struct VpnPeer { + pub client_public_key: String, + pub client_tunnel_ip: String, +} + pub struct ApiClient { client: Client, gateway_url: String, @@ -327,6 +333,38 @@ impl ApiClient { .context("Failed to parse resource group peers response") } + pub async fn fetch_resource_group_vpn_peers( + &self, + api_key: &str, + resource_group_id: &str, + ) -> Result> { + let url = format!( + "{}/api/resource-groups/{}/vpn-peers", + self.gateway_url, resource_group_id + ); + + let resp = self + .client + .get(&url) + .header("X-API-Key", api_key) + .send() + .await + .context("Failed to fetch resource group vpn peers")?; + + if !resp.status().is_success() { + let status = resp.status(); + anyhow::bail!( + "Failed to fetch resource group vpn peers status={} {}", + status, + resp.text().await.unwrap_or_default() + ); + } + + resp.json::>() + .await + .context("Failed to parse resource group vpn peers response") + } + pub async fn fetch_active_resource_group_ids(&self, api_key: &str) -> Result> { let url = format!("{}/api/resource-groups/agent/active-ids", self.gateway_url); diff --git a/agent/src/main.rs b/agent/src/main.rs index f259c350..28459e9e 100644 --- a/agent/src/main.rs +++ b/agent/src/main.rs @@ -57,7 +57,9 @@ async fn main() -> Result<()> { .and_then(|v| v.parse().ok()) .unwrap_or(60); - let wg_endpoint = std::env::var("CSFX_WG_ENDPOINT").ok(); + let wg_endpoint = std::env::var("CSFX_WG_ENDPOINT") + .ok() + .or_else(detect_wg_endpoint); let wg_tunnel_ip = if config::is_registered() { config::load_config().ok().and_then(|cfg| cfg.wg_tunnel_ip) } else { @@ -109,6 +111,18 @@ async fn main() -> Result<()> { info!(agent_id = %agent_id, "Agent registered, starting heartbeat loop"); + let mgmt_tunnel_ip = config::load_config().ok().and_then(|cfg| cfg.wg_tunnel_ip); + if let Some(ref tunnel_ip) = mgmt_tunnel_ip { + if let Err(e) = + wireguard::ensure_mgmt_interface(&wg_identity.private_key_b64, MGMT_WG_PORT, tunnel_ip) + .await + { + warn!(error = %e, "Failed to bring up management WireGuard interface"); + } + } else { + warn!("No management tunnel IP available, VPN peering disabled for this agent"); + } + if let Err(e) = nftables::ensure_table_and_chain().await { warn!(error = %e, "Failed to initialize nftables resource group isolation"); } @@ -252,6 +266,7 @@ async fn run_heartbeat_loop( process_volumes(client, agent_id, api_key, &mounted_volumes).await; let resource_group_ids = process_workloads(client, api_key, &docker, &running_containers, &workload_phases, &mounted_volumes, &restart_counts, wg_private_key_b64).await; sync_wireguard_peers(client, api_key, agent_id, &resource_group_ids).await; + sync_vpn_peers(client, api_key, &resource_group_ids).await; cleanup_stale_resource_groups(client, api_key, wg_private_key_b64).await; let statuses = build_container_statuses(&docker, &running_containers, &workload_phases).await; @@ -680,6 +695,32 @@ async fn sync_wireguard_peers( } } +async fn sync_vpn_peers(client: &client::ApiClient, api_key: &str, resource_group_ids: &[String]) { + let mut wg_peers: Vec = Vec::new(); + + for resource_group_id in resource_group_ids { + match client + .fetch_resource_group_vpn_peers(api_key, resource_group_id) + .await + { + Ok(peers) => { + wg_peers.extend(peers.into_iter().map(|p| wireguard::Peer { + public_key: p.client_public_key, + endpoint: None, + allowed_ips: format!("{}/32", p.client_tunnel_ip), + })); + } + Err(e) => { + warn!(resource_group_id = %resource_group_id, error = %e, "Failed to fetch resource group vpn peers"); + } + } + } + + if let Err(e) = wireguard::reconcile_peers(wireguard::MGMT_INTERFACE_NAME, &wg_peers).await { + warn!(error = %e, "Failed to sync VPN client peers"); + } +} + async fn cleanup_stale_resource_groups( client: &client::ApiClient, api_key: &str, @@ -806,3 +847,12 @@ async fn push_workload_stats( warn!(error = %e, "Failed to push workload stats"); } } + +const MGMT_WG_PORT: u16 = 51820; + +fn detect_wg_endpoint() -> Option { + let socket = std::net::UdpSocket::bind("0.0.0.0:0").ok()?; + socket.connect("8.8.8.8:80").ok()?; + let local_ip = socket.local_addr().ok()?.ip(); + Some(format!("{}:{}", local_ip, MGMT_WG_PORT)) +} diff --git a/agent/src/wireguard.rs b/agent/src/wireguard.rs index 8c89a360..59cb70fa 100644 --- a/agent/src/wireguard.rs +++ b/agent/src/wireguard.rs @@ -9,6 +9,28 @@ pub struct Peer { pub allowed_ips: String, } +pub const MGMT_INTERFACE_NAME: &str = "wgmgmt0"; + +pub async fn ensure_mgmt_interface( + private_key_b64: &str, + listen_port: u16, + tunnel_ip: &str, +) -> Result<()> { + let iface = MGMT_INTERFACE_NAME; + + if !interface_exists(iface).await? { + run_ip(&["link", "add", "dev", iface, "type", "wireguard"]).await?; + set_private_key(iface, private_key_b64).await?; + run_wg(&["set", iface, "listen-port", &listen_port.to_string()]).await?; + run_ip(&["address", "add", &format!("{}/32", tunnel_ip), "dev", iface]).await?; + run_ip(&["link", "set", "up", "dev", iface]).await?; + + info!(iface = %iface, tunnel_ip = %tunnel_ip, "Management WireGuard interface ready"); + } + + Ok(()) +} + pub fn rg_interface_name(resource_group_id: &str) -> String { const FNV_OFFSET_BASIS: u32 = 0x811c9dc5; const FNV_PRIME: u32 = 0x01000193; @@ -111,6 +133,38 @@ pub async fn set_peers(iface: &str, peers: &[Peer]) -> Result<()> { Ok(()) } +pub async fn reconcile_peers(iface: &str, peers: &[Peer]) -> Result<()> { + let current = list_peer_public_keys(iface).await?; + let desired: std::collections::HashSet<&str> = + peers.iter().map(|p| p.public_key.as_str()).collect(); + + for stale_key in current.iter().filter(|k| !desired.contains(k.as_str())) { + run_wg(&["set", iface, "peer", stale_key, "remove"]).await?; + info!(iface = %iface, public_key = %stale_key, "Removed stale WireGuard peer"); + } + + set_peers(iface, peers).await +} + +async fn list_peer_public_keys(iface: &str) -> Result> { + let output = Command::new("wg") + .args(["show", iface, "peers"]) + .output() + .await + .context("failed to execute wg show peers")?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!("wg show peers failed iface={} stderr={}", iface, stderr); + } + + Ok(String::from_utf8_lossy(&output.stdout) + .lines() + .map(|l| l.trim().to_string()) + .filter(|l| !l.is_empty()) + .collect()) +} + async fn interface_exists(iface: &str) -> Result { let output = Command::new("ip") .args(["link", "show", iface]) diff --git a/control-plane/Dockerfile b/control-plane/Dockerfile index 246cfc1f..8d42b63b 100644 --- a/control-plane/Dockerfile +++ b/control-plane/Dockerfile @@ -21,6 +21,7 @@ FROM base AS dependencies COPY Cargo.toml Cargo.lock ./ COPY agent/Cargo.toml ./agent/ +COPY agent/csfx-guest-init/Cargo.toml ./agent/csfx-guest-init/ COPY control-plane/api-gateway/Cargo.toml ./control-plane/api-gateway/ COPY control-plane/scheduler/Cargo.toml ./control-plane/scheduler/ COPY control-plane/failover-controller/Cargo.toml ./control-plane/failover-controller/ @@ -34,6 +35,7 @@ COPY control-plane/csfx-migrate/Cargo.toml ./control-plane/csfx-migrate/ COPY control-plane/csfx-updater/Cargo.toml ./control-plane/csfx-updater/ RUN mkdir -p agent/src \ + agent/csfx-guest-init/src \ control-plane/api-gateway/src \ control-plane/scheduler/src \ control-plane/failover-controller/src \ @@ -46,6 +48,7 @@ RUN mkdir -p agent/src \ control-plane/csfx-migrate/src \ control-plane/csfx-updater/src \ && echo "fn main() {}" > agent/src/main.rs \ + && echo "fn main() {}" > agent/csfx-guest-init/src/main.rs \ && echo "fn main() {}" > control-plane/api-gateway/src/main.rs \ && echo "fn main() {}" > control-plane/scheduler/src/main.rs \ && echo "fn main() {}" > control-plane/failover-controller/src/main.rs \ diff --git a/control-plane/api-gateway/Cargo.toml b/control-plane/api-gateway/Cargo.toml index 222d966b..7cb88f0f 100644 --- a/control-plane/api-gateway/Cargo.toml +++ b/control-plane/api-gateway/Cargo.toml @@ -31,6 +31,7 @@ http = { workspace = true } rcgen = { workspace = true } rustls = { workspace = true } ring = { workspace = true } +x25519-dalek = { workspace = true } # Serialization serde = { workspace = true } diff --git a/control-plane/api-gateway/src/routes/agents.rs b/control-plane/api-gateway/src/routes/agents.rs index 30f1f49b..eabe6093 100644 --- a/control-plane/api-gateway/src/routes/agents.rs +++ b/control-plane/api-gateway/src/routes/agents.rs @@ -16,31 +16,6 @@ use crate::auth::agent::AgentApiKey; use crate::auth::rbac::{CanManageSystem, CanViewAgents}; use crate::AppState; -#[derive(Debug, Serialize, Deserialize)] -pub struct AgentRegistration { - pub agent_id: Uuid, - pub name: String, - pub hostname: String, - pub os_type: String, - pub os_version: String, - pub architecture: String, - pub agent_version: String, - pub tags: Option, -} - -#[derive(Debug, Serialize, Deserialize)] -pub struct RegistrationResponse { - pub success: bool, - pub message: String, -} - -#[derive(Debug, Serialize, Deserialize)] -pub struct Heartbeat { - pub agent_id: Uuid, - pub timestamp: chrono::DateTime, - pub status: String, -} - #[derive(Debug, Serialize, Deserialize)] pub struct SystemMetrics { pub agent_id: Uuid, @@ -109,115 +84,6 @@ impl From for AgentResponse { } } -/// Register a new agent or update existing one -pub async fn register_agent( - State(state): State, - Json(registration): Json, -) -> Result { - // Check if agent already exists - let existing_agent = agents::Entity::find() - .filter(agents::Column::Id.eq(registration.agent_id)) - .one(&state.db_conn) - .await - .map_err(|e| { - tracing::error!("Database error: {}", e); - StatusCode::INTERNAL_SERVER_ERROR - })?; - - if let Some(agent) = existing_agent { - // Update existing agent - let mut active_model: agents::ActiveModel = agent.into(); - active_model.name = ActiveValue::Set(registration.name); - active_model.hostname = ActiveValue::Set(registration.hostname); - active_model.os_type = ActiveValue::Set(registration.os_type); - active_model.os_version = ActiveValue::Set(registration.os_version); - active_model.architecture = ActiveValue::Set(registration.architecture); - active_model.agent_version = ActiveValue::Set(registration.agent_version); - active_model.status = ActiveValue::Set("online".to_string()); - active_model.last_heartbeat = ActiveValue::Set(Some(chrono::Utc::now().naive_utc())); - active_model.updated_at = ActiveValue::Set(Some(chrono::Utc::now().naive_utc())); - if let Some(tags) = registration.tags { - active_model.tags = ActiveValue::Set(Some(tags)); - } - - active_model.update(&state.db_conn).await.map_err(|e| { - tracing::error!("Failed to update agent: {}", e); - StatusCode::INTERNAL_SERVER_ERROR - })?; - - Ok(Json(RegistrationResponse { - success: true, - message: "Agent updated successfully".to_string(), - })) - } else { - // Create new agent - let new_agent = agents::ActiveModel { - id: ActiveValue::Set(registration.agent_id), - name: ActiveValue::Set(registration.name), - hostname: ActiveValue::Set(registration.hostname), - ip_address: ActiveValue::Set(None), - agent_version: ActiveValue::Set(registration.agent_version), - os_type: ActiveValue::Set(registration.os_type), - os_version: ActiveValue::Set(registration.os_version), - architecture: ActiveValue::Set(registration.architecture), - status: ActiveValue::Set("online".to_string()), - last_heartbeat: ActiveValue::Set(Some(chrono::Utc::now().naive_utc())), - registered_at: ActiveValue::Set(chrono::Utc::now().naive_utc()), - updated_at: ActiveValue::Set(None), - organization_id: ActiveValue::Set(None), - tags: ActiveValue::Set(registration.tags), - capabilities: ActiveValue::Set(None), - public_key_pem: ActiveValue::Set(None), - wg_public_key: ActiveValue::Set(None), - wg_endpoint: ActiveValue::Set(None), - wg_tunnel_ip: ActiveValue::Set(None), - kvm_capable: ActiveValue::Set(false), - cordoned: ActiveValue::Set(false), - }; - - new_agent.insert(&state.db_conn).await.map_err(|e| { - tracing::error!("Failed to create agent: {}", e); - StatusCode::INTERNAL_SERVER_ERROR - })?; - - Ok(Json(RegistrationResponse { - success: true, - message: "Agent registered successfully".to_string(), - })) - } -} - -/// Receive heartbeat from agent -pub async fn heartbeat( - State(state): State, - Json(heartbeat): Json, -) -> Result { - let agent = agents::Entity::find() - .filter(agents::Column::Id.eq(heartbeat.agent_id)) - .one(&state.db_conn) - .await - .map_err(|e| { - tracing::error!("Database error: {}", e); - StatusCode::INTERNAL_SERVER_ERROR - })?; - - if let Some(agent) = agent { - let mut active_model: agents::ActiveModel = agent.into(); - active_model.status = ActiveValue::Set(heartbeat.status); - active_model.last_heartbeat = ActiveValue::Set(Some(heartbeat.timestamp.naive_utc())); - active_model.updated_at = ActiveValue::Set(Some(chrono::Utc::now().naive_utc())); - - active_model.update(&state.db_conn).await.map_err(|e| { - tracing::error!("Failed to update heartbeat: {}", e); - StatusCode::INTERNAL_SERVER_ERROR - })?; - - Ok(StatusCode::OK) - } else { - Err(StatusCode::NOT_FOUND) - } -} - /// Receive metrics from agent pub async fn receive_metrics( State(state): State, @@ -546,8 +412,6 @@ pub fn agents_routes() -> Router { pub fn agents_unmetered_routes() -> Router { Router::new() - .route("/agents/register", post(register_agent)) - .route("/agents/heartbeat", post(heartbeat)) .route("/agents/metrics", post(receive_metrics)) .route("/agents/self/workloads", get(get_self_workloads)) .route( diff --git a/control-plane/api-gateway/src/routes/resource_groups.rs b/control-plane/api-gateway/src/routes/resource_groups.rs index e4be5739..53a96db6 100644 --- a/control-plane/api-gateway/src/routes/resource_groups.rs +++ b/control-plane/api-gateway/src/routes/resource_groups.rs @@ -8,8 +8,8 @@ use axum::{ use base64::{engine::general_purpose::STANDARD as B64, Engine}; use chrono::Utc; use entity::{ - entities::{agents, networks, resource_groups, volumes, workloads}, - Agents, Networks, ResourceGroups, Volumes, Workloads, + entities::{agents, networks, resource_group_vpn_peers, resource_groups, volumes, workloads}, + Agents, Networks, ResourceGroupVpnPeers, ResourceGroups, Volumes, Workloads, }; use ring::rand::{SecureRandom, SystemRandom}; use sea_orm::{ @@ -486,6 +486,38 @@ pub async fn list_resource_group_peers( Ok((StatusCode::OK, Json(json!(peers)))) } +#[derive(Debug, Serialize)] +pub struct VpnPeerInfo { + pub client_public_key: String, + pub client_tunnel_ip: String, +} + +pub async fn list_resource_group_vpn_peers( + _agent: crate::auth::agent::AgentApiKey, + State(state): State, + Path(id): Path, +) -> Result)> { + let peers: Vec = ResourceGroupVpnPeers::find() + .filter(resource_group_vpn_peers::Column::ResourceGroupId.eq(id)) + .all(&state.db_conn) + .await + .map_err(|e| { + tracing::error!(error = %e, "failed to list resource group vpn peers"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "error": "database error" })), + ) + })? + .into_iter() + .map(|p| VpnPeerInfo { + client_public_key: p.client_public_key, + client_tunnel_ip: p.client_tunnel_ip, + }) + .collect(); + + Ok((StatusCode::OK, Json(json!(peers)))) +} + pub async fn get_vpn_config( CanViewResourceGroups(_claims): CanViewResourceGroups, State(state): State, @@ -511,7 +543,40 @@ pub async fn get_vpn_config( ) })?; + let hosting_agents = Workloads::find() + .filter(workloads::Column::ResourceGroupId.eq(id)) + .filter(workloads::Column::AssignedAgentId.is_not_null()) + .filter( + workloads::Column::Status + .eq("scheduled") + .or(workloads::Column::Status.eq("running")), + ) + .all(&state.db_conn) + .await + .map_err(|e| { + tracing::error!(error = %e, "failed to list resource group workloads"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "error": "database error" })), + ) + })?; + + let mut agent_ids: Vec = hosting_agents + .into_iter() + .filter_map(|w| w.assigned_agent_id) + .collect(); + agent_ids.sort_unstable(); + agent_ids.dedup(); + + if agent_ids.is_empty() { + return Err(( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({ "error": "no WireGuard-enabled agent online" })), + )); + } + let gateway_agent = Agents::find() + .filter(agents::Column::Id.is_in(agent_ids)) .filter(agents::Column::Status.eq("Online")) .filter(agents::Column::WgPublicKey.is_not_null()) .filter(agents::Column::WgEndpoint.is_not_null()) @@ -542,11 +607,57 @@ pub async fn get_vpn_config( ) })?; + let client_public_key = derive_wg_public_key(&client_private_key).ok_or_else(|| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "error": "failed to derive public key" })), + ) + })?; + + let existing_ips: Vec = ResourceGroupVpnPeers::find() + .filter(resource_group_vpn_peers::Column::ResourceGroupId.eq(id)) + .all(&state.db_conn) + .await + .map_err(|e| { + tracing::error!(error = %e, "failed to list existing vpn peers"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "error": "database error" })), + ) + })? + .into_iter() + .map(|p| p.client_tunnel_ip) + .collect(); + + let client_tunnel_ip = + allocate_vpn_peer_ip(&group.internal_cidr, &existing_ips).ok_or_else(|| { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({ "error": "no free address in resource group cidr" })), + ) + })?; + + let peer = resource_group_vpn_peers::ActiveModel { + id: Set(Uuid::new_v4()), + resource_group_id: Set(id), + client_public_key: Set(client_public_key), + client_tunnel_ip: Set(client_tunnel_ip.clone()), + created_at: Set(Utc::now().naive_utc()), + }; + peer.insert(&state.db_conn).await.map_err(|e| { + tracing::error!(error = %e, "failed to store vpn peer"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "error": "database error" })), + ) + })?; + let dns = first_host_ip(&group.internal_cidr).unwrap_or_else(|| "1.1.1.1".to_string()); let config = format!( - "[Interface]\nPrivateKey = {client_private_key}\nAddress = {dns}/32\nDNS = {dns}\n\n[Peer]\nPublicKey = {server_pubkey}\nEndpoint = {endpoint}\nAllowedIPs = {cidr}\nPersistentKeepalive = 25\n", + "[Interface]\nPrivateKey = {client_private_key}\nAddress = {client_tunnel_ip}/32\nDNS = {dns}\n\n[Peer]\nPublicKey = {server_pubkey}\nEndpoint = {endpoint}\nAllowedIPs = {cidr}\nPersistentKeepalive = 25\n", client_private_key = client_private_key, + client_tunnel_ip = client_tunnel_ip, dns = dns, server_pubkey = server_pubkey, endpoint = endpoint, @@ -579,6 +690,34 @@ fn generate_wg_key() -> Result { Ok(B64.encode(key_bytes)) } +fn derive_wg_public_key(private_key_b64: &str) -> Option { + let bytes = B64.decode(private_key_b64).ok()?; + let bytes: [u8; 32] = bytes.try_into().ok()?; + let secret = x25519_dalek::StaticSecret::from(bytes); + let public = x25519_dalek::PublicKey::from(&secret); + Some(B64.encode(public.as_bytes())) +} + +fn allocate_vpn_peer_ip(cidr: &str, taken: &[String]) -> Option { + let parsed = parse_cidr(cidr)?; + let host_count = 1u32 << (32 - parsed.prefix_len); + let dns_host = parsed.network + 1; + + for offset in 2..host_count.saturating_sub(1) { + let candidate_bits = parsed.network + offset; + if candidate_bits == dns_host { + continue; + } + let [a, b, c, d] = candidate_bits.to_be_bytes(); + let candidate = format!("{}.{}.{}.{}", a, b, c, d); + if !taken.iter().any(|ip| ip == &candidate) { + return Some(candidate); + } + } + + None +} + struct Cidr { network: u32, prefix_len: u8, @@ -672,4 +811,8 @@ pub fn resource_groups_routes() -> Router { "/resource-groups/{id}/peers", get(list_resource_group_peers), ) + .route( + "/resource-groups/{id}/vpn-peers", + get(list_resource_group_vpn_peers), + ) } diff --git a/control-plane/registry/src/db/agents.rs b/control-plane/registry/src/db/agents.rs index 00c27d19..18fcd454 100644 --- a/control-plane/registry/src/db/agents.rs +++ b/control-plane/registry/src/db/agents.rs @@ -106,7 +106,7 @@ pub async fn update_heartbeat( wg_tunnel_ip: Option, agent_version: Option, kvm_capable: bool, -) -> Result<()> { +) -> Result { let mut agent: agents::ActiveModel = agents::Entity::find_by_id(agent_id) .one(db) .await? @@ -129,9 +129,7 @@ pub async fn update_heartbeat( agent.agent_version = Set(agent_version); } agent.kvm_capable = Set(kvm_capable); - agent.update(db).await?; - - Ok(()) + Ok(agent.update(db).await?) } pub async fn mark_degraded_by_timeout( diff --git a/control-plane/registry/src/handlers/agent.rs b/control-plane/registry/src/handlers/agent.rs index d5ce548c..098e9ace 100644 --- a/control-plane/registry/src/handlers/agent.rs +++ b/control-plane/registry/src/handlers/agent.rs @@ -221,6 +221,7 @@ pub async fn heartbeat( request.wg_tunnel_ip.clone(), request.agent_version.clone(), request.kvm_capable, + &state.mgmt_ipam, ) .await { @@ -355,10 +356,37 @@ async fn forward_metrics( let url = format!("{}/api/agents/metrics", state.gateway_url); - if let Err(e) = state.http_client.post(&url).json(&payload).send().await { - crate::log_warn!( - "agent_handler", - &format!("Failed to forward metrics to gateway err={}", e) - ); + for attempt in 1..=2 { + match state.http_client.post(&url).json(&payload).send().await { + Ok(resp) if resp.status().is_success() => return, + Ok(resp) => { + crate::log_warn!( + "agent_handler", + &format!( + "Metrics forward rejected by gateway attempt={} status={}", + attempt, + resp.status() + ) + ); + } + Err(e) => { + crate::log_warn!( + "agent_handler", + &format!( + "Failed to forward metrics to gateway attempt={} err={}", + attempt, e + ) + ); + } + } + + if attempt == 1 { + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + } } + + crate::log_error!( + "agent_handler", + &format!("Metrics forward permanently failed agent_id={}", agent_id) + ); } diff --git a/control-plane/registry/src/services/registry.rs b/control-plane/registry/src/services/registry.rs index 6401b980..41afecf1 100644 --- a/control-plane/registry/src/services/registry.rs +++ b/control-plane/registry/src/services/registry.rs @@ -245,8 +245,9 @@ impl AgentRegistry { wg_tunnel_ip: Option, agent_version: Option, kvm_capable: bool, + mgmt_ipam: &MgmtIpamService, ) -> Result<(), String> { - crate::db::agents::update_heartbeat( + let db_agent = crate::db::agents::update_heartbeat( &self.db, agent_id, "Online".to_string(), @@ -259,6 +260,21 @@ impl AgentRegistry { .await .map_err(|e| format!("Failed to update heartbeat: {}", e))?; + if db_agent.wg_tunnel_ip.is_none() { + let ip = mgmt_ipam + .allocate(agent_id) + .await + .map_err(|e| format!("Failed to backfill management tunnel IP: {}", e))?; + crate::db::agents::set_wg_tunnel_ip(&self.db, agent_id, &ip) + .await + .map_err(|e| format!("Failed to persist management tunnel IP: {}", e))?; + + crate::log_info!( + "agent_registry", + &format!("Backfilled management tunnel IP agent={} ip={}", agent_id, ip) + ); + } + crate::log_debug!( "agent_registry", &format!("Heartbeat received agent={}", agent_id) diff --git a/control-plane/scheduler/src/db/agents.rs b/control-plane/scheduler/src/db/agents.rs index ce5f635d..ba96976b 100644 --- a/control-plane/scheduler/src/db/agents.rs +++ b/control-plane/scheduler/src/db/agents.rs @@ -25,13 +25,6 @@ pub async fn get_online_agents_with_resources( .await?; let Some(m) = latest_metrics else { - result.push(AgentResources { - agent_id: agent.id, - free_cpu_millicores: 0, - free_memory_bytes: 0, - free_disk_bytes: 0, - kvm_capable: agent.kvm_capable, - }); continue; }; diff --git a/control-plane/scheduler/src/db/workloads.rs b/control-plane/scheduler/src/db/workloads.rs index 6aaab9d8..f956cc1a 100644 --- a/control-plane/scheduler/src/db/workloads.rs +++ b/control-plane/scheduler/src/db/workloads.rs @@ -1,6 +1,9 @@ use chrono::Utc; use entity::entities::workloads; -use sea_orm::{ActiveModelTrait, ActiveValue::Set, DatabaseConnection, EntityTrait, ModelTrait}; +use sea_orm::{ + ActiveModelTrait, ActiveValue::Set, ColumnTrait, DatabaseConnection, EntityTrait, ModelTrait, + QueryFilter, +}; use uuid::Uuid; use crate::models::workload::{ @@ -84,6 +87,14 @@ pub async fn get_all(db: &DatabaseConnection) -> Result, s Ok(rows.into_iter().map(into_response).collect()) } +pub async fn get_pending(db: &DatabaseConnection) -> Result, sea_orm::DbErr> { + workloads::Entity::find() + .filter(workloads::Column::Status.eq(WorkloadStatus::Pending.as_str())) + .filter(workloads::Column::AssignedAgentId.is_null()) + .all(db) + .await +} + pub async fn get_by_id( db: &DatabaseConnection, workload_id: Uuid, diff --git a/control-plane/scheduler/src/main.rs b/control-plane/scheduler/src/main.rs index e821ea86..79cf4ff8 100644 --- a/control-plane/scheduler/src/main.rs +++ b/control-plane/scheduler/src/main.rs @@ -52,7 +52,27 @@ async fn main() -> anyhow::Result<()> { scheduler, }; - let app = server::create_router(state); + let app = server::create_router(state.clone()); + + let retry_scheduler = state.scheduler.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(30)); + loop { + interval.tick().await; + match retry_scheduler.retry_pending().await { + Ok(0) => {} + Ok(count) => { + log_info!( + "main", + &format!("Retried pending workloads placed={}", count) + ); + } + Err(e) => { + log_error!("main", &format!("Pending retry failed err={}", e)); + } + } + } + }); let port = std::env::var("SCHEDULER_PORT") .ok() diff --git a/control-plane/scheduler/src/services/scheduler.rs b/control-plane/scheduler/src/services/scheduler.rs index 3eb29959..5b5ad188 100644 --- a/control-plane/scheduler/src/services/scheduler.rs +++ b/control-plane/scheduler/src/services/scheduler.rs @@ -448,6 +448,84 @@ impl SchedulerService { Ok(responses) } + pub async fn retry_pending(&self) -> Result { + let pending = crate::db::workloads::get_pending(&self.db) + .await + .map_err(|e| format!("Failed to fetch pending workloads: {}", e))?; + + if pending.is_empty() { + return Ok(0); + } + + let mut agents = crate::db::agents::get_online_agents_with_resources(&self.db) + .await + .map_err(|e| format!("Failed to fetch agent resources: {}", e))?; + + for agent in agents.iter_mut() { + let (reserved_cpu, reserved_mem, reserved_disk) = + crate::db::agents::get_assigned_workload_resources(&self.db, agent.agent_id) + .await + .map_err(|e| format!("Failed to fetch reserved resources: {}", e))?; + + agent.free_cpu_millicores -= reserved_cpu; + agent.free_memory_bytes -= reserved_mem; + agent.free_disk_bytes -= reserved_disk; + } + + let mut placed_count = 0; + + for workload in pending { + let runtime_class = + crate::models::workload::RuntimeClass::from_str(&workload.runtime_class); + + let Some(agent_id) = self.first_fit_resources( + workload.cpu_millicores, + workload.memory_bytes, + workload.disk_bytes, + runtime_class.requires_kvm(), + &agents, + ) else { + continue; + }; + + crate::db::workloads::assign(&self.db, workload.id, agent_id) + .await + .map_err(|e| format!("Failed to assign workload: {}", e))?; + + let record = PlacementRecord { + workload_id: workload.id, + agent_id, + image: workload.image.clone(), + cpu_millicores: workload.cpu_millicores, + memory_bytes: workload.memory_bytes, + disk_bytes: workload.disk_bytes, + scheduled_at: Utc::now().to_rfc3339(), + stack_id: workload.stack_id, + service_name: workload.service_name.clone(), + runtime_class: workload.runtime_class.clone(), + }; + put_placement(&self.etcd, &record).await?; + + if let Some(agent) = agents.iter_mut().find(|a| a.agent_id == agent_id) { + agent.free_cpu_millicores -= workload.cpu_millicores; + agent.free_memory_bytes -= workload.memory_bytes; + agent.free_disk_bytes -= workload.disk_bytes; + } + + crate::log_info!( + "scheduler", + &format!( + "Pending workload scheduled workload_id={} agent_id={}", + workload.id, agent_id + ) + ); + + placed_count += 1; + } + + Ok(placed_count) + } + pub async fn delete_workload(&self, workload_id: Uuid) -> Result<(), String> { crate::db::workloads::delete(&self.db, workload_id) .await diff --git a/control-plane/shared/entity/src/entities/mod.rs b/control-plane/shared/entity/src/entities/mod.rs index 8b1b761b..343c9f68 100644 --- a/control-plane/shared/entity/src/entities/mod.rs +++ b/control-plane/shared/entity/src/entities/mod.rs @@ -14,6 +14,7 @@ pub mod networks; pub mod organization; pub mod permission; pub mod registry_tokens; +pub mod resource_group_vpn_peers; pub mod resource_groups; pub mod role; pub mod role_permission; @@ -42,6 +43,7 @@ pub use networks::Entity as Networks; pub use organization::Entity as Organization; pub use permission::Entity as Permission; pub use registry_tokens::Entity as RegistryTokens; +pub use resource_group_vpn_peers::Entity as ResourceGroupVpnPeers; pub use resource_groups::Entity as ResourceGroups; pub use role::Entity as Role; pub use role_permission::Entity as RolePermission; diff --git a/control-plane/shared/entity/src/entities/resource_group_vpn_peers.rs b/control-plane/shared/entity/src/entities/resource_group_vpn_peers.rs new file mode 100644 index 00000000..57f20e61 --- /dev/null +++ b/control-plane/shared/entity/src/entities/resource_group_vpn_peers.rs @@ -0,0 +1,32 @@ +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)] +#[sea_orm(table_name = "resource_group_vpn_peers")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub id: Uuid, + pub resource_group_id: Uuid, + pub client_public_key: String, + pub client_tunnel_ip: String, + pub created_at: chrono::NaiveDateTime, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::resource_groups::Entity", + from = "Column::ResourceGroupId", + to = "super::resource_groups::Column::Id", + on_delete = "Cascade" + )] + ResourceGroups, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::ResourceGroups.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/control-plane/shared/entity/src/entities/resource_groups.rs b/control-plane/shared/entity/src/entities/resource_groups.rs index 1b8febac..0a153d50 100644 --- a/control-plane/shared/entity/src/entities/resource_groups.rs +++ b/control-plane/shared/entity/src/entities/resource_groups.rs @@ -32,6 +32,8 @@ pub enum Relation { Networks, #[sea_orm(has_many = "super::workload_stacks::Entity")] WorkloadStacks, + #[sea_orm(has_many = "super::resource_group_vpn_peers::Entity")] + ResourceGroupVpnPeers, } impl Related for Entity { @@ -64,4 +66,10 @@ impl Related for Entity { } } +impl Related for Entity { + fn to() -> RelationDef { + Relation::ResourceGroupVpnPeers.def() + } +} + impl ActiveModelBehavior for ActiveModel {} diff --git a/control-plane/shared/migration/src/lib.rs b/control-plane/shared/migration/src/lib.rs index dfc46a98..a23d2b9c 100644 --- a/control-plane/shared/migration/src/lib.rs +++ b/control-plane/shared/migration/src/lib.rs @@ -28,6 +28,7 @@ mod m20260710_000000_add_workload_restart_policy; mod m20260711_000000_add_runtime_class; mod m20260711_010000_add_agent_cordoned; mod m20260712_000000_add_workload_lifecycle; +mod m20260712_010000_add_resource_group_vpn_peers; pub struct Migrator; @@ -63,6 +64,7 @@ impl MigratorTrait for Migrator { Box::new(m20260711_000000_add_runtime_class::Migration), Box::new(m20260711_010000_add_agent_cordoned::Migration), Box::new(m20260712_000000_add_workload_lifecycle::Migration), + Box::new(m20260712_010000_add_resource_group_vpn_peers::Migration), ] } } diff --git a/control-plane/shared/migration/src/m20260712_010000_add_resource_group_vpn_peers.rs b/control-plane/shared/migration/src/m20260712_010000_add_resource_group_vpn_peers.rs new file mode 100644 index 00000000..986131f9 --- /dev/null +++ b/control-plane/shared/migration/src/m20260712_010000_add_resource_group_vpn_peers.rs @@ -0,0 +1,97 @@ +use sea_orm_migration::prelude::*; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .create_table( + Table::create() + .table(ResourceGroupVpnPeers::Table) + .if_not_exists() + .col( + ColumnDef::new(ResourceGroupVpnPeers::Id) + .uuid() + .not_null() + .primary_key(), + ) + .col( + ColumnDef::new(ResourceGroupVpnPeers::ResourceGroupId) + .uuid() + .not_null(), + ) + .col( + ColumnDef::new(ResourceGroupVpnPeers::ClientPublicKey) + .string() + .not_null(), + ) + .col( + ColumnDef::new(ResourceGroupVpnPeers::ClientTunnelIp) + .string() + .not_null(), + ) + .col( + ColumnDef::new(ResourceGroupVpnPeers::CreatedAt) + .date_time() + .not_null(), + ) + .foreign_key( + ForeignKey::create() + .from( + ResourceGroupVpnPeers::Table, + ResourceGroupVpnPeers::ResourceGroupId, + ) + .to(ResourceGroups::Table, ResourceGroups::Id) + .on_delete(ForeignKeyAction::Cascade), + ) + .to_owned(), + ) + .await?; + + manager + .create_index( + Index::create() + .table(ResourceGroupVpnPeers::Table) + .col(ResourceGroupVpnPeers::ResourceGroupId) + .name("idx_resource_group_vpn_peers_resource_group_id") + .to_owned(), + ) + .await?; + + manager + .create_index( + Index::create() + .table(ResourceGroupVpnPeers::Table) + .col(ResourceGroupVpnPeers::ClientTunnelIp) + .col(ResourceGroupVpnPeers::ResourceGroupId) + .name("idx_resource_group_vpn_peers_tunnel_ip") + .unique() + .to_owned(), + ) + .await + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .drop_table(Table::drop().table(ResourceGroupVpnPeers::Table).to_owned()) + .await + } +} + +#[derive(DeriveIden)] +enum ResourceGroupVpnPeers { + Table, + Id, + ResourceGroupId, + ClientPublicKey, + ClientTunnelIp, + CreatedAt, +} + +#[derive(DeriveIden)] +enum ResourceGroups { + Table, + Id, +}