fdb.rs

#![cfg(test)]

use foundationdb::api::NetworkAutoStop;
use foundationdb::Database;
use log::{error, info};
use rand::{thread_rng, Rng};
use std::borrow::Cow;
use std::collections::HashMap;
use std::io::Write;
use std::ops::Deref;
use std::sync::{Arc, LazyLock, RwLock};
use tempfile::NamedTempFile;
use testcontainers::core::{CmdWaitFor, ContainerPort, ContainerState, ExecCommand, WaitFor};
use testcontainers::runners::SyncRunner;
use testcontainers::{Container, Image, ImageExt, TestcontainersError};
use thiserror::Error;

static IMAGE_NAME: &str = "foundationdb/foundationdb";
static IMAGE_TAG: &str = "7.3.51";

struct FdbImage {
    env_vars: HashMap<String, String>,
    exposed_ports: Vec<ContainerPort>,
}

impl FdbImage {
    fn get_port(&self) -> ContainerPort {
        *self
            .exposed_ports
            .first()
            .expect("Unable to get exposed port")
    }
}

fn cluster_file(port: ContainerPort, container: &Container<FdbImage>) -> NamedTempFile {
    let mut cluster_file = NamedTempFile::new().expect("Unable to create cluster file");

    let host = container
        .get_host()
        .expect("Unable to get container IP")
        .to_string();

    let port = container.get_host_port_ipv4(port).expect("");

    info!("{host}:{port}");

    cluster_file
        .write_all(format!("docker:docker@{host}:{port}").as_bytes())
        .expect("Unable to write cluster file");
    cluster_file.flush().expect("unable to flush cluster file");

    info!("{:?}", cluster_file.path());

    cluster_file
}

impl Default for FdbImage {
    fn default() -> Self {
        let port = ContainerPort::Tcp(thread_rng().gen_range(49152..65535));
        let env_vars = HashMap::from([
            ("FDB_PORT".to_string(), port.as_u16().to_string()),
            ("FDB_NETWORKING_MODE".to_string(), "host".to_string()),
        ]);

        Self {
            exposed_ports: vec![port],
            env_vars,
        }
    }
}

impl Image for FdbImage {
    fn name(&self) -> &str {
        IMAGE_NAME
    }

    fn tag(&self) -> &str {
        IMAGE_TAG
    }

    fn ready_conditions(&self) -> Vec<WaitFor> {
        vec![WaitFor::message_on_stdout("FDBD joined cluster.")]
    }

    fn env_vars(
        &self,
    ) -> impl IntoIterator<Item = (impl Into<Cow<'_, str>>, impl Into<Cow<'_, str>>)> {
        &self.env_vars
    }

    fn expose_ports(&self) -> &[ContainerPort] {
        &self.exposed_ports
    }

    fn exec_after_start(
        &self,
        _cs: ContainerState,
    ) -> Result<Vec<ExecCommand>, TestcontainersError> {
        Ok(vec![ExecCommand::new(vec![
            "bash",
            "-c",
            "\"\"/usr/bin/fdbcli -C /var/fdb/fdb.cluster --exec 'configure new single memory'\"\"",
        ])
        .with_cmd_ready_condition(CmdWaitFor::message_on_stdout(
            "Database created",
        ))])
    }
}

pub struct Details {
    database: Arc<Database>,
    _cluster_file: NamedTempFile,
    container: Option<Container<FdbImage>>,
}

pub struct FdbTestContainer {
    details: RwLock<Details>,
    _guard: Arc<NetworkAutoStop>,
    counter: Arc<()>,
}

impl FdbTestContainer {
    pub fn new(guard: Arc<NetworkAutoStop>, details: Details) -> Self {
        Self {
            details: RwLock::new(details),
            _guard: guard,
            counter: Default::default(),
        }
    }

    pub fn get(&self) -> DatabaseGuard {
        let details = self.details.read().expect("Unable to get details");
        if details.container.is_some() {
            let db = details.database.clone();
            return DatabaseGuard {
                test_container: self,
                database: db,
                _marker: Some(self.counter.clone()),
            };
        }

        // Explicitly release the lock on details
        drop(details);

        let new_details = create_test_container_loop().expect("Unable to create container");
        let db = new_details.database.clone();

        // Swap details
        *self.details.write().expect("Unable to write details") = new_details;

        DatabaseGuard {
            test_container: self,
            database: db,
            _marker: Some(self.counter.clone()),
        }
    }

    fn finish(&self) {
        if Arc::strong_count(&self.counter) == 1 {
            self.details
                .write()
                .expect("Unable to write container")
                .container
                .take();
        }
    }
}

pub struct DatabaseGuard<'a> {
    test_container: &'a FdbTestContainer,
    database: Arc<Database>,
    _marker: Option<Arc<()>>,
}

impl Drop for DatabaseGuard<'_> {
    fn drop(&mut self) {
        // Destroy the marker reference to allow container cleanup if needed
        self._marker.take();
        self.test_container.finish()
    }
}

impl Deref for DatabaseGuard<'_> {
    type Target = Arc<Database>;

    fn deref(&self) -> &Self::Target {
        &self.database
    }
}

#[derive(Debug, Error)]
enum Error {
    #[error("A test container error occurs : {0}")]
    TestContainer(#[from] TestcontainersError),
    #[error("Unable to create cluster file path")]
    ClusterFile,
}

fn create_fdb_test_container() -> Result<Details, Error> {
    let image = FdbImage::default();
    let port = image.get_port();
    let container = image.with_mapped_port(port.as_u16(), port).start()?;

    let cluster_file = cluster_file(port, &container);

    let cluster_file_path = cluster_file
        .path()
        .as_os_str()
        .to_str()
        .ok_or(Error::ClusterFile)?;

    info!("path => {cluster_file_path}");

    let database = Database::new(Some(cluster_file_path)).expect("Unable to create database");

    Ok(Details {
        database: Arc::new(database),
        _cluster_file: cluster_file,
        container: Some(container),
    })
}

fn create_test_container_loop() -> Result<Details, Error> {
    let mut counter = 3;

    loop {
        match create_fdb_test_container() {
            Ok(test_container) => {
                return Ok(test_container);
            }
            Err(error) => {
                error!("{error}");
                counter -= 1;
                if counter <= 0 {
                    panic!("Unable to start container after 3 attempts")
                }
            }
        }
    }
}

pub static DATABASE: LazyLock<FdbTestContainer> = LazyLock::new(|| {
    let guard = Arc::new(unsafe { foundationdb::boot() });

    ctrlc::set_handler(|| {
        info!("register signal called");
        DATABASE
            .details
            .write()
            .expect("Unable to write container")
            .container
            .take();
    })
    .expect("Unable to set signal handler");

    let details = create_test_container_loop().expect("Unable to create FDB test container");
    FdbTestContainer::new(guard, details)
});