lazaret/crates/lazaret-storage-catalog/tests/catalog.rs

use std::fs::OpenOptions;
use std::sync::{Arc, Barrier};
use std::thread;
use std::time::Duration;

use lazaret_core::codec::encode;
use lazaret_core::identity::{
    DeviceId, DomainId, GenerationId, HostPath, IdMapDigest, InodeId, PolicyId, PrepareOperationId,
    PreparedArtifactId, SelectedManifestDigest, SourceIdentity,
};
use lazaret_core::records::{
    ArtifactLease, ArtifactState, PrepareOperation, PrepareStageState, PreparedArtifact,
};
use lazaret_linux::filesystem::try_exclusive_flock;
use lazaret_storage_catalog::{Catalog, CatalogError, CatalogRow};
use redb::{Database, TableDefinition};
use serde::Serialize;
use tempfile::TempDir;

const METADATA: TableDefinition<u8, &[u8]> = TableDefinition::new("catalog_metadata");

struct Fixture {
    directory: TempDir,
    catalog: Catalog,
}

impl Fixture {
    fn new() -> Self {
        let directory = TempDir::new().unwrap();
        let catalog = Catalog::new(
            directory.path().join("catalog.redb"),
            directory.path().join("storage-catalog.lock"),
        );
        catalog.initialize().unwrap();
        Self { directory, catalog }
    }

    fn artifact(seed: u8) -> PreparedArtifact {
        PreparedArtifact {
            prepared_artifact_id: PreparedArtifactId::new([seed; 32]),
            source_identity: SourceIdentity::new(vec![seed, 1]),
            selected_manifest_digest: SelectedManifestDigest::new([seed; 32]),
            policy_id: PolicyId::new([seed.wrapping_add(1); 32]),
            id_map_digest: IdMapDigest::new([seed.wrapping_add(2); 32]),
            path: HostPath::new(format!("artifacts/{seed}").into_bytes()),
            device: DeviceId::new(u64::from(seed)),
            inode: InodeId::new(u64::from(seed) + 100),
            state: ArtifactState::Available,
        }
    }

    fn prepare(artifact: &PreparedArtifact, seed: u8) -> PrepareOperation {
        PrepareOperation {
            operation_id: PrepareOperationId::new([seed; 32]),
            expected_artifact_id: artifact.prepared_artifact_id,
            source_identity: artifact.source_identity.clone(),
            selected_manifest_digest: artifact.selected_manifest_digest,
            policy_id: artifact.policy_id,
            id_map_digest: artifact.id_map_digest,
            stage_path: HostPath::new(format!("stages/{seed}").into_bytes()),
            stage_state: PrepareStageState::Owned,
            device: Some(artifact.device),
            inode: Some(artifact.inode),
        }
    }

    fn publish(&self, seed: u8) -> PreparedArtifact {
        let artifact = Self::artifact(seed);
        let operation = Self::prepare(&artifact, seed);
        let intent = PrepareOperation {
            stage_state: PrepareStageState::Intent,
            device: None,
            inode: None,
            ..operation.clone()
        };
        self.catalog.record_prepare(&intent).unwrap();
        self.catalog
            .bind_prepare_stage(&intent, &operation)
            .unwrap();
        self.catalog.publish(&operation, &artifact).unwrap();
        artifact
    }
}

fn lease(artifact: &PreparedArtifact, domain: u8, generation: u64) -> ArtifactLease {
    ArtifactLease {
        prepared_artifact_id: artifact.prepared_artifact_id,
        domain_id: DomainId::new([domain; 32]),
        generation_id: GenerationId::new(generation),
    }
}

#[test]
fn catalog_reopens_with_committed_rows() {
    let fixture = Fixture::new();
    let artifact = fixture.publish(1);
    let first_lease = lease(&artifact, 10, 1);
    fixture.catalog.acquire_lease(&first_lease).unwrap();

    let reopened = Catalog::new(
        fixture.directory.path().join("catalog.redb"),
        fixture.directory.path().join("storage-catalog.lock"),
    );
    let snapshot = reopened.snapshot().unwrap();
    assert_eq!(snapshot.prepared_artifacts, vec![artifact]);
    assert_eq!(snapshot.artifact_leases, vec![first_lease]);
    assert!(snapshot.prepare_operations.is_empty());
}

#[test]
fn publication_is_one_exact_transaction() {
    let fixture = Fixture::new();
    let artifact = Fixture::artifact(2);
    let operation = Fixture::prepare(&artifact, 2);
    let intent = PrepareOperation {
        stage_state: PrepareStageState::Intent,
        device: None,
        inode: None,
        ..operation.clone()
    };
    fixture.catalog.record_prepare(&intent).unwrap();
    fixture
        .catalog
        .bind_prepare_stage(&intent, &operation)
        .unwrap();

    let mut forged = operation.clone();
    forged.inode = Some(InodeId::new(999));
    assert!(matches!(
        fixture.catalog.publish(&forged, &artifact),
        Err(CatalogError::OwnershipMismatch {
            row: CatalogRow::PrepareOperation
        })
    ));
    let before = fixture.catalog.snapshot().unwrap();
    assert_eq!(before.prepare_operations, vec![operation.clone()]);
    assert!(before.prepared_artifacts.is_empty());

    fixture.catalog.publish(&operation, &artifact).unwrap();
    let after = fixture.catalog.snapshot().unwrap();
    assert!(after.prepare_operations.is_empty());
    assert_eq!(after.prepared_artifacts, vec![artifact]);
}

