Cargo.toml

[package]
name = "fbd_unit_test"
version = "0.1.0"
edition = "2021"

[dependencies]
foundationdb = { version = "0.9.0", features = ["embedded-fdb-include", "fdb-7_3"] }
smol = "2.0.2"
smol-potat = "1.1.2"
tokio = { version = "1.40.0", features = ["macros", "rt-multi-thread"] }
testcontainers = { version = "0.23.1", features = ["blocking"] }
rand = "0.8.5"
tempfile = "3.13.0"
log = "0.4.22"
env_logger = "0.11.5"
thiserror = "1.0.64"
ulid = "1.1.3"
ctrlc = "3.4.5"
async-lazy = { version = "0.1.0", features = ["parking_lot"] }
tokio-async-drop = "0.1.0"
async-trait = "0.1.83"
futures = "0.3.31"

common.rs

use crate::errors::Error;
use crate::fdb::FdbImage;
use foundationdb::Database;
use log::error;
use std::io::Write;
use std::sync::Arc;
use tempfile::NamedTempFile;
use testcontainers::core::ContainerPort;
use testcontainers::runners::AsyncRunner;
use testcontainers::{ContainerAsync, ImageExt};

pub(crate) struct Details {
    pub(crate) database: Arc<Database>,
    _cluster_file: NamedTempFile,
    pub(crate) container: Option<ContainerAsync<FdbImage>>,
}

/// Create an FDB cluster file from container information
async fn cluster_file(
    port: ContainerPort,
    container: &ContainerAsync<FdbImage>,
) -> Result<NamedTempFile, Error> {
    let mut cluster_file = NamedTempFile::new().expect("Unable to create cluster file");

    let host = container.get_host().await?.to_string();
    let port = container.get_host_port_ipv4(port).await?;

    cluster_file.write_all(format!("docker:docker@{host}:{port}").as_bytes())?;
    cluster_file.flush()?;

    Ok(cluster_file)
}

/// Create an FDB container and its context
async 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().await?;

    let cluster_file = cluster_file(port, &container).await?;

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

    let database = Database::new(Some(cluster_file_path))?;

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

/// Attempt to create a FDB container, retries 3 times before giving up
pub(crate) async fn create_test_container_loop() -> Details {
    let mut counter = 3;

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

errors.rs

use foundationdb::FdbError;
use log::error;
use std::io;
use testcontainers::TestcontainersError;
use thiserror::Error;

#[derive(Debug, Error)]
pub(crate) enum Error {
    #[error("A test container error occurs : {0}")]
    TestContainer(#[from] TestcontainersError),
    #[error("Unable to create cluster file path")]
    ClusterFile,
    #[error("FDB error occurred : {0}")]
    Fdb(#[from] FdbError),
    #[error("I/O error occurred: {0}")]
    Io(#[from] io::Error),
}

merged.rs

use crate::common::{create_test_container_loop, Details};
use foundationdb::api::NetworkAutoStop;
use foundationdb::Database;
use futures::FutureExt;
use std::future::Future;
use std::ops::Deref;
use std::panic::{resume_unwind, AssertUnwindSafe};
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio_async_drop::tokio_async_drop;

struct FdbTestContainer {
    details: RwLock<Details>,
    _guard: Arc<NetworkAutoStop>,
}

static COUNTER: AtomicU8 = AtomicU8::new(0);

async fn shutdown_container() {
    CONTEXT.force().await.details.write().await.container.take();
}

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

    async fn get(&self) -> Arc<Database> {
        let details = self.details.read().await;
        if details.container.is_some() {
            let db = details.database.clone();
            return db;
        }

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

        let new_details = create_test_container_loop().await;
        let db = new_details.database.clone();

        // Swap details
        *self.details.write().await = new_details;

        db
    }
}

pub struct DatabaseGuardOnce {
    details: Details,
}

impl Deref for DatabaseGuardOnce {
    type Target = Arc<Database>;

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

pub struct DatabaseGuard {
    database: Arc<Database>,
}

impl Drop for DatabaseGuard {
    fn drop(&mut self) {
        if COUNTER.fetch_sub(1, Ordering::Acquire) == 1 {
            tokio_async_drop!({ shutdown_container().await })
        }
    }
}

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

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

pub async fn get_db_context<F, Fut>(closure: F)
where
    F: FnOnce(Arc<Database>) -> Fut,
    Fut: Future<Output = ()>,
{
    COUNTER.fetch_add(1, Ordering::Release);
    let db = CONTEXT.force().await.get().await;

    let closure_result = AssertUnwindSafe(closure(db)).catch_unwind().await;

    if COUNTER.fetch_sub(1, Ordering::Acquire) == 1 {
        shutdown_container().await
    }

    if let Err(err) = closure_result {
        resume_unwind(err)
    }
}

pub async fn get_db() -> DatabaseGuard {
    if let tokio::runtime::RuntimeFlavor::CurrentThread =
        tokio::runtime::Handle::current().runtime_flavor()
    {
        shutdown_container().await;
        panic!(
            r#"Can't run FDB test container RAII in a mono_threaded flavor, please switch to #[tokio::test(flavor = "multi_thread")] "#
        )
    }

    COUNTER.fetch_add(1, Ordering::Release);
    let database = CONTEXT.force().await.get().await;
    DatabaseGuard { database }
}

pub async fn get_db_once() -> DatabaseGuardOnce {
    let details = create_test_container_loop().await;
    DatabaseGuardOnce { details }
}

static CONTEXT: async_lazy::Lazy<FdbTestContainer> = async_lazy::Lazy::const_new(|| {
    Box::pin(async {
        let guard = Arc::new(unsafe { foundationdb::boot() });
        let details = create_test_container_loop().await;
        FdbTestContainer::new(guard, details)
    })
});