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},
    Context, 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>,
    deployment: bool,
}

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

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,
            deployment: true,
        })
    }

    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
    }

    pub fn with_deployment(&mut self, deployment: bool) -> &mut Self {
        self.deployment = deployment;

        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 client = self.client.pipeline("build-base");
        let rust_src = RustSource::new(client.pipeline("load-source"));

        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(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 = 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?.pipeline("build-release");
        let client = self.client.pipeline("build-release");

        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(client.container().from("debian:bookworm"));

        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?.pipeline("build-test");

        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, _ctx: &mut Context) -> eyre::Result<()> {
        self.build_test().await?;

        Ok(())
    }
}

const IMAGE_TAG: &str = "RUST_SERVICE_IMAGE_TAG";

pub trait RustServiceContext {
    fn set_image_tag(&mut self, tag: impl Into<String>) -> eyre::Result<()>;
    fn get_image_tag(&self) -> eyre::Result<Option<String>>;
}

impl RustServiceContext for Context {
    fn get_image_tag(&self) -> eyre::Result<Option<String>> {
        Ok(self.get(IMAGE_TAG).cloned())
    }

    fn set_image_tag(&mut self, tag: impl Into<String>) -> eyre::Result<()> {
        let tag = tag.into();

        self.insert(IMAGE_TAG.to_string(), tag);

        Ok(())
    }
}

#[async_trait]
impl MainAction for RustService {
    async fn execute_main(&self, ctx: &mut Context) -> eyre::Result<()> {
        let container = self.build_release().await?.pipeline("main");
        let timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let tag = format!(
            "docker.io/kasperhermansen/{}:main-{}",
            self.bin_name, timestamp,
        );

        container.publish(&tag).await?;
        ctx.set_image_tag(format!("main-{}", &timestamp.to_string()))?;

        if self.deployment {
            let update_deployments_docker_image =
                "docker.io/kasperhermansen/update-deployment:1701123940";
            let dep = self
                .client
                .container()
                .from(update_deployments_docker_image);

            let dep = if let Ok(sock) = std::env::var("SSH_AUTH_SOCK") {
                dep.with_unix_socket("/tmp/ssh_sock", self.client.host().unix_socket(sock))
                    .with_env_variable("SSH_AUTH_SOCK", "/tmp/ssh_sock")
                    .with_exec(vec![
                        "update-deployment",
                        "--repo",
                        &format!(
                            "git@git.front.kjuulh.io:kjuulh/{}-deployment.git",
                            self.bin_name
                        ),
                        "--service",
                        &self.bin_name,
                        "--image",
                        &format!("kasperhermansen/{}:main-{}", self.bin_name, timestamp),
                    ])
            } else {
                dep.with_env_variable("GIT_USERNAME", "kjuulh")
                    .with_env_variable(
                        "GIT_PASSWORD",
                        std::env::var("GIT_PASSWORD").expect("GIT_PASSWORD to be set"),
                    )
                    .with_exec(vec![
                        "update-deployment",
                        "--repo",
                        &format!(
                            "https://git.front.kjuulh.io/kjuulh/{}-deployment.git",
                            self.bin_name
                        ),
                        "--service",
                        &self.bin_name,
                        "--image",
                        &format!("kasperhermansen/{}:main-{}", self.bin_name, timestamp),
                    ])
            };

            dep.sync().await?;
        }
        Ok(())
    }
}

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

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

mod apt;
mod apt_ca_certificates;
mod assets;
mod cargo_binstall;
mod cargo_clean;
mod clap_sanity_test;
mod cuddle_cli;
mod cuddle_file;
mod dagger_bin;
mod docker_cache;
mod docker_cli;
mod kubectl;
mod mold;
mod rust_workspace;
mod sqlx;
mod ssh_agent;

pub mod extensions {
    pub use super::apt::*;
    pub use super::apt_ca_certificates::*;
    pub use super::assets::*;
    pub use super::cargo_binstall::*;
    pub use super::cargo_clean::*;
    pub use super::clap_sanity_test::*;
    pub use super::cuddle_cli::*;
    pub use super::cuddle_file::*;
    pub use super::dagger_bin::*;
    pub use super::docker_cache::*;
    pub use super::docker_cli::*;
    pub use super::kubectl::*;
    pub use super::mold::*;
    pub use super::rust_workspace::*;
    pub use super::sqlx::*;
    pub use super::ssh_agent::*;
}

#[cfg(test)]
mod test {

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

    #[tokio::test]
    #[cfg(any(feature = "dagger", feature = "integration"))]
    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(())
    }
}