fix: better asset canister upload chunking (#4442)

Author: viviveevee

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 # UNRELEASED
 
+### Improve frontend canister sync logic
+
+Previously, committing frontend canister changes happened in multiple batches defined by simple heuristics that would likely not exceed the ingress message size limit.
+Now, the ingress message size limit is respected more explicitly, and also a limit of total content size per batch since all content in the batch newly gets hashed in the canister.
+
 ## Dependencies
 
 ### Motoko

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

@@ -59,6 +59,7 @@ pub enum Mode {
 
 type IdMapping = BTreeMap<usize, Nat>;
 type UploadQueue = Vec<(usize, Vec<u8>)>;
+type CanisterChunkSizeMap = HashMap<Nat, usize>;
 pub(crate) struct ChunkUploader<'agent> {
     canister: Canister<'agent>,
     batch_id: Nat,
@@ -68,6 +69,8 @@ pub(crate) struct ChunkUploader<'agent> {
     // maps uploader_chunk_id to canister_chunk_id
     id_mapping: Arc<Mutex<IdMapping>>,
     upload_queue: Arc<Mutex<UploadQueue>>,
+    // maps canister_chunk_id to chunk size
+    canister_chunk_sizes: Arc<Mutex<CanisterChunkSizeMap>>,
 }
 
 impl<'agent> ChunkUploader<'agent> {
@@ -80,6 +83,7 @@ impl<'agent> ChunkUploader<'agent> {
             bytes: Arc::new(AtomicUsize::new(0)),
             id_mapping: Arc::new(Mutex::new(BTreeMap::new())),
             upload_queue: Arc::new(Mutex::new(vec![])),
+            canister_chunk_sizes: Arc::new(Mutex::new(HashMap::new())),
         }
     }
 
@@ -93,8 +97,10 @@ impl<'agent> ChunkUploader<'agent> {
         progress: Option<&dyn AssetSyncProgressRenderer>,
     ) -> Result<usize, CreateChunkError> {
         let uploader_chunk_id = self.chunks.fetch_add(1, Ordering::SeqCst);
-        self.bytes.fetch_add(contents.len(), Ordering::SeqCst);
-        if contents.len() == MAX_CHUNK_SIZE || self.api_version < 2 {
+        let chunk_size = contents.len();
+        self.bytes.fetch_add(chunk_size, Ordering::SeqCst);
+
+        if chunk_size == MAX_CHUNK_SIZE || self.api_version < 2 {
             let canister_chunk_id = create_chunk(
                 &self.canister,
                 &self.batch_id,
@@ -103,8 +109,15 @@ impl<'agent> ChunkUploader<'agent> {
                 progress,
             )
             .await?;
-            let mut map = self.id_mapping.lock().await;
-            map.insert(uploader_chunk_id, canister_chunk_id);
+
+            self.id_mapping
+                .lock()
+                .await
+                .insert(uploader_chunk_id, canister_chunk_id.clone());
+            self.canister_chunk_sizes
+                .lock()
+                .await
+                .insert(canister_chunk_id, chunk_size);
 
             Ok(uploader_chunk_id)
         } else {
@@ -146,6 +159,15 @@ impl<'agent> ChunkUploader<'agent> {
         self.chunks.load(Ordering::SeqCst)
     }
 
+    /// Get total size of chunks by their canister chunk IDs
+    pub(crate) async fn get_canister_chunk_total_size(&self, canister_chunk_ids: &[Nat]) -> usize {
+        let sizes = self.canister_chunk_sizes.lock().await;
+        canister_chunk_ids
+            .iter()
+            .filter_map(|id| sizes.get(id))
+            .sum()
+    }
+
     /// Call only after `finalize_upload` has completed.
     pub(crate) async fn uploader_ids_to_canister_chunk_ids(
         &self,
@@ -212,15 +234,27 @@ impl<'agent> ChunkUploader<'agent> {
         }
 
         try_join_all(batches.into_iter().map(|chunks| async move {
-            let (uploader_chunk_ids, chunks): (Vec<_>, Vec<_>) = chunks.into_iter().unzip();
-            let canister_chunk_ids =
-                create_chunks(&self.canister, &self.batch_id, chunks, semaphores, progress).await?;
-            let mut map = self.id_mapping.lock().await;
-            for (uploader_id, canister_id) in uploader_chunk_ids
+            let (uploader_chunk_ids, chunk_contents): (Vec<_>, Vec<_>) = chunks.into_iter().unzip();
+            let chunk_sizes: Vec<usize> = chunk_contents.iter().map(|c| c.len()).collect();
+
+            let canister_chunk_ids = create_chunks(
+                &self.canister,
+                &self.batch_id,
+                chunk_contents,
+                semaphores,
+                progress,
+            )
+            .await?;
+
+            let mut id_map = self.id_mapping.lock().await;
+            let mut canister_sizes = self.canister_chunk_sizes.lock().await;
+            for ((uploader_id, canister_id), chunk_size) in uploader_chunk_ids
                 .into_iter()
                 .zip(canister_chunk_ids.into_iter())
+                .zip(chunk_sizes.into_iter())
             {
-                map.insert(uploader_id, canister_id);
+                id_map.insert(uploader_id, canister_id.clone());
+                canister_sizes.insert(canister_id, chunk_size);
             }
             Ok(())
         }))

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

@@ -45,15 +45,15 @@ use walkdir::WalkDir;
 const KNOWN_DIRECTORIES: [&str; 1] = [".well-known"];
 
 /// Sets the contents of the asset canister to the contents of a directory, including deleting old assets.
-pub async fn upload_content_and_assemble_sync_operations(
-    canister: &Canister<'_>,
+pub async fn upload_content_and_assemble_sync_operations<'a>(
+    canister: &Canister<'a>,
     canister_api_version: u16,
     dirs: &[&Path],
     no_delete: bool,
     mode: batch_upload::plumbing::Mode,
     logger: &Logger,
     progress: Option<&dyn AssetSyncProgressRenderer>,
-) -> Result<CommitBatchArguments, UploadContentError> {
+) -> Result<(CommitBatchArguments, ChunkUploader<'a>), UploadContentError> {
     if let Some(progress) = progress {
         progress.set_state(AssetSyncState::GatherAssetDescriptors);
     }
@@ -144,7 +144,7 @@ pub async fn upload_content_and_assemble_sync_operations(
     // -vv
     trace!(logger, "Value of CommitBatch: {:?}", commit_batch_args);
 
-    Ok(commit_batch_args)
+    Ok((commit_batch_args, chunk_uploader))
 }
 
 /// Sets the contents of the asset canister to the contents of a directory, including deleting old assets.
@@ -156,7 +156,7 @@ pub async fn sync(
     progress: Option<&dyn AssetSyncProgressRenderer>,
 ) -> Result<(), SyncError> {
     let canister_api_version = api_version(canister).await;
-    let commit_batch_args = upload_content_and_assemble_sync_operations(
+    let (commit_batch_args, chunk_uploader) = upload_content_and_assemble_sync_operations(
         canister,
         canister_api_version,
         dirs,
@@ -180,64 +180,134 @@ pub async fn sync(
             warn!(logger, "The asset canister is running an old version of the API. It will not be able to set assets properties.");
             commit_batch(canister, commit_batch_args_v0).await
         }
-        BATCH_UPLOAD_API_VERSION.. => commit_in_stages(canister, commit_batch_args, logger, progress).await,
+        BATCH_UPLOAD_API_VERSION.. => commit_in_stages(canister, commit_batch_args, &chunk_uploader, logger, progress).await,
     }.map_err(CommitBatchFailed)?;
     if let Some(progress) = progress {
         progress.set_state(AssetSyncState::Done);
     }
     Ok(())
 }
 
-async fn commit_in_stages(
+/// Creates batches of operations for committing.
+///
+/// Batches are created based on three conditions (any of which triggers a new batch):
+/// 1. 500 operations reached -  generally respected limit to avoid too much cert tree work
+/// 2. 1.5MB of header map data reached - headers are the largest part of ingress message size
+/// 3. 100MB of total chunk size reached - the full asset content gets hashed in the commit message
+async fn create_commit_batches<'a>(
+    operations: Vec<BatchOperationKind>,
+    chunk_uploader: &ChunkUploader<'a>,
+) -> Vec<Vec<BatchOperationKind>> {
+    const MAX_OPERATIONS_PER_BATCH: usize = 500; // empirically this works good enough
+    const MAX_HEADER_MAP_SIZE: usize = 1_500_000; // 1.5 MB leaves plenty of room for other data that we do not calculate precisely
+    const MAX_ASSET_CONTENT_SIZE: usize = 100_000_000; // 100 MB is ~20% of how much data we can hash in a single message: 40b instructions per update call, measured best case of 80 instructions per byte hashed -> ~500MB limit
+
+    let mut batches = Vec::new();
+    let mut current_batch = Vec::new();
+    let mut operation_count = 0;
+    let mut header_map_size = 0;
+    let mut content_size = 0;
+
+    for operation in operations {
+        let operation_header_size = calculate_header_size(&operation);
+        let operation_chunk_size = calculate_content_size(&operation, chunk_uploader).await;
+
+        // Check if adding this operation would exceed any limits
+        let would_exceed_operation_limit = operation_count >= MAX_OPERATIONS_PER_BATCH;
+        let would_exceed_header_limit =
+            header_map_size + operation_header_size >= MAX_HEADER_MAP_SIZE;
+        let would_exceed_chunk_limit =
+            content_size + operation_chunk_size >= MAX_ASSET_CONTENT_SIZE;
+
+        if (would_exceed_operation_limit || would_exceed_header_limit || would_exceed_chunk_limit)
+            && !current_batch.is_empty()
+        {
+            // Start a new batch
+            batches.push(current_batch);
+            current_batch = Vec::new();
+            operation_count = 0;
+            header_map_size = 0;
+            content_size = 0;
+        }
+
+        // Add operation to current batch
+        current_batch.push(operation);
+        operation_count += 1;
+        header_map_size += operation_header_size;
+        content_size += operation_chunk_size;
+    }
+
+    // Add the last batch if it has any operations
+    if !current_batch.is_empty() {
+        batches.push(current_batch);
+    }
+
+    batches
+}
+
+/// Calculate the size in bytes of header data for an operation.
+fn calculate_header_size(operation: &BatchOperationKind) -> usize {
+    match operation {
+        BatchOperationKind::CreateAsset(args) => args.headers.as_ref().map_or(0, |headers| {
+            headers.iter().map(|(k, v)| k.len() + v.len()).sum()
+        }),
+        BatchOperationKind::SetAssetProperties(args) => args
+            .headers
+            .as_ref()
+            .and_then(|h| h.as_ref())
+            .map_or(0, |headers| {
+                headers.iter().map(|(k, v)| k.len() + v.len()).sum()
+            }),
+        _ => 0,
+    }
+}
+
+/// Calculate the size in bytes of chunk data for an operation.
+/// This includes both:
+/// - Chunks referenced by `chunk_ids` (looked up from ChunkUploader)
+/// - The `last_chunk` field which is included directly in the commit message
+async fn calculate_content_size<'a>(
+    operation: &BatchOperationKind,
+    chunk_uploader: &ChunkUploader<'a>,
+) -> usize {
+    match operation {
+        BatchOperationKind::SetAssetContent(args) => {
+            let chunk_ids_size = chunk_uploader
+                .get_canister_chunk_total_size(&args.chunk_ids)
+                .await;
+            let last_chunk_size = args.last_chunk.as_ref().map_or(0, |chunk| chunk.len());
+            chunk_ids_size + last_chunk_size
+        }
+        _ => 0,
+    }
+}
+
+async fn commit_in_stages<'a>(
     canister: &Canister<'_>,
     commit_batch_args: CommitBatchArguments,
+    chunk_uploader: &ChunkUploader<'a>,
     logger: &Logger,
     progress: Option<&dyn AssetSyncProgressRenderer>,
 ) -> Result<(), AgentError> {
     if let Some(progress) = progress {
         progress.set_total_batch_operations(commit_batch_args.operations.len());
     }
-    // Note that SetAssetProperties operations are only generated for assets that
-    // already exist, since CreateAsset operations set all properties.
-    let (set_properties_operations, other_operations): (Vec<_>, Vec<_>) = commit_batch_args
-        .operations
-        .into_iter()
-        .partition(|op| matches!(op, BatchOperationKind::SetAssetProperties(_)));
-
-    // This part seems reasonable in general as a separate batch
-    for operations in set_properties_operations.chunks(500) {
-        debug!(logger, "Setting properties of {} assets.", operations.len());
-        commit_batch(
-            canister,
-            CommitBatchArguments {
-                batch_id: Nat::from(0_u8),
-                operations: operations.into(),
-            },
-        )
-        .await?;
-        if let Some(progress) = progress {
-            progress.add_committed_batch_operations(operations.len());
-        }
-    }
 
-    // Seen to work at 800 ({"SetAssetContent": 932, "Delete": 47, "CreateAsset": 58})
-    // so 500 shouldn't exceed per-message instruction limit
-    for operations in other_operations.chunks(500) {
-        debug!(
-            logger,
-            "Committing batch with {} operations.",
-            operations.len()
-        );
+    let batches = create_commit_batches(commit_batch_args.operations, chunk_uploader).await;
+
+    for operations in batches {
+        let op_amount = operations.len();
+        debug!(logger, "Committing batch with {op_amount} operations.");
         commit_batch(
             canister,
             CommitBatchArguments {
                 batch_id: Nat::from(0_u8),
-                operations: operations.into(),
+                operations,
             },
         )
         .await?;
         if let Some(progress) = progress {
-            progress.add_committed_batch_operations(operations.len());
+            progress.add_committed_batch_operations(op_amount);
         }
     }
 
@@ -260,7 +330,7 @@ pub async fn prepare_sync_for_proposal(
     progress: Option<&dyn AssetSyncProgressRenderer>,
 ) -> Result<(Nat, ByteBuf), PrepareSyncForProposalError> {
     let canister_api_version = api_version(canister).await;
-    let arg = upload_content_and_assemble_sync_operations(
+    let (arg, _chunk_uploader) = upload_content_and_assemble_sync_operations(
         canister,
         canister_api_version,
         dirs,