canic-control-plane 0.93.19

use crate::{
    dto::template::{
        TemplateChunkResponse, TemplateChunkSetInfoResponse, TemplateChunkSetPrepareInput,
        TemplateManifestInput, WasmStoreCatalogEntryResponse, WasmStoreStatusResponse,
    },
    ids::{TemplateId, TemplateVersion},
};
use candid::{CandidType, utils::ArgumentEncoder};
use canic_core::cdk::types::Principal;
use canic_core::{
    control_plane_support::{
        error::InternalError,
        ops::{cost_guard::CostGuardPermit, ic::call::CallOps},
    },
    dto::error::Error,
    protocol,
};

///
/// WasmStoreInternalClient
///
pub(in crate::workflow::runtime::template) struct WasmStoreInternalClient {
    store_pid: Principal,
}

impl WasmStoreInternalClient {
    const BEGIN_GC: &str = protocol::CANIC_WASM_STORE_BEGIN_GC;
    const CATALOG: &str = protocol::CANIC_WASM_STORE_CATALOG;
    const CHUNK: &str = protocol::CANIC_WASM_STORE_CHUNK;
    const COMPLETE_GC: &str = protocol::CANIC_WASM_STORE_COMPLETE_GC;
    const INFO: &str = protocol::CANIC_WASM_STORE_INFO;
    const PREPARE: &str = protocol::CANIC_WASM_STORE_PREPARE;
    const PREPARE_GC: &str = protocol::CANIC_WASM_STORE_PREPARE_GC;
    const PUBLISH_CHUNK: &str = protocol::CANIC_WASM_STORE_PUBLISH_CHUNK;
    const STAGE_MANIFEST: &str = protocol::CANIC_WASM_STORE_STAGE_MANIFEST;
    const STATUS: &str = protocol::CANIC_WASM_STORE_STATUS;
    #[cfg(test)]
    const ENDPOINTS: &[&str] = &[
        Self::BEGIN_GC,
        Self::CATALOG,
        Self::CHUNK,
        Self::COMPLETE_GC,
        Self::INFO,
        Self::PREPARE,
        Self::PREPARE_GC,
        Self::PUBLISH_CHUNK,
        Self::STAGE_MANIFEST,
        Self::STATUS,
    ];

    pub(super) const fn new(store_pid: Principal) -> Self {
        Self { store_pid }
    }

    pub(super) async fn catalog(
        &self,
    ) -> Result<Vec<WasmStoreCatalogEntryResponse>, InternalError> {
        self.call_result(Self::CATALOG, ()).await
    }

    pub(super) async fn info(
        &self,
        template_id: &TemplateId,
        version: &TemplateVersion,
    ) -> Result<TemplateChunkSetInfoResponse, InternalError> {
        self.call_result(
            Self::INFO,
            (
                template_id.as_str().to_string(),
                version.as_str().to_string(),
            ),
        )
        .await
    }

    pub(super) async fn status(&self) -> Result<WasmStoreStatusResponse, InternalError> {
        self.call_result(Self::STATUS, ()).await
    }

    pub(super) async fn prepare_chunk_set(
        &self,
        _publication_permit: &CostGuardPermit,
        request: TemplateChunkSetPrepareInput,
    ) -> Result<TemplateChunkSetInfoResponse, InternalError> {
        self.call_result(Self::PREPARE, (request,)).await
    }

    pub(super) async fn stage_manifest(
        &self,
        _publication_permit: &CostGuardPermit,
        request: TemplateManifestInput,
    ) -> Result<(), InternalError> {
        self.call_result(Self::STAGE_MANIFEST, (request,)).await
    }

    pub(super) async fn publish_chunk(
        &self,
        _publication_permit: &CostGuardPermit,
        template_id: &TemplateId,
        version: &TemplateVersion,
        chunk_index: u32,
        bytes: &[u8],
    ) -> Result<(), InternalError> {
        self.call_result(
            Self::PUBLISH_CHUNK,
            (TemplateChunkInputRef {
                template_id,
                version,
                chunk_index,
                bytes,
            },),
        )
        .await
    }

    pub(super) async fn prepare_gc(&self) -> Result<(), InternalError> {
        self.call_result(Self::PREPARE_GC, ()).await
    }

    pub(super) async fn begin_gc(&self) -> Result<(), InternalError> {
        self.call_result(Self::BEGIN_GC, ()).await
    }

    pub(super) async fn complete_gc(&self) -> Result<(), InternalError> {
        self.call_result(Self::COMPLETE_GC, ()).await
    }

    pub(super) async fn chunk(
        &self,
        template_id: &TemplateId,
        version: &TemplateVersion,
        chunk_index: u32,
    ) -> Result<Vec<u8>, InternalError> {
        let response: TemplateChunkResponse = self
            .call_result(
                Self::CHUNK,
                (
                    template_id.as_str().to_string(),
                    version.as_str().to_string(),
                    chunk_index,
                ),
            )
            .await?;

        Ok(response.bytes)
    }

    async fn call_result<T, A>(&self, method: &'static str, arg: A) -> Result<T, InternalError>
    where
        T: CandidType + serde::de::DeserializeOwned,
        A: ArgumentEncoder,
    {
        let call = CallOps::bounded_wait(self.store_pid, method)
            .with_args(arg)
            .map_err(|err| InternalError::public(Error::invariant(err.to_string())))?
            .execute()
            .await
            .map_err(|err| InternalError::public(Error::unavailable(err.to_string())))?;
        let call_res: Result<T, Error> = call
            .candid::<Result<T, Error>>()
            .map_err(|err| InternalError::public(Error::invariant(err.to_string())))?;

        call_res.map_err(InternalError::public)
    }
}

// Borrowed chunk publish input for store-side chunk staging.
#[derive(CandidType)]
struct TemplateChunkInputRef<'a> {
    pub template_id: &'a TemplateId,
    pub version: &'a TemplateVersion,
    pub chunk_index: u32,
    pub bytes: &'a [u8],
}

#[cfg(test)]
mod tests {
    use super::WasmStoreInternalClient;
    use canic_core::protocol;
    use std::collections::BTreeSet;

    #[test]
    fn typed_client_endpoint_table_matches_protocol_manifests() {
        let root_updates = WasmStoreInternalClient::ENDPOINTS
            .iter()
            .filter(|method| protocol::CANIC_WASM_STORE_ROOT_UPDATE_METHODS.contains(method))
            .copied()
            .collect::<BTreeSet<_>>();
        let structural = WasmStoreInternalClient::ENDPOINTS
            .iter()
            .filter(|method| protocol::CANIC_WASM_STORE_STRUCTURAL_QUERY_METHODS.contains(method))
            .copied()
            .collect::<BTreeSet<_>>();
        let all = WasmStoreInternalClient::ENDPOINTS
            .iter()
            .copied()
            .collect::<BTreeSet<_>>();

        assert_eq!(
            root_updates,
            protocol::CANIC_WASM_STORE_ROOT_UPDATE_METHODS
                .iter()
                .copied()
                .collect::<BTreeSet<_>>()
        );
        assert_eq!(
            structural,
            protocol::CANIC_WASM_STORE_STRUCTURAL_QUERY_METHODS
                .iter()
                .copied()
                .collect::<BTreeSet<_>>()
        );
        assert_eq!(
            all.len(),
            WasmStoreInternalClient::ENDPOINTS.len(),
            "typed wasm-store client endpoint methods must be unique"
        );
    }
}