feat(asset-canister): upload and commit in same message (#3954)

CHANGELOG.md

@@ -13,15 +13,20 @@ Valid settings are:
 ### feat: batch upload assets
 
 The frontend canister sync now tries to batch multiple small content chunks into a single call using the `create_chunks` method added earlier.
+And for small amounts of uploaded data the asset sync can now skip chunk creation entirely.
 This should lead to significantly faster upload times for frontends with many small files.
 
 ## Dependencies
 
 ### Frontend canister
 
+`SetAssetContentArguments` has a new field `last_chunk: opt blob` which can be used in addition to `chunk_ids` so that small assets can be uploaded as part of `commit_batch`,
+skipping the need to await a separate `create_chunk` call.
+
 Bumped `api_version` to `2` for the previous addition of `create_chunks` since the improved file sync relies on it.
 
-- Module hash: 9e4485d4358dd910aebcc025843547d05604cf28c6dc7c2cc2f8c76d083112e8
+- Module hash: 296d1ad1a7f8b15f90ff8b728658646b649cabd159f360f1b427297f4c76763e
+- https://github.com/dfinity/sdk/pull/3954
 - https://github.com/dfinity/sdk/pull/3947
 
 # 0.24.1

docs/design/asset-canister-interface.md

@@ -483,11 +483,14 @@ type SetAssetContentArguments = record {
   key: Key;
   content_encoding: text;
   chunk_ids: vec ChunkId;
+  last_chunk: opt blob;
   sha256: opt blob;
 };
 ```
 
 This operation adds or changes a single content encoding for an asset.  It also updates the modification time of the content encoding.
+The content of the encoding can be specified with `chunk_ids` and `last_chunk`.
+If `last_chunk` is not `null`, then its content is used as the last chunk of the encoding.
 
 If `sha256` is not passed, the asset canister will compute the hash of the content.
 

src/canisters/frontend/ic-asset/src/batch_upload/operations.rs

@@ -7,6 +7,7 @@ use crate::canister_api::types::batch_upload::common::{
     UnsetAssetContentArguments,
 };
 use crate::canister_api::types::batch_upload::v1::{BatchOperationKind, CommitBatchArguments};
+use crate::error::{AssembleCommitBatchArgumentError, SetEncodingError};
 use candid::Nat;
 use std::collections::HashMap;
 
@@ -20,7 +21,7 @@ pub(crate) async fn assemble_batch_operations(
     canister_assets: HashMap<String, AssetDetails>,
     asset_deletion_reason: AssetDeletionReason,
     canister_asset_properties: HashMap<String, AssetProperties>,
-) -> Vec<BatchOperationKind> {
+) -> Result<Vec<BatchOperationKind>, AssembleCommitBatchArgumentError> {
     let mut canister_assets = canister_assets;
 
     let mut operations = vec![];
@@ -33,10 +34,12 @@ pub(crate) async fn assemble_batch_operations(
     );
     create_new_assets(&mut operations, project_assets, &canister_assets);
     unset_obsolete_encodings(&mut operations, project_assets, &canister_assets);
-    set_encodings(&mut operations, chunk_uploader, project_assets).await;
+    set_encodings(&mut operations, chunk_uploader, project_assets)
+        .await
+        .map_err(AssembleCommitBatchArgumentError::SetEncodingFailed)?;
     update_properties(&mut operations, project_assets, &canister_asset_properties);
 
-    operations
+    Ok(operations)
 }
 
 pub(crate) async fn assemble_commit_batch_arguments(
@@ -46,19 +49,19 @@ pub(crate) async fn assemble_commit_batch_arguments(
     asset_deletion_reason: AssetDeletionReason,
     canister_asset_properties: HashMap<String, AssetProperties>,
     batch_id: Nat,
-) -> CommitBatchArguments {
+) -> Result<CommitBatchArguments, AssembleCommitBatchArgumentError> {
     let operations = assemble_batch_operations(
         Some(chunk_uploader),
         &project_assets,
         canister_assets,
         asset_deletion_reason,
         canister_asset_properties,
     )
-    .await;
-    CommitBatchArguments {
+    .await?;
+    Ok(CommitBatchArguments {
         operations,
         batch_id,
-    }
+    })
 }
 
 pub(crate) enum AssetDeletionReason {
@@ -163,29 +166,32 @@ pub(crate) async fn set_encodings(
     operations: &mut Vec<BatchOperationKind>,
     chunk_uploader: Option<&ChunkUploader<'_>>,
     project_assets: &HashMap<String, ProjectAsset>,
-) {
+) -> Result<(), SetEncodingError> {
     for (key, project_asset) in project_assets {
         for (content_encoding, v) in &project_asset.encodings {
             if v.already_in_place {
                 continue;
             }
-            let chunk_ids = if let Some(uploader) = chunk_uploader {
-                uploader
-                    .uploader_ids_to_canister_chunk_ids(&v.uploader_chunk_ids)
-                    .await
-            } else {
-                vec![]
+            let (chunk_ids, last_chunk) = match chunk_uploader {
+                Some(uploader) => {
+                    uploader
+                        .uploader_ids_to_canister_chunk_ids(&v.uploader_chunk_ids)
+                        .await?
+                }
+                None => (vec![], None),
             };
             operations.push(BatchOperationKind::SetAssetContent(
                 SetAssetContentArguments {
                     key: key.clone(),
                     content_encoding: content_encoding.clone(),
                     chunk_ids,
+                    last_chunk,
                     sha256: Some(v.sha256.clone()),
                 },
             ));
         }
     }
+    Ok(())
 }
 
 pub(crate) fn update_properties(

src/canisters/frontend/ic-asset/src/batch_upload/plumbing.rs

@@ -9,6 +9,7 @@ use crate::error::CreateChunkError;
 use crate::error::CreateEncodingError;
 use crate::error::CreateEncodingError::EncodeContentFailed;
 use crate::error::CreateProjectAssetError;
+use crate::error::SetEncodingError;
 use candid::Nat;
 use futures::future::try_join_all;
 use futures::TryFutureExt;
@@ -99,8 +100,7 @@ impl<'agent> ChunkUploader<'agent> {
             //  - Tested with: `for i in $(seq 1 50); do dd if=/dev/urandom of="src/hello_frontend/assets/file_$i.bin" bs=$(shuf -i 1-2000000 -n 1) count=1; done && dfx deploy hello_frontend`
             //  - Result: Roughly 15% of batches under 90% full.
             // With other byte ranges (e.g. `shuf -i 1-3000000 -n 1`) stats improve significantly
-            self.upload_chunks(4 * MAX_CHUNK_SIZE, usize::MAX, semaphores)
-                .await?;
+            self.upload_chunks(4 * MAX_CHUNK_SIZE, semaphores).await?;
             Ok(uploader_chunk_id)
         }
     }
@@ -109,8 +109,8 @@ impl<'agent> ChunkUploader<'agent> {
         &self,
         semaphores: &Semaphores,
     ) -> Result<(), CreateChunkError> {
-        self.upload_chunks(0, 0, semaphores).await?;
-        Ok(())
+        // Crude estimate: If `MAX_CHUNK_SIZE / 2` bytes are added as data to the `commit_batch` args the message won't be above the message size limit.
+        self.upload_chunks(MAX_CHUNK_SIZE / 2, semaphores).await
     }
 
     pub(crate) fn bytes(&self) -> usize {
@@ -120,21 +120,32 @@ impl<'agent> ChunkUploader<'agent> {
         self.chunks.load(Ordering::SeqCst)
     }
 
-    /// Call only after `finalize_upload` has completed
+    /// Call only after `finalize_upload` has completed.
     pub(crate) async fn uploader_ids_to_canister_chunk_ids(
         &self,
         uploader_ids: &[usize],
-    ) -> Vec<Nat> {
+    ) -> Result<(Vec<Nat>, Option<Vec<u8>>), SetEncodingError> {
+        let mut chunk_ids = vec![];
+        let mut last_chunk: Option<Vec<u8>> = None;
         let mapping = self.id_mapping.lock().await;
-        uploader_ids
-            .iter()
-            .map(|id| {
-                mapping
-                    .get(id)
-                    .expect("Chunk uploader did not upload all chunks. This is a bug.")
-                    .clone()
-            })
-            .collect()
+        let queue = self.upload_queue.lock().await;
+        for uploader_id in uploader_ids {
+            if let Some(item) = mapping.get(uploader_id) {
+                chunk_ids.push(item.clone());
+            } else if let Some(last_chunk_data) =
+                queue
+                    .iter()
+                    .find_map(|(id, data)| if id == uploader_id { Some(data) } else { None })
+            {
+                match last_chunk.as_mut() {
+                    Some(existing_data) => existing_data.extend(last_chunk_data.iter()),
+                    None => last_chunk = Some(last_chunk_data.clone()),
+                }
+            } else {
+                return Err(SetEncodingError::UnknownUploaderChunkId(*uploader_id));
+            }
+        }
+        Ok((chunk_ids, last_chunk))
     }
 
     async fn add_to_upload_queue(&self, uploader_chunk_id: usize, contents: &[u8]) {
@@ -148,7 +159,6 @@ impl<'agent> ChunkUploader<'agent> {
     async fn upload_chunks(
         &self,
         max_retained_bytes: usize,
-        max_retained_chunks: usize,
         semaphores: &Semaphores,
     ) -> Result<(), CreateChunkError> {
         let mut queue = self.upload_queue.lock().await;
@@ -159,7 +169,6 @@ impl<'agent> ChunkUploader<'agent> {
             .map(|(_, content)| content.len())
             .sum::<usize>()
             > max_retained_bytes
-            || queue.len() > max_retained_chunks
         {
             // Greedily fills batch with the largest chunk that fits
             queue.sort_unstable_by_key(|(_, content)| content.len());

src/canisters/frontend/ic-asset/src/canister_api/types/batch_upload/common.rs

@@ -77,8 +77,10 @@ pub struct SetAssetContentArguments {
     pub key: String,
     /// The content encoding for which this content applies
     pub content_encoding: String,
-    /// The chunks to assign to this content
+    /// The chunks to assign to this content encoding
     pub chunk_ids: Vec<Nat>,
+    /// Appends this chunk to the data supplied in `chunk_ids`
+    pub last_chunk: Option<Vec<u8>>,
     /// The sha256 of the entire content
     pub sha256: Option<Vec<u8>>,
 }

src/canisters/frontend/ic-asset/src/error/assemble_commit_batch_argument.rs

@@ -0,0 +1,11 @@
+use thiserror::Error;
+
+use super::SetEncodingError;
+
+/// Errors related to assembling commit_batch arguments for the asset canister.
+#[derive(Error, Debug)]
+pub enum AssembleCommitBatchArgumentError {
+    /// Failed to set encoding.
+    #[error("Failed to set encoding: {0}")]
+    SetEncodingFailed(#[from] SetEncodingError),
+}

src/canisters/frontend/ic-asset/src/error/compute_evidence.rs

@@ -5,9 +5,15 @@ use crate::error::hash_content::HashContentError;
 use ic_agent::AgentError;
 use thiserror::Error;
 
+use super::AssembleCommitBatchArgumentError;
+
 /// Errors related to computing evidence for a proposed update.
 #[derive(Error, Debug)]
 pub enum ComputeEvidenceError {
+    /// Failed when assembling commit_batch argument.
+    #[error(transparent)]
+    AssembleCommitBatchArgumentFailed(#[from] AssembleCommitBatchArgumentError),
+
     /// Failed when inspecting assets to be updated.
     #[error(transparent)]
     ProcessProjectAsset(#[from] CreateProjectAssetError),

src/canisters/frontend/ic-asset/src/error/mod.rs

@@ -1,5 +1,6 @@
 //! Error types
 
+mod assemble_commit_batch_argument;
 mod compatibility;
 mod compute_evidence;
 mod create_chunk;
@@ -13,10 +14,12 @@ mod hash_content;
 mod load_config;
 mod load_rule;
 mod prepare_sync_for_proposal;
+mod set_encoding;
 mod sync;
 mod upload;
 mod upload_content;
 
+pub use assemble_commit_batch_argument::AssembleCommitBatchArgumentError;
 pub use compatibility::CompatibilityError;
 pub use compute_evidence::ComputeEvidenceError;
 pub use create_chunk::CreateChunkError;
@@ -30,6 +33,7 @@ pub use hash_content::HashContentError;
 pub use load_config::AssetLoadConfigError;
 pub use load_rule::LoadRuleError;
 pub use prepare_sync_for_proposal::PrepareSyncForProposalError;
+pub use set_encoding::SetEncodingError;
 pub use sync::SyncError;
 pub use upload::UploadError;
 pub use upload_content::UploadContentError;

src/canisters/frontend/ic-asset/src/error/set_encoding.rs

@@ -0,0 +1,9 @@
+use thiserror::Error;
+
+/// Errors encountered while setting the encodings.
+#[derive(Error, Debug)]
+pub enum SetEncodingError {
+    /// Failed when attempting to translate an uploader chunk id to a canister chunk id because the id is unknown.
+    #[error("Unknown uploader chunk id: {0}")]
+    UnknownUploaderChunkId(usize),
+}

src/canisters/frontend/ic-asset/src/error/upload.rs

@@ -1,3 +1,4 @@
+use super::AssembleCommitBatchArgumentError;
 use crate::error::compatibility::CompatibilityError;
 use crate::error::create_project_asset::CreateProjectAssetError;
 use ic_agent::AgentError;
@@ -18,6 +19,10 @@ pub enum UploadError {
     #[error("Create batch failed: {0}")]
     CreateBatchFailed(AgentError),
 
+    /// Failed when assembling commit_batch argument.
+    #[error("Failed to assemble commit_batch argument: {0}")]
+    AssembleCommitBatchArgumentFailed(#[from] AssembleCommitBatchArgumentError),
+
     /// Failed when creating project assets.
     #[error("Failed to create project asset: {0}")]
     CreateProjectAssetFailed(#[from] CreateProjectAssetError),

src/canisters/frontend/ic-asset/src/error/upload_content.rs

@@ -4,9 +4,15 @@ use crate::error::get_asset_properties::GetAssetPropertiesError;
 use ic_agent::AgentError;
 use thiserror::Error;
 
+use super::AssembleCommitBatchArgumentError;
+
 /// Errors related to uploading content to the asset canister.
 #[derive(Error, Debug)]
 pub enum UploadContentError {
+    /// Failed when assembling commit_batch argument.
+    #[error("Failed to assemble commit_batch argument: {0}")]
+    AssembleCommitBatchArgumentFailed(AssembleCommitBatchArgumentError),
+
     /// Failed when calling create_batch.
     #[error("Failed to create batch: {0}")]
     CreateBatchFailed(AgentError),

src/canisters/frontend/ic-asset/src/evidence/mod.rs

@@ -65,7 +65,8 @@ pub async fn compute_evidence(
         Obsolete,
         canister_asset_properties,
     )
-    .await;
+    .await
+    .map_err(ComputeEvidenceError::AssembleCommitBatchArgumentFailed)?;
     operations.sort();
 
     let mut sha = Sha256::new();

src/canisters/frontend/ic-asset/src/sync.rs

@@ -84,7 +84,8 @@ pub async fn upload_content_and_assemble_sync_operations(
         &canister_assets,
         logger,
     )
-    .await?;
+    .await
+    .map_err(UploadContentError::CreateProjectAssetError)?;
 
     let commit_batch_args = batch_upload::operations::assemble_commit_batch_arguments(
         &chunk_uploader,
@@ -97,7 +98,8 @@ pub async fn upload_content_and_assemble_sync_operations(
         canister_asset_properties,
         batch_id,
     )
-    .await;
+    .await
+    .map_err(UploadContentError::AssembleCommitBatchArgumentFailed)?;
 
     // -v
     debug!(

src/canisters/frontend/ic-asset/src/upload.rs

@@ -12,8 +12,7 @@ use crate::canister_api::methods::{
 };
 use crate::canister_api::types::batch_upload::v0;
 use crate::error::CompatibilityError::DowngradeV1TOV0Failed;
-use crate::error::UploadError;
-use crate::error::UploadError::{CommitBatchFailed, CreateBatchFailed, ListAssetsFailed};
+use crate::error::UploadError::{self, CommitBatchFailed, CreateBatchFailed, ListAssetsFailed};
 use ic_utils::Canister;
 use slog::{info, Logger};
 use std::collections::HashMap;
@@ -62,7 +61,8 @@ pub async fn upload(
         HashMap::new(),
         batch_id,
     )
-    .await;
+    .await
+    .map_err(UploadError::AssembleCommitBatchArgumentFailed)?;
 
     let canister_api_version = api_version(canister).await;
     info!(logger, "Committing batch.");