#[test]
fn shared_artifact_waits_for_every_lease() {
    let fixture = Fixture::new();
    let artifact = fixture.publish(3);
    let first = lease(&artifact, 1, 1);
    let second = lease(&artifact, 2, 9);
    fixture.catalog.acquire_lease(&first).unwrap();
    fixture.catalog.acquire_lease(&second).unwrap();

    assert!(matches!(
        fixture.catalog.begin_reclaim(&artifact),
        Err(CatalogError::LeasesPresent)
    ));
    fixture.catalog.release_lease(&first).unwrap();
    assert!(matches!(
        fixture.catalog.begin_reclaim(&artifact),
        Err(CatalogError::LeasesPresent)
    ));
    fixture.catalog.release_lease(&second).unwrap();
    let reclaiming = fixture.catalog.begin_reclaim(&artifact).unwrap();
    assert_eq!(reclaiming.state, ArtifactState::Reclaiming);
    fixture.catalog.finish_reclaim(&reclaiming).unwrap();
    assert!(
        fixture
            .catalog
            .snapshot()
            .unwrap()
            .prepared_artifacts
            .is_empty()
    );
}

#[test]
fn one_domain_generation_cannot_bind_two_artifacts() {
    let fixture = Fixture::new();
    let first_artifact = fixture.publish(5);
    let second_artifact = fixture.publish(6);
    let first = lease(&first_artifact, 4, 11);
    let conflicting = lease(&second_artifact, 4, 11);

    fixture.catalog.acquire_lease(&first).unwrap();
    assert!(matches!(
        fixture.catalog.acquire_lease(&conflicting),
        Err(CatalogError::AlreadyExists {
            row: CatalogRow::ArtifactLease
        })
    ));
    assert_eq!(
        fixture.catalog.snapshot().unwrap().artifact_leases,
        vec![first]
    );
}

#[test]
fn lease_and_reclaim_have_exactly_one_transaction_winner() {
    for seed in 10..30 {
        let fixture = Fixture::new();
        let artifact = fixture.publish(seed);
        let candidate = lease(&artifact, seed, u64::from(seed));
        let barrier = Arc::new(Barrier::new(3));

        let lease_catalog = fixture.catalog.clone();
        let lease_barrier = Arc::clone(&barrier);
        let lease_handle = thread::spawn(move || {
            lease_barrier.wait();
            lease_catalog.acquire_lease(&candidate)
        });

        let reclaim_catalog = fixture.catalog.clone();
        let reclaim_barrier = Arc::clone(&barrier);
        let reclaim_artifact = artifact.clone();
        let reclaim_handle = thread::spawn(move || {
            reclaim_barrier.wait();
            reclaim_catalog.begin_reclaim(&reclaim_artifact)
        });

        barrier.wait();
        let lease_result = lease_handle.join().unwrap();
        let reclaim_result = reclaim_handle.join().unwrap();
        assert_ne!(lease_result.is_ok(), reclaim_result.is_ok());
        if lease_result.is_ok() {
            assert!(matches!(reclaim_result, Err(CatalogError::LeasesPresent)));
        } else {
            assert!(matches!(
                lease_result,
                Err(CatalogError::ArtifactUnavailable)
            ));
        }
    }
}

#[test]
fn exact_row_deletion_refuses_forged_ownership() {
    let fixture = Fixture::new();
    let artifact = fixture.publish(4);
    let first = lease(&artifact, 1, 7);
    fixture.catalog.acquire_lease(&first).unwrap();

    let forged_lease = lease(&artifact, 2, 7);
    assert!(matches!(
        fixture.catalog.release_lease(&forged_lease),
        Err(CatalogError::NotFound {
            row: CatalogRow::ArtifactLease
        })
    ));
    assert_eq!(
        fixture.catalog.snapshot().unwrap().artifact_leases,
        vec![first]
    );

    fixture.catalog.release_lease(&first).unwrap();
    let reclaiming = fixture.catalog.begin_reclaim(&artifact).unwrap();
    let mut forged_artifact = reclaiming.clone();
    forged_artifact.inode = InodeId::new(999);
    assert!(matches!(
        fixture.catalog.finish_reclaim(&forged_artifact),
        Err(CatalogError::OwnershipMismatch {
            row: CatalogRow::PreparedArtifact
        })
    ));
    assert_eq!(
        fixture.catalog.snapshot().unwrap().prepared_artifacts,
        vec![reclaiming]
    );
}

#[derive(Serialize)]
struct UnknownMetadata<'a> {
    schema_id: &'a str,
    schema_version: u32,
}

#[test]
fn unknown_version_refuses_before_rows() {
    let fixture = Fixture::new();
    let database = Database::open(fixture.directory.path().join("catalog.redb")).unwrap();
    let transaction = database.begin_write().unwrap();
    {
        let mut metadata = transaction.open_table(METADATA).unwrap();
        let encoded = encode(&UnknownMetadata {
            schema_id: "lazaret-storage-catalog-v1",
            schema_version: 2,
        })
        .unwrap();
        metadata.insert(0, encoded.as_slice()).unwrap();
    }
    transaction.commit().unwrap();
    drop(database);

    assert!(matches!(
        fixture.catalog.snapshot(),
        Err(CatalogError::UnknownSchema)
    ));
}

#[test]
fn lock_wait_is_bounded() {
    let fixture = Fixture::new();
    let lock_file = OpenOptions::new()
        .read(true)
        .write(true)
        .open(fixture.directory.path().join("storage-catalog.lock"))
        .unwrap();
    let _held = try_exclusive_flock(lock_file).unwrap();
    let impatient = fixture
        .catalog
        .clone()
        .with_lock_wait(Duration::from_millis(10));
    assert!(matches!(
        impatient.snapshot(),
        Err(CatalogError::LockTimeout)
    ));
}