dagger-components/crates/cuddle-ci/src/rust_service.rs

use std::{path::PathBuf, sync::Arc};

use async_trait::async_trait;
use dagger_rust::source::RustSource;
use dagger_sdk::Container;
use futures::{stream, StreamExt};

use crate::{
    dagger_middleware::{DaggerMiddleware, DynMiddleware},
    MainAction, PullRequestAction,
};

use self::architecture::{Architecture, Os};

#[derive(Clone)]
pub enum RustServiceStage {
    BeforeDeps(DynMiddleware),
    AfterDeps(DynMiddleware),
    BeforeBase(DynMiddleware),
    AfterBase(DynMiddleware),
    BeforeBuild(DynMiddleware),
    AfterBuild(DynMiddleware),
    BeforePackage(DynMiddleware),
    AfterPackage(DynMiddleware),
    BeforeRelease(DynMiddleware),
    AfterRelease(DynMiddleware),
}

#[derive(Clone)]
pub struct RustService {
    client: dagger_sdk::Query,
    base_image: Option<dagger_sdk::Container>,
    final_image: Option<dagger_sdk::Container>,
    stages: Vec<RustServiceStage>,
    source: Option<PathBuf>,
    crates: Vec<String>,
    bin_name: String,
    arch: Option<Architecture>,
    os: Option<Os>,
}

impl From<dagger_sdk::Query> for RustService {
    fn from(value: dagger_sdk::Query) -> Self {
        Self {
            client: value,
            base_image: None,
            final_image: None,
            stages: Vec::new(),
            source: None,
            crates: Vec::new(),
            bin_name: String::new(),
            arch: None,
            os: None,
        }
    }
}

impl RustService {
    pub async fn new() -> eyre::Result<Self> {
        Ok(Self {
            client: dagger_sdk::connect().await?,
            base_image: None,
            final_image: None,
            stages: Vec::new(),
            source: None,
            crates: Vec::new(),
            bin_name: String::new(),
            arch: None,
            os: None,
        })
    }

    pub fn with_base_image(&mut self, base: dagger_sdk::Container) -> &mut Self {
        self.base_image = Some(base);

        self
    }

    pub fn with_stage(&mut self, stage: RustServiceStage) -> &mut Self {
        self.stages.push(stage);

        self
    }

    pub fn with_source(&mut self, path: impl Into<PathBuf>) -> &mut Self {
        self.source = Some(path.into());

        self
    }

    pub fn with_bin_name(&mut self, bin_name: impl Into<String>) -> &mut Self {
        self.bin_name = bin_name.into();

        self
    }

    pub fn with_crates(
        &mut self,
        crates: impl IntoIterator<Item = impl Into<String>>,
    ) -> &mut Self {
        self.crates = crates.into_iter().map(|c| c.into()).collect();

        self
    }

    pub fn with_arch(&mut self, arch: Architecture) -> &mut Self {
        self.arch = Some(arch);

        self
    }

    pub fn with_os(&mut self, os: Os) -> &mut Self {
        self.os = Some(os);

        self
    }

    fn get_src(&self) -> PathBuf {
        self.source
            .clone()
            .unwrap_or(std::env::current_dir().unwrap())
    }

    fn get_arch(&self) -> Architecture {
        self.arch
            .clone()
            .unwrap_or_else(|| match std::env::consts::ARCH {
                "x86" | "x86_64" | "amd64" => Architecture::Amd64,
                "arm" => Architecture::Arm64,
                arch => panic!("unsupported architecture: {arch}"),
            })
    }

    fn get_os(&self) -> Os {
        self.os
            .clone()
            .unwrap_or_else(|| match std::env::consts::OS {
                "linux" => Os::Linux,
                "macos" => Os::MacOS,
                os => panic!("unsupported os: {os}"),
            })
    }

    async fn run_stage(
        &self,
        stages: impl IntoIterator<Item = &Arc<dyn DaggerMiddleware + Send + Sync>>,
        container: Container,
    ) -> eyre::Result<Container> {
        let before_deps_stream = stream::iter(stages.into_iter().map(Ok));
        let res = StreamExt::fold(before_deps_stream, Ok(container), |base, m| async move {
            match (base, m) {
                (Ok(base), Ok(m)) => m.handle(base).await,
                (_, Err(e)) | (Err(e), _) => eyre::bail!("failed with {e}"),
            }
        })
        .await?;

        Ok(res)
    }

