lazaret/crates/lazaret-host-coordinator/tests/lifecycle.rs

use std::cell::Cell;
use std::fs::OpenOptions;
use std::sync::{Arc, Mutex};

use lazaret_core::identity::{
    ClientId, ConfigurationFingerprint, CoordinatorId, DomainId, GenerationId, HostPath,
    NetworkSidecarId, SandboxId, SessionAnchorId, SessionId,
};
use lazaret_core::records::{Domain, DurableDecision, Generation, Session, SessionPhase};
use lazaret_host_coordinator::cold_start::BootstrapOwnerGuard;
use lazaret_host_coordinator::domain_store::{DomainStoreMutation, HostCoordinatorDomainStore};
use lazaret_host_coordinator::lifecycle::{
    DomainPhase, HostCoordinatorLifecycle, LifecycleAction, LifecycleAdapter, LifecycleError,
    SessionAdmissionIdentity, SessionLifecycleAdapter,
};
use lazaret_linux::filesystem::try_exclusive_flock;
use rustix::pipe::{PipeFlags, pipe_with};

#[derive(Clone, Default)]
struct RecordingAdapter {
    actions: Arc<Mutex<Vec<LifecycleAction>>>,
}

impl LifecycleAdapter for RecordingAdapter {
    type Error = &'static str;

    fn execute(
        &mut self,
        action: LifecycleAction,
        _generation: &Generation,
    ) -> Result<(), Self::Error> {
        self.actions.lock().unwrap().push(action);
        Ok(())
    }
}

struct PublicationFailureAdapter;

impl LifecycleAdapter for PublicationFailureAdapter {
    type Error = &'static str;

    fn execute(
        &mut self,
        _action: LifecycleAction,
        _generation: &Generation,
    ) -> Result<(), Self::Error> {
        Ok(())
    }

    fn after_commit(
        &mut self,
        action: LifecycleAction,
        _generation: &Generation,
    ) -> Result<(), Self::Error> {
        if action == LifecycleAction::PublishDomain {
            Err("exposure failed")
        } else {
            Ok(())
        }
    }
}

#[derive(Clone)]
struct TestAdmission {
    session: Session,
    fingerprint: ConfigurationFingerprint,
}

#[derive(Clone)]
struct RecordingSessionAdapter {
    actions: Arc<Mutex<Vec<LifecycleAction>>>,
}

impl SessionLifecycleAdapter for RecordingSessionAdapter {
    type Error = &'static str;
    type Admission = TestAdmission;

    fn admission_identity(&self, admission: &Self::Admission) -> SessionAdmissionIdentity {
        SessionAdmissionIdentity {
            generation_id: admission.session.generation_id,
            session_id: admission.session.session_id,
            configuration_fingerprint: admission.fingerprint,
        }
    }

    fn register_session(
        &self,
        store: &HostCoordinatorDomainStore,
        entry_action: LifecycleAction,
        admission: Self::Admission,
    ) -> Result<Session, Self::Error> {
        let snapshot = store.snapshot().map_err(|_| "snapshot")?;
        if snapshot.sessions.iter().any(|current| {
            current.session_id == admission.session.session_id
                || (matches!(
                    current.session_phase,
                    SessionPhase::Waiting
                        | SessionPhase::Registering
                        | SessionPhase::Attached
                        | SessionPhase::Prompting
                        | SessionPhase::HostExecuting
                ) && (current.cli_id == admission.session.cli_id
                    || current.session_anchor_id == admission.session.session_anchor_id))
        }) && entry_action == LifecycleAction::StartReadySession
        {
            return Err("duplicate identity");
        }
        self.actions.lock().unwrap().push(entry_action);
        self.actions
            .lock()
            .unwrap()
            .push(LifecycleAction::RegisterSession);
        let registering = Session {
            session_phase: SessionPhase::Registering,
            ..admission.session
        };
        let attached = Session {
            session_phase: SessionPhase::Attached,
            ..registering.clone()
        };
        store
            .apply(&[
                DomainStoreMutation::PutSession(registering),
                DomainStoreMutation::PutSession(attached.clone()),
            ])
            .map_err(|_| "registration")?;
        Ok(attached)
    }

