reconciliation.rs

//! Server-startup reconciliation: sync SQLite state with live Docker state.
//!
//! Called once during [`run_server`] startup, before accepting connections.
//!
//! ## What it does
//!
//! 1. **Dead sessions** — sessions with state `Starting` or `Running` in DB
//!    whose container no longer exists (or has exited) are marked `Stopped`.
//!
//! 2. **Orphan containers** — containers labelled `relay.managed=true` that
//!    have no corresponding DB session are force-removed.
//!
//! 3. **Orphan networks** — Docker networks labelled `relay.managed=true`
//!    that are not attached to any managed container are removed.
//!
//! All errors during reconciliation are logged and tolerated — a partial
//! reconciliation is better than refusing to start.

use std::collections::HashSet;

use tracing::{error, info, warn};

use crate::docker::DockerOrchestrator;
use crate::store::Store;

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// Run the full reconciliation pass.
///
/// Errors from individual reconciliation steps are logged but not propagated;
/// the server should start even if cleanup is partial.
pub async fn reconcile(store: &Store, docker: &DockerOrchestrator) {
    info!("reconciliation: starting");

    // Fetch the container list once and share it between steps 1 and 2, which
    // can run in parallel.  Step 3 re-fetches after step 2 has removed orphan
    // containers so the network check sees the updated state.
    let containers = match docker.list_managed_containers().await {
        Ok(c) => c,
        Err(e) => {
            error!(error = %e, "reconciliation: failed to list Docker containers");
            info!(
                dead_sessions = 0,
                orphan_containers = 0,
                orphan_networks = 0,
                "reconciliation: complete (skipped — Docker unavailable)"
            );
            return;
        }
    };

    let (dead_sessions, orphan_containers) = tokio::join!(
        reconcile_dead_sessions(store, &containers),
        reconcile_orphan_containers(store, docker, &containers),
    );

    let orphan_networks = reconcile_orphan_networks(docker).await;

    info!(
        dead_sessions,
        orphan_containers, orphan_networks, "reconciliation: complete"
    );
}

// ---------------------------------------------------------------------------
// Step 1 — mark dead sessions as Stopped
// ---------------------------------------------------------------------------

/// Mark `Starting`/`Running` sessions whose container is absent or not running
/// as `Stopped`.
///
/// Returns the number of sessions updated.
async fn reconcile_dead_sessions(
    store: &Store,
    containers: &[bollard::models::ContainerSummary],
) -> u32 {
    // Collect the set of container IDs that are currently running in Docker.
    let running_ids: HashSet<String> = containers
        .iter()
        .filter_map(|c| {
            // Only containers in a "running" state count as alive.
            let is_running = matches!(
                c.state,
                Some(bollard::models::ContainerSummaryStateEnum::RUNNING)
            );
            if is_running {
                c.id.clone()
            } else {
                None
            }
        })
        .collect();

    let active_sessions = match store.list_active_sessions().await {
        Ok(s) => s,
        Err(e) => {
            error!(error = %e, "reconciliation: failed to list active sessions from DB");
            return 0;
        }
    };

    let mut stopped = 0u32;

    for session in active_sessions {
        let container_alive = session
            .container_id
            .as_deref()
            .is_some_and(|cid| running_ids.contains(cid));

        if !container_alive {
            let container_id = session.container_id.as_deref().unwrap_or("<none>");
            match store
                .force_stop_session(&session.id, "container not found after runner restart")
                .await
            {
                Ok(true) => {
                    info!(
                        session_id = %session.id,
                        container_id,
                        "reconciliation: session marked Stopped (container dead)"
                    );
                    stopped += 1;
                }
                Ok(false) => {
                    // Already stopped by a concurrent writer — nothing to do.
                }
                Err(e) => {
                    warn!(
                        session_id = %session.id,
                        error = %e,
                        "reconciliation: failed to stop dead session"
                    );
                }
            }
        }
    }

    stopped
}

// ---------------------------------------------------------------------------
// Step 2 — remove orphan containers
// ---------------------------------------------------------------------------

/// Remove managed containers that have no corresponding DB session.
///
/// Returns the number of containers removed.
async fn reconcile_orphan_containers(
    store: &Store,
    docker: &DockerOrchestrator,
    containers: &[bollard::models::ContainerSummary],
) -> u32 {
    let mut removed = 0u32;

    for container in containers {
        let container_id = match &container.id {
            Some(id) => id.clone(),
            None => continue,
        };

        // Extract the session ID from the container's labels.
        let session_id = container
            .labels
            .as_ref()
            .and_then(|l| l.get("relay.session_id"))
            .cloned();

        let is_orphan = match &session_id {
            None => {
                // No session label at all — treat as orphan.
                true
            }
            Some(sid) => match store.get_session(sid).await {
                Ok(Some(_)) => false, // Session exists in DB — keep container.
                Ok(None) => true,     // No DB record — orphan.
                Err(e) => {
                    warn!(
                        container_id = %container_id,
                        session_id = %sid,
                        error = %e,
                        "reconciliation: DB lookup failed, skipping container"
                    );
                    false // Skip on error to avoid false-positive removal.
                }
            },
        };

        if is_orphan {
            match docker.remove_container(&container_id).await {
                Ok(()) => {
                    info!(
                        container_id = %container_id,
                        session_id = ?session_id,
                        "reconciliation: orphan container removed"
                    );
                    removed += 1;
                }
                Err(e) => {
                    warn!(
                        container_id = %container_id,
                        error = %e,
                        "reconciliation: failed to remove orphan container"
                    );
                }
            }
        }
    }

    removed
}