    pub async fn build_base(&self) -> eyre::Result<Container> {
        let rust_src = RustSource::new(self.client.clone());

        let (src, dep_src) = rust_src
            .get_rust_src(Some(&self.get_src()), self.crates.clone())
            .await?;

        let base_image = self
            .base_image
            .clone()
            .unwrap_or(self.client.container().from("rustlang/rust:nightly"));

        let before_deps = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::BeforeDeps(middleware) => Some(middleware),
                _ => None,
            })
            .collect::<Vec<_>>();
        let image = self.run_stage(before_deps, base_image).await?;

        let after_deps = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::AfterDeps(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let image = self.run_stage(after_deps, image).await?;

        let before_base = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::BeforeBase(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let image = self.run_stage(before_base, image).await?;

        let cache = self.client.cache_volume("rust_target_cache");

        let rust_prebuild = image
            .with_workdir("/mnt/src")
            .with_directory("/mnt/src", dep_src)
            .with_exec(vec!["cargo", "build", "--release", "--bin", &self.bin_name])
            .with_mounted_cache("/mnt/src/target/", cache);

        let incremental_dir = rust_src
            .get_rust_target_src(&self.get_src(), rust_prebuild.clone(), self.crates.clone())
            .await?;

        let rust_with_src = image
            .with_workdir("/mnt/src")
            .with_directory(
                "/usr/local/cargo",
                rust_prebuild.directory("/usr/local/cargo"),
            )
            .with_directory("/mnt/src/target", incremental_dir)
            .with_directory("/mnt/src/", src);

        let after_base = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::AfterBase(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let image = self.run_stage(after_base, rust_with_src).await?;

        Ok(image)
    }

    pub async fn build_release(&self) -> eyre::Result<Container> {
        let base = self.build_base().await?;

        let before_build = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::BeforeBuild(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let base = self.run_stage(before_build, base).await?;

        let binary_build =
            base.with_exec(vec!["cargo", "build", "--release", "--bin", &self.bin_name]);

        let after_build = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::AfterBuild(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let binary_build = self.run_stage(after_build, binary_build).await?;

        let dest = self
            .final_image
            .clone()
            .unwrap_or(self.client.container().from("debian:bullseye"));

        let before_package = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::BeforePackage(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let dest = self.run_stage(before_package, dest).await?;

        let final_image = dest.with_workdir("/mnt/app").with_file(
            format!("/usr/local/bin/{}", self.bin_name),
            binary_build.file(format!("/mnt/src/target/release/{}", self.bin_name)),
        );

        let after_package = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::AfterPackage(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let final_image = self.run_stage(after_package, final_image).await?;

        Ok(final_image)
    }

    pub async fn build_test(&self) -> eyre::Result<()> {
        let base = self.build_base().await?;

        let before_build = self
            .stages
            .iter()
            .filter_map(|s| match s {
                RustServiceStage::BeforeBuild(m) => Some(m),
                _ => None,
            })
            .collect::<Vec<_>>();
        let base = self.run_stage(before_build, base).await?;

        base.with_exec(vec!["cargo", "test", "--release"])
            .sync()
            .await?;

        Ok(())
    }
}

#[async_trait]
impl PullRequestAction for RustService {
    async fn execute_pull_request(&self) -> eyre::Result<()> {
        self.build_test().await?;

        Ok(())
    }
}

#[async_trait]
impl MainAction for RustService {
    async fn execute_main(&self) -> eyre::Result<()> {
        let container = self.build_release().await?;

        container
            .publish(format!("docker.io//{}", self.bin_name))
            .await?;

        Ok(())
    }
}

pub mod architecture {
    #[derive(Debug, Clone)]
    pub enum Architecture {
        Amd64,
        Arm64,
    }

    #[derive(Debug, Clone)]
    pub enum Os {
        Linux,
        MacOS,
    }
}

mod apt;
mod cargo_binstall;
mod cargo_clean;
mod clap_sanity_test;
mod mold;
mod sqlx;

pub mod extensions {
    pub use super::apt::*;
    pub use super::cargo_binstall::*;
    pub use super::cargo_clean::*;
    pub use super::clap_sanity_test::*;
    pub use super::mold::*;
    pub use super::sqlx::*;
}

#[cfg(test)]
mod test {
    use futures::FutureExt;

    use crate::{
        dagger_middleware::middleware,
        rust_service::{
            apt::AptExt,
            architecture::{Architecture, Os},
            cargo_binstall::CargoBInstallExt,
            clap_sanity_test::ClapSanityTestExt,
            mold::MoldActionExt,
            RustService, RustServiceStage,
        },
    };

    #[tokio::test]
    async fn test_can_build_rust() -> eyre::Result<()> {
        let client = dagger_sdk::connect().await?;

        let root_dir = std::path::PathBuf::from("../../").canonicalize()?;

        let container = RustService::from(client.clone())
            .with_arch(Architecture::Amd64)
            .with_os(Os::Linux)
            .with_source(root_dir)
            .with_bin_name("ci")
            .with_crates(["crates/*", "examples/*", "ci"])
            .with_apt(&["git"])
            .with_cargo_binstall("latest", ["sqlx-cli"])
            .with_mold("2.3.3")
            .with_stage(RustServiceStage::BeforeDeps(middleware(|c| {
                async move {
                    // Noop
                    Ok(c)
                }
                .boxed()
            })))
            .with_clap_sanity_test()
            .build_release()
            .await?;

        container.sync().await?;

        Ok(())
    }
}