    fn disconnect_session(
        &self,
        store: &HostCoordinatorDomainStore,
        session: &Session,
    ) -> Result<(), Self::Error> {
        self.actions
            .lock()
            .unwrap()
            .push(LifecycleAction::DisconnectSession);
        store
            .apply(&[DomainStoreMutation::PutSession(Session {
                session_phase: SessionPhase::Closed,
                ..session.clone()
            })])
            .map_err(|_| "disconnect")
    }
}

fn adapters() -> (RecordingAdapter, RecordingSessionAdapter) {
    let actions = Arc::new(Mutex::new(Vec::new()));
    (
        RecordingAdapter {
            actions: Arc::clone(&actions),
        },
        RecordingSessionAdapter { actions },
    )
}

fn admission(session: Session) -> TestAdmission {
    TestAdmission {
        session,
        fingerprint: fingerprint(2),
    }
}

fn bytes(value: u8) -> [u8; 32] {
    [value; 32]
}

fn domain() -> Domain {
    Domain {
        domain_id: DomainId::new(bytes(1)),
        profile_name: "trusted".into(),
        canonical_workspace: HostPath::new(b"/workspace/project".to_vec()),
    }
}

fn fingerprint(value: u8) -> ConfigurationFingerprint {
    ConfigurationFingerprint::new(bytes(value))
}

fn generation(id: u64) -> Generation {
    Generation {
        domain_id: domain().domain_id,
        generation_id: GenerationId::new(id),
        configuration_fingerprint: fingerprint(2),
        coordinator_id: CoordinatorId::new(bytes(u8::try_from(id * 3).unwrap())),
        network_sidecar_id: NetworkSidecarId::new(bytes(u8::try_from(id * 3 + 1).unwrap())),
        sandbox_id: SandboxId::new(bytes(u8::try_from(id * 3 + 2).unwrap())),
        durable_decision: DurableDecision::None,
    }
}

fn session(id: u8, generation_id: u64) -> Session {
    Session {
        generation_id: GenerationId::new(generation_id),
        session_id: SessionId::new(bytes(id)),
        cli_id: ClientId::new(bytes(id.wrapping_add(32))),
        session_anchor_id: SessionAnchorId::new(bytes(id.wrapping_add(64))),
        session_phase: SessionPhase::New,
    }
}

fn open_store(directory: &tempfile::TempDir) -> HostCoordinatorDomainStore {
    let lock_file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(false)
        .open(directory.path().join("domain.lock"))
        .unwrap();
    let lock = try_exclusive_flock(lock_file).unwrap();
    HostCoordinatorDomainStore::open(lock, &directory.path().join("domain.redb"), &domain())
        .unwrap()
}

fn complete_startup(
    lifecycle: &mut HostCoordinatorLifecycle<'_, RecordingAdapter, RecordingSessionAdapter>,
) {
    lifecycle.create_trust_anchor().unwrap();
    lifecycle.start_runtime_supervisor().unwrap();
    lifecycle.lease_prepared_lower().unwrap();
    lifecycle.mount_state_anchor().unwrap();
    lifecycle.start_sidecar().unwrap();
    lifecycle.start_connectivity_helper().unwrap();
    lifecycle.mark_network_ready().unwrap();
    lifecycle.arm_runtime_causal_edge().unwrap();
    lifecycle.start_sandbox().unwrap();
    lifecycle.publish_domain().unwrap();
}

#[test]
fn bootstrap_owner_guard_failure_uses_exact_modeled_disconnect_transition() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let actions = Arc::clone(&domain_adapter.actions);
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let checks = Cell::new(0_u8);
    let owner = session(10, 1);
    let published = lifecycle
        .bootstrap_new_domain_while_owner_live(generation(1), owner.clone(), || {
            checks.set(checks.get() + 1);
            checks.get() < 3
        })
        .unwrap();
    assert!(!published);
    assert_eq!(lifecycle.phase(), DomainPhase::Failed);
    assert!(!lifecycle.is_published());
    assert_eq!(
        lifecycle.session_phase(owner.session_id),
        Some(SessionPhase::Closed)
    );
    assert!(
        !actions
            .lock()
            .unwrap()
            .contains(&LifecycleAction::PublishDomain)
    );
}