// ---------------------------------------------------------------------------
// Step 3 — remove orphan networks
// ---------------------------------------------------------------------------

/// Remove managed networks that are not connected to any managed container.
///
/// Returns the number of networks removed.
async fn reconcile_orphan_networks(docker: &DockerOrchestrator) -> u32 {
    // Re-fetch containers (after orphan container removal) to get the live set.
    let containers = match docker.list_managed_containers().await {
        Ok(c) => c,
        Err(e) => {
            error!(error = %e, "reconciliation: failed to list containers for network check");
            return 0;
        }
    };

    // Build the set of network names that still have at least one container.
    let occupied_networks: HashSet<String> = containers
        .into_iter()
        .filter_map(|c| c.network_settings)
        .filter_map(|ns| ns.networks)
        .flat_map(std::collections::HashMap::into_keys)
        .collect();

    let networks = match docker.list_managed_networks().await {
        Ok(n) => n,
        Err(e) => {
            error!(error = %e, "reconciliation: failed to list managed networks");
            return 0;
        }
    };

    let mut removed = 0u32;

    for network in networks {
        let name = match &network.name {
            Some(n) => n.clone(),
            None => continue,
        };

        if !occupied_networks.contains(&name) {
            match docker.remove_network_if_exists(&name).await {
                Ok(()) => {
                    info!(network = %name, "reconciliation: orphan network removed");
                    removed += 1;
                }
                Err(e) => {
                    warn!(
                        network = %name,
                        error = %e,
                        "reconciliation: failed to remove orphan network"
                    );
                }
            }
        }
    }

    removed
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use crate::store::{CreateProject, CreateSession, SessionState, Store};

    async fn make_store() -> Store {
        Store::for_tests().await
    }

    #[tokio::test]
    async fn list_active_sessions_returns_starting_and_running() {
        let store = make_store().await;

        // Create project → workspace → 3 sessions in different states.
        let proj = store
            .create_project(CreateProject {
                name: "p".into(),
                git_url: "https://example.com/repo.git".into(),
                local_path: "/tmp/p".into(),
            })
            .await
            .expect("create project");

        // Session 1: Starting (default after create_session)
        let s1 = store
            .create_session(CreateSession {
                project_id: proj.id.clone(),
                branch: "main".into(),
                image: "ubuntu:24.04".into(),
                profile: "default".into(),
                work_dir: String::new(),
            })
            .await
            .expect("create session s1");

        // Session 2: Running
        let s2 = store
            .create_session(CreateSession {
                project_id: proj.id.clone(),
                branch: "main".into(),
                image: "ubuntu:24.04".into(),
                profile: "default".into(),
                work_dir: String::new(),
            })
            .await
            .expect("create session s2");
        store
            .update_session_state(
                &s2.id,
                SessionState::Running,
                Some("cid-running"),
                None,
                &s2.updated_at.to_rfc3339(),
            )
            .await
            .expect("update s2 to Running");

        // Session 3: Stopped
        let s3 = store
            .create_session(CreateSession {
                project_id: proj.id.clone(),
                branch: "main".into(),
                image: "ubuntu:24.04".into(),
                profile: "default".into(),
                work_dir: String::new(),
            })
            .await
            .expect("create session s3");
        store
            .update_session_state(
                &s3.id,
                SessionState::Stopped,
                None,
                None,
                &s3.updated_at.to_rfc3339(),
            )
            .await
            .expect("update s3 to Stopped");

        let active = store
            .list_active_sessions()
            .await
            .expect("list active sessions");
        let ids: Vec<&str> = active.iter().map(|s| s.id.as_str()).collect();

        assert!(
            ids.contains(&s1.id.as_str()),
            "Starting session should be active"
        );
        assert!(
            ids.contains(&s2.id.as_str()),
            "Running session should be active"
        );
        assert!(
            !ids.contains(&s3.id.as_str()),
            "Stopped session should not be active"
        );
    }

    #[tokio::test]
    async fn force_stop_session_updates_state() {
        let store = make_store().await;

        let proj = store
            .create_project(CreateProject {
                name: "p2".into(),
                git_url: "https://example.com/repo.git".into(),
                local_path: "/tmp/p2".into(),
            })
            .await
            .expect("create project");

        let session = store
            .create_session(CreateSession {
                project_id: proj.id.clone(),
                branch: "main".into(),
                image: "ubuntu:24.04".into(),
                profile: "default".into(),
                work_dir: String::new(),
            })
            .await
            .expect("create session");

        // First call should succeed (state is Starting).
        let updated = store
            .force_stop_session(&session.id, "container not found after runner restart")
            .await
            .expect("force_stop_session first call");
        assert!(
            updated,
            "force_stop_session should return true when updating Starting"
        );

        // Verify state is now Stopped.
        let fetched = store
            .get_session(&session.id)
            .await
            .expect("get_session query")
            .expect("session should exist");
        assert_eq!(fetched.state, SessionState::Stopped);
        assert_eq!(
            fetched.error_reason.as_deref(),
            Some("container not found after runner restart")
        );

        // Second call on already-stopped session should return false.
        let updated_again = store
            .force_stop_session(&session.id, "should not change")
            .await
            .expect("force_stop_session second call");
        assert!(
            !updated_again,
            "force_stop_session should return false when already Stopped"
        );
    }
}