Fix/small bug fixes

agent/src/agent_stream.rs

@@ -75,7 +75,14 @@ async fn run(gateway_url: String, api_key: String, agent_id: Uuid, tx: mpsc::Unb
             },
         );
 
-        match tokio_tungstenite::connect_async_tls_with_config(request, None, false, connector.clone()).await {
+        match tokio_tungstenite::connect_async_tls_with_config(
+            request,
+            None,
+            false,
+            connector.clone(),
+        )
+        .await
+        {
             Ok((socket, _)) => {
                 info!(agent_id = %agent_id, "agent stream connected");
                 delay = MIN_RECONNECT_DELAY;

agent/src/docker.rs

@@ -251,30 +251,30 @@ impl DockerRuntime {
 
         let existing = self
             .docker
-            .inspect_container(&container_name, None::<InspectContainerOptionsBuilder>.map(|b| b.build()))
+            .inspect_container(
+                &container_name,
+                None::<InspectContainerOptionsBuilder>.map(|b| b.build()),
+            )
             .await;
 
-        match existing {
-            Ok(inspect) => {
-                let running = inspect
-                    .state
-                    .and_then(|s| s.status)
-                    .map(|status| status == ContainerStateStatusEnum::RUNNING)
-                    .unwrap_or(false);
-                if running {
-                    return Ok(());
-                }
-                info!(resource_group_id = %resource_group_id, "Resource group dns container exists but not running, recreating");
-                let remove_options =
-                    bollard::query_parameters::RemoveContainerOptionsBuilder::default()
-                        .force(true)
-                        .build();
-                self.docker
-                    .remove_container(&container_name, Some(remove_options))
-                    .await
-                    .context("Failed to remove stale resource group dns container")?;
+        if let Ok(inspect) = existing {
+            let running = inspect
+                .state
+                .and_then(|s| s.status)
+                .map(|status| status == ContainerStateStatusEnum::RUNNING)
+                .unwrap_or(false);
+            if running {
+                return Ok(());
             }
-            Err(_) => {}
+            info!(resource_group_id = %resource_group_id, "Resource group dns container exists but not running, recreating");
+            let remove_options =
+                bollard::query_parameters::RemoveContainerOptionsBuilder::default()
+                    .force(true)
+                    .build();
+            self.docker
+                .remove_container(&container_name, Some(remove_options))
+                .await
+                .context("Failed to remove stale resource group dns container")?;
         }
 
         crate::rg_dns::write_corefile(resource_group_id, &dns_ip)
@@ -310,10 +310,7 @@ impl DockerRuntime {
             )])),
         };
 
-        let corefile_container_path = format!(
-            "/zones/{}.Corefile",
-            resource_group_id
-        );
+        let corefile_container_path = format!("/zones/{}.Corefile", resource_group_id);
 
         let config = ContainerCreateBody {
             image: Some(DNS_IMAGE.to_string()),
@@ -822,8 +819,7 @@ impl crate::runtime::Runtime for DockerRuntime {
         workload_handle: &str,
         network_name: &str,
     ) -> Result<Option<String>> {
-        self.inspect_network_ip(workload_handle, network_name)
-            .await
+        self.inspect_network_ip(workload_handle, network_name).await
     }
 
     async fn list_managed_workloads(&self) -> Result<Vec<(String, String)>> {

agent/src/main.rs

@@ -273,7 +273,19 @@ async fn reconcile_tick(
     current_flake_rev: &mut String,
 ) {
     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, service_dns_registry, rg_dns_registry, wg_private_key_b64).await;
+    let resource_group_ids = process_workloads(
+        client,
+        api_key,
+        docker,
+        running_containers,
+        workload_phases,
+        mounted_volumes,
+        restart_counts,
+        service_dns_registry,
+        rg_dns_registry,
+        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;
@@ -282,7 +294,10 @@ async fn reconcile_tick(
     push_workload_stats(client, api_key, docker, running_containers).await;
     let metrics = system::collect_metrics();
 
-    match client.heartbeat(agent_id, api_key, Some(statuses), Some(metrics)).await {
+    match client
+        .heartbeat(agent_id, api_key, Some(statuses), Some(metrics))
+        .await
+    {
         Ok(resp) => {
             if *failure_count > 0 {
                 info!(agent_id = %agent_id, "Heartbeat recovered after {} failures", failure_count);
@@ -465,7 +480,10 @@ async fn reap_stale_containers(
 
         let dns_entry = service_dns_registry.lock().await.remove(&workload_id);
         if let Some((resource_group_id, service_name)) = dns_entry {
-            if let Err(e) = rg_dns_registry.remove(&resource_group_id, &service_name).await {
+            if let Err(e) = rg_dns_registry
+                .remove(&resource_group_id, &service_name)
+                .await
+            {
                 warn!(workload_id = %workload_id, error = %e, "Failed to remove service dns record");
             }
         }
@@ -572,7 +590,10 @@ async fn process_workloads(
                 for (workload_id, container_id) in existing {
                     containers.insert(workload_id, container_id);
                 }
-                info!(count = containers.len(), "Reconciled running containers from docker state");
+                info!(
+                    count = containers.len(),
+                    "Reconciled running containers from docker state"
+                );
             }
             Err(e) => {
                 warn!(error = %e, "Failed to reconcile running containers from docker state");
@@ -692,6 +713,17 @@ async fn start_or_restart_workload(
         }
     }
 
+    if let Some(max) = workload.max_restarts {
+        let count = *restart_counts.lock().await.get(&workload.id).unwrap_or(&0);
+        if count as i32 >= max {
+            workload_phases
+                .lock()
+                .await
+                .insert(workload.id.clone(), "failed".to_string());
+            return false;
+        }
+    }
+
     workload_phases
         .lock()
         .await
@@ -754,8 +786,14 @@ async fn start_or_restart_workload(
             restart_requested
         }
         Err(e) => {
-            workload_phases.lock().await.remove(&workload.id);
-            warn!(workload_id = %workload.id, error = ?e, "Failed to start workload");
+            let mut counts = restart_counts.lock().await;
+            let count = counts.entry(workload.id.clone()).or_insert(0);
+            *count += 1;
+            warn!(workload_id = %workload.id, error = ?e, attempt = *count, "Failed to start workload");
+            workload_phases
+                .lock()
+                .await
+                .insert(workload.id.clone(), "failed".to_string());
             false
         }
     }
@@ -770,7 +808,10 @@ async fn register_service_dns(
 ) {
     let network_name = docker::rg_network_name(resource_group_id);
 
-    let ip_address = match runtime.service_network_ip(container_id, &network_name).await {
+    let ip_address = match runtime
+        .service_network_ip(container_id, &network_name)
+        .await
+    {
         Ok(Some(ip)) => ip,
         Ok(None) => {
             warn!(container_id = %container_id, network = %network_name, "No network ip found for service dns registration");

agent/src/rg_dns.rs

@@ -62,10 +62,7 @@ impl RgDnsRegistry {
     }
 }
 
-async fn write_zone_file(
-    resource_group_id: &str,
-    records: &HashMap<String, String>,
-) -> Result<()> {
+async fn write_zone_file(resource_group_id: &str, records: &HashMap<String, String>) -> Result<()> {
     fs::create_dir_all(DNS_DIR)
         .await
         .context("Failed to create dns zone directory")?;
@@ -97,11 +94,12 @@ pub async fn write_corefile(resource_group_id: &str, listen_ip: &str) -> Result<
         .context("Failed to create dns zone directory")?;
 
     let zone = zone_name(resource_group_id);
+    let zone_path = zone_file_path(resource_group_id).display().to_string();
     let contents = format!(
-        "{zone}:53 {{\n    bind {listen_ip}\n    file /zones/{rg}.zone\n    reload 2s\n    log\n    errors\n}}\n",
+        "{zone}:53 {{\n    bind {listen_ip}\n    file {zone_path}\n    reload 2s\n    log\n    errors\n}}\n",
         zone = zone,
         listen_ip = listen_ip,
-        rg = resource_group_id,
+        zone_path = zone_path,
     );
 
     let path = corefile_path(resource_group_id);