#[test]
fn published_bootstrap_owner_eof_disconnects_waiting_session_and_tears_down() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    complete_startup(&mut lifecycle);
    let (reader, writer) = pipe_with(PipeFlags::CLOEXEC).unwrap();
    let owner_guard = BootstrapOwnerGuard::from_inherited(reader).unwrap();
    drop(writer);

    lifecycle
        .reconcile_published_bootstrap_owner(owner.session_id, owner_guard.is_live().unwrap())
        .unwrap();

    assert_eq!(lifecycle.phase(), DomainPhase::Absent);
    assert!(!lifecycle.is_published());
    assert_eq!(
        lifecycle.session_phase(owner.session_id),
        Some(SessionPhase::Closed)
    );
}

#[test]
fn admitted_bootstrap_owner_ignores_later_bootstrap_pipe_eof() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    complete_startup(&mut lifecycle);
    lifecycle.admit_session(admission(owner.clone())).unwrap();
    let (reader, writer) = pipe_with(PipeFlags::CLOEXEC).unwrap();
    let owner_guard = BootstrapOwnerGuard::from_inherited(reader).unwrap();
    drop(writer);

    lifecycle
        .reconcile_published_bootstrap_owner(owner.session_id, owner_guard.is_live().unwrap())
        .unwrap();

    assert_eq!(lifecycle.phase(), DomainPhase::Ready);
    assert!(lifecycle.is_published());
    assert_eq!(
        lifecycle.session_phase(owner.session_id),
        Some(SessionPhase::Attached)
    );
}

#[test]
fn exact_start_join_disconnect_and_teardown_actions_are_ordered() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    let waiter = session(11, 1);

    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    lifecycle.wait_for_domain(waiter.clone()).unwrap();
    complete_startup(&mut lifecycle);
    lifecycle.admit_session(admission(owner.clone())).unwrap();
    lifecycle.admit_session(admission(waiter.clone())).unwrap();
    lifecycle.disconnect_session(owner.session_id).unwrap();
    assert!(!lifecycle.is_teardown_ready());
    assert!(!lifecycle.teardown_if_idle().unwrap());
    lifecycle.disconnect_session(waiter.session_id).unwrap();
    assert!(lifecycle.is_teardown_ready());
    assert!(lifecycle.teardown_if_idle().unwrap());

    assert_eq!(lifecycle.phase(), DomainPhase::Absent);
    assert!(!lifecycle.is_published());
    assert_eq!(
        lifecycle.into_adapters().0.actions.lock().unwrap().clone(),
        vec![
            LifecycleAction::StartNewDomain,
            LifecycleAction::WaitForDomain,
            LifecycleAction::CreateTrustAnchor,
            LifecycleAction::StartRuntimeSupervisor,
            LifecycleAction::LeasePreparedLower,
            LifecycleAction::MountStateAnchor,
            LifecycleAction::StartSidecar,
            LifecycleAction::StartConnectivityHelper,
            LifecycleAction::MarkNetworkReady,
            LifecycleAction::ArmRuntimeCausalEdge,
            LifecycleAction::StartSandbox,
            LifecycleAction::PublishDomain,
            LifecycleAction::AdvanceWaitingSession,
            LifecycleAction::RegisterSession,
            LifecycleAction::AdvanceWaitingSession,
            LifecycleAction::RegisterSession,
            LifecycleAction::DisconnectSession,
            LifecycleAction::DisconnectSession,
            LifecycleAction::BeginTeardown,
            LifecycleAction::StopSandbox,
            LifecycleAction::StopSidecar,
            LifecycleAction::UnmountStateAnchor,
            LifecycleAction::ReleasePreparedLowerLease,
            LifecycleAction::GracefulSupervisorExit,
            LifecycleAction::FinishTeardown,
        ]
    );
    let snapshot = store.snapshot().unwrap();
    assert_eq!(
        snapshot.generations[0].durable_decision,
        DurableDecision::Clean
    );
    assert!(
        snapshot
            .sessions
            .iter()
            .all(|row| row.session_phase == SessionPhase::Closed)
    );
}

#[test]
fn cold_start_entry_reaches_publication_in_exact_modeled_order() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);

    lifecycle
        .bootstrap_new_domain(generation(1), session(10, 1))
        .unwrap();

    assert_eq!(lifecycle.phase(), DomainPhase::Ready);
    assert!(lifecycle.is_published());
    assert_eq!(
        lifecycle.into_adapters().0.actions.lock().unwrap().clone(),
        vec![
            LifecycleAction::StartNewDomain,
            LifecycleAction::CreateTrustAnchor,
            LifecycleAction::StartRuntimeSupervisor,
            LifecycleAction::LeasePreparedLower,
            LifecycleAction::MountStateAnchor,
            LifecycleAction::StartSidecar,
            LifecycleAction::StartConnectivityHelper,
            LifecycleAction::MarkNetworkReady,
            LifecycleAction::ArmRuntimeCausalEdge,
            LifecycleAction::StartSandbox,
            LifecycleAction::PublishDomain,
        ]
    );
}

#[test]
fn publication_exposure_failure_retains_recoverable_durable_decision() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (_, session_adapter) = adapters();
    let mut lifecycle =
        HostCoordinatorLifecycle::new(&store, PublicationFailureAdapter, session_adapter);

    assert!(
        lifecycle
            .bootstrap_new_domain(generation(1), session(10, 1))
            .is_err()
    );
    assert_eq!(lifecycle.phase(), DomainPhase::Failed);
    assert!(!lifecycle.is_published());
    assert_eq!(
        store.snapshot().unwrap().generations[0].durable_decision,
        DurableDecision::Published
    );
}

#[test]
fn publication_and_invalid_transitions_fail_before_adapter_execution() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    lifecycle
        .start_new_domain(generation(1), session(10, 1))
        .unwrap();

    assert!(matches!(
        lifecycle.publish_domain(),
        Err(LifecycleError::InvalidTransition {
            action: LifecycleAction::PublishDomain,
            ..
        })
    ));
    assert!(matches!(
        lifecycle.start_sandbox(),
        Err(LifecycleError::InvalidTransition {
            action: LifecycleAction::StartSandbox,
            ..
        })
    ));
    assert!(matches!(
        lifecycle.mount_state_anchor(),
        Err(LifecycleError::InvalidTransition {
            action: LifecycleAction::MountStateAnchor,
            ..
        })
    ));
    assert_eq!(
        lifecycle.into_adapters().0.actions.lock().unwrap().clone(),
        vec![LifecycleAction::StartNewDomain]
    );
    assert_eq!(
        store.snapshot().unwrap().generations[0].durable_decision,
        DurableDecision::None
    );
}

#[test]
fn aliased_generation_role_identities_are_refused_before_adapter_or_store() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let mut aliased = generation(1);
    aliased.network_sidecar_id = NetworkSidecarId::new(*aliased.coordinator_id.as_bytes());

    assert!(matches!(
        lifecycle.start_new_domain(aliased, session(10, 1)),
        Err(LifecycleError::InvalidTransition {
            action: LifecycleAction::StartNewDomain,
            ..
        })
    ));
    assert!(store.snapshot().unwrap().generations.is_empty());
    assert!(
        lifecycle
            .into_adapters()
            .0
            .actions
            .lock()
            .unwrap()
            .is_empty()
    );
}

#[test]
fn configuration_mismatch_is_refused_without_registration() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    complete_startup(&mut lifecycle);
    let mismatched = session(11, 1);

    assert!(matches!(
        lifecycle.admit_session(TestAdmission {
            session: mismatched.clone(),
            fingerprint: fingerprint(9),
        }),
        Err(LifecycleError::InvalidTransition {
            action: LifecycleAction::StartReadySession,
            ..
        })
    ));
    lifecycle
        .reject_new_configuration_mismatch(mismatched.clone(), fingerprint(9))
        .unwrap();
    assert_eq!(
        lifecycle.session_phase(mismatched.session_id),
        Some(SessionPhase::Failed)
    );
    assert_eq!(lifecycle.registered_sessions(), 0);
}

#[test]
fn waiting_session_restarts_after_normal_teardown() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    complete_startup(&mut lifecycle);
    lifecycle.admit_session(admission(owner.clone())).unwrap();
    lifecycle.disconnect_session(owner.session_id).unwrap();
    lifecycle.begin_teardown().unwrap();
    let waiter = session(11, 1);
    let second_waiter = session(12, 1);
    lifecycle.wait_for_domain(waiter.clone()).unwrap();
    lifecycle.wait_for_domain(second_waiter.clone()).unwrap();
    lifecycle.stop_sandbox().unwrap();
    lifecycle.stop_sidecar().unwrap();
    lifecycle.unmount_state_anchor().unwrap();
    lifecycle.release_prepared_lower_lease().unwrap();
    lifecycle.graceful_supervisor_exit().unwrap();
    lifecycle.finish_teardown().unwrap();
    let mut aliased_successor = generation(2);
    aliased_successor.sandbox_id = SandboxId::new(*aliased_successor.coordinator_id.as_bytes());
    assert!(matches!(
        lifecycle.restart_waiting_domain(aliased_successor, waiter.session_id),
        Err(LifecycleError::InvalidTransition {
            action: LifecycleAction::RestartWaitingDomain,
            ..
        })
    ));
    lifecycle
        .restart_waiting_domain(generation(2), waiter.session_id)
        .unwrap();

    assert_eq!(lifecycle.phase(), DomainPhase::Starting);
    let snapshot = store.snapshot().unwrap();
    assert!(snapshot.sessions.iter().any(|row| {
        row.generation_id == GenerationId::new(1)
            && row.session_id == waiter.session_id
            && row.session_phase == SessionPhase::Closed
    }));
    assert!(snapshot.sessions.iter().any(|row| {
        row.generation_id == GenerationId::new(2)
            && row.session_id == waiter.session_id
            && row.session_phase == SessionPhase::Waiting
    }));
    assert!(snapshot.sessions.iter().any(|row| {
        row.generation_id == GenerationId::new(2)
            && row.session_id == second_waiter.session_id
            && row.session_phase == SessionPhase::Waiting
    }));
    drop(snapshot);
    complete_startup(&mut lifecycle);
    lifecycle
        .admit_session(admission(Session {
            generation_id: GenerationId::new(2),
            ..second_waiter
        }))
        .unwrap();
    assert_eq!(
        lifecycle.session_phase(SessionId::new(bytes(12))),
        Some(SessionPhase::Attached)
    );
}

#[test]
fn admission_adapter_refuses_duplicate_cli_and_anchor_identities() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    complete_startup(&mut lifecycle);
    lifecycle.admit_session(admission(owner.clone())).unwrap();
    let duplicate = Session {
        session_id: SessionId::new(bytes(11)),
        session_anchor_id: SessionAnchorId::new(bytes(99)),
        ..owner
    };

    assert!(matches!(
        lifecycle.admit_session(admission(duplicate)),
        Err(LifecycleError::Adapter("duplicate identity"))
    ));
    assert_eq!(lifecycle.registered_sessions(), 1);
    assert_eq!(
        store
            .snapshot()
            .unwrap()
            .sessions
            .iter()
            .filter(|row| row.session_phase == SessionPhase::Attached)
            .count(),
        1
    );
}

#[test]
fn serialized_concurrent_joins_register_distinct_sessions() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    complete_startup(&mut lifecycle);
    let lifecycle = Mutex::new(lifecycle);

    std::thread::scope(|scope| {
        for id in 11..=18 {
            let lifecycle = &lifecycle;
            scope.spawn(move || {
                let candidate = session(id, 1);
                let mut guard = lifecycle.lock().unwrap();
                guard.admit_session(admission(candidate)).unwrap();
            });
        }
    });

    let lifecycle = lifecycle.into_inner().unwrap();
    assert_eq!(lifecycle.registered_sessions(), 8);
    assert_eq!(
        store
            .snapshot()
            .unwrap()
            .sessions
            .iter()
            .filter(|row| row.session_phase == SessionPhase::Attached)
            .count(),
        8
    );
}

#[test]
fn bootstrap_owner_disconnect_enters_failed_without_teardown_or_recovery() {
    let directory = tempfile::tempdir().unwrap();
    let store = open_store(&directory);
    let (domain_adapter, session_adapter) = adapters();
    let mut lifecycle = HostCoordinatorLifecycle::new(&store, domain_adapter, session_adapter);
    let owner = session(10, 1);
    lifecycle
        .start_new_domain(generation(1), owner.clone())
        .unwrap();
    lifecycle.disconnect_session(owner.session_id).unwrap();

    assert_eq!(lifecycle.phase(), DomainPhase::Failed);
    assert!(!lifecycle.is_published());
    assert!(matches!(
        lifecycle.begin_teardown(),
        Err(LifecycleError::InvalidTransition {
            action: LifecycleAction::BeginTeardown,
            ..
        })
    ));
}