Skip to content

Commit

Permalink
feat: PocketIC library functions to manage WASM chunk store and insta…
Browse files Browse the repository at this point in the history 
…ll chunked code (#2546)

This PR contains the following changes:
- adds the functions `PocketIc::upload_chunk`,
`PocketIc::stored_chunks`, and `PocketIc::clear_chunk_store` to manage
the WASM chunk store of a canister;
- adds the function `PocketIc::install_chunked_canister` to install a
canister from WASM chunks in the WASM chunk store of a canister.
  • Loading branch information
@mraszyk
mraszyk authored Nov 11, 2024
1 parent 1c9fe92 commit fd4bbd4
Show file tree
Hide file tree
Showing 4 changed files with 249 additions and 53 deletions.
2 changes: 2 additions & 0 deletions packages/pocket-ic/CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- The function `get_default_effective_canister_id` to retrieve a default effective canister id for canister creation on a PocketIC instance.
- The function `PocketIc::get_controllers` to get the controllers of a canister.
- Functions `PocketIc::take_canister_snapshot`, `PocketIc::load_canister_snapshot`, `PocketIc::list_canister_snapshots`, and `PocketIc::delete_canister_snapshot` to manage canister snapshots.
- Functions `PocketIc::upload_chunk`, `PocketIc::stored_chunks`, and `PocketIc::clear_chunk_store` to manage the WASM chunk store of a canister.
- The function `PocketIc::install_chunked_canister` to install a canister from WASM chunks in the WASM chunk store of a canister.

### Removed
- Functions `PocketIc::from_config`, `PocketIc::from_config_and_max_request_time`, and `PocketIc::from_config_and_server_url`.
Expand Down
69 changes: 68 additions & 1 deletion packages/pocket-ic/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ use crate::{
InstanceId, MockCanisterHttpResponse, RawEffectivePrincipal, RawMessageId, SubnetId,
SubnetKind, SubnetSpec, Topology,
},
management_canister::{CanisterId, CanisterStatusResult, Snapshot},
management_canister::{CanisterId, CanisterInstallMode, CanisterStatusResult, Snapshot},
nonblocking::PocketIc as PocketIcAsync,
};
use candid::{
Expand Down Expand Up @@ -739,6 +739,73 @@ impl PocketIc {
})
}

/// Upload a WASM chunk to the WASM chunk store of a canister.
/// Returns the WASM chunk hash.
#[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub fn upload_chunk(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
chunk: Vec<u8>,
) -> Result<Vec<u8>, CallError> {
let runtime = self.runtime.clone();
runtime.block_on(async {
self.pocket_ic
.upload_chunk(canister_id, sender, chunk)
.await
})
}

/// List WASM chunk hashes in the WASM chunk store of a canister.
#[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub fn stored_chunks(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<Vec<Vec<u8>>, CallError> {
let runtime = self.runtime.clone();
runtime.block_on(async { self.pocket_ic.stored_chunks(canister_id, sender).await })
}

/// Clear the WASM chunk store of a canister.
#[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub fn clear_chunk_store(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<(), CallError> {
let runtime = self.runtime.clone();
runtime.block_on(async { self.pocket_ic.clear_chunk_store(canister_id, sender).await })
}

/// Install a WASM module assembled from chunks on an existing canister.
#[instrument(skip(self, mode, chunk_hashes_list, wasm_module_hash, arg), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string(), store_canister_id = %store_canister_id.to_string(), arg_len = %arg.len()))]
pub fn install_chunked_canister(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
mode: CanisterInstallMode,
store_canister_id: CanisterId,
chunk_hashes_list: Vec<Vec<u8>>,
wasm_module_hash: Vec<u8>,
arg: Vec<u8>,
) -> Result<(), CallError> {
let runtime = self.runtime.clone();
runtime.block_on(async {
self.pocket_ic
.install_chunked_canister(
canister_id,
sender,
mode,
store_canister_id,
chunk_hashes_list,
wasm_module_hash,
arg,
)
.await
})
}

/// Install a WASM module on an existing canister.
#[instrument(skip(self, wasm_module, arg), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub fn install_canister(
Expand Down
158 changes: 109 additions & 49 deletions packages/pocket-ic/src/nonblocking.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ use crate::management_canister::{
CanisterInstallModeUpgradeInnerWasmMemoryPersistenceInner, CanisterSettings,
CanisterStatusResult, ChunkHash, DeleteCanisterSnapshotArgs, InstallChunkedCodeArgs,
InstallCodeArgs, LoadCanisterSnapshotArgs, ProvisionalCreateCanisterWithCyclesArgs, Snapshot,
TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs,
StoredChunksResult, TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs,
UploadChunkResult,
};
pub use crate::DefaultEffectiveCanisterIdError;
use crate::{CallError, PocketIcBuilder, UserError, WasmResult};
Expand Down Expand Up @@ -730,6 +731,98 @@ impl PocketIc {
canister_id
}

/// Upload a WASM chunk to the WASM chunk store of a canister.
/// Returns the WASM chunk hash.
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn upload_chunk(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
chunk: Vec<u8>,
) -> Result<Vec<u8>, CallError> {
call_candid_as::<_, (UploadChunkResult,)>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"upload_chunk",
(UploadChunkArgs { canister_id, chunk },),
)
.await
.map(|responses| responses.0.hash)
}

/// List WASM chunk hashes in the WASM chunk store of a canister.
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn stored_chunks(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<Vec<Vec<u8>>, CallError> {
call_candid_as::<_, (StoredChunksResult,)>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"stored_chunks",
(CanisterIdRecord { canister_id },),
)
.await
.map(|responses| responses.0.into_iter().map(|chunk| chunk.hash).collect())
}

/// Clear the WASM chunk store of a canister.
#[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))]
pub async fn clear_chunk_store(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
) -> Result<(), CallError> {
call_candid_as(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"clear_chunk_store",
(CanisterIdRecord { canister_id },),
)
.await
}

/// Install a WASM module assembled from chunks on an existing canister.
#[instrument(skip(self, mode, chunk_hashes_list, wasm_module_hash, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string(), store_canister_id = %store_canister_id.to_string(), arg_len = %arg.len()))]
pub async fn install_chunked_canister(
&self,
canister_id: CanisterId,
sender: Option<Principal>,
mode: CanisterInstallMode,
store_canister_id: CanisterId,
chunk_hashes_list: Vec<Vec<u8>>,
wasm_module_hash: Vec<u8>,
arg: Vec<u8>,
) -> Result<(), CallError> {
call_candid_as(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"install_chunked_code",
(InstallChunkedCodeArgs {
mode,
target_canister: canister_id,
store_canister: Some(store_canister_id),
chunk_hashes_list: chunk_hashes_list
.into_iter()
.map(|hash| ChunkHash { hash })
.collect(),
wasm_module_hash,
arg,
sender_canister_version: None,
},),
)
.await
}

async fn install_canister_helper(
&self,
mode: CanisterInstallMode,
Expand All @@ -755,60 +848,27 @@ impl PocketIc {
)
.await
} else {
self.clear_chunk_store(canister_id, sender).await.unwrap();
let chunks: Vec<_> = wasm_module.chunks(INSTALL_CODE_CHUNK_SIZE).collect();
let hashes: Vec<_> = chunks
.iter()
.map(|c| {
let mut hasher = Sha256::new();
hasher.update(c);
ChunkHash {
hash: hasher.finalize().to_vec(),
}
})
.collect();
call_candid_as::<_, ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"clear_chunk_store",
(CanisterIdRecord { canister_id },),
)
.await
.unwrap();
let mut hashes = vec![];
for chunk in chunks {
call_candid_as::<_, ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"upload_chunk",
(UploadChunkArgs {
canister_id,
chunk: chunk.to_vec(),
},),
)
.await
.unwrap();
let hash = self
.upload_chunk(canister_id, sender, chunk.to_vec())
.await
.unwrap();
hashes.push(hash);
}
let mut hasher = Sha256::new();
hasher.update(wasm_module);
let wasm_module_hash = hasher.finalize().to_vec();
call_candid_as::<_, ()>(
self,
Principal::management_canister(),
RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
sender.unwrap_or(Principal::anonymous()),
"install_chunked_code",
(InstallChunkedCodeArgs {
mode,
target_canister: canister_id,
store_canister: None,
chunk_hashes_list: hashes,
wasm_module_hash,
arg,
sender_canister_version: None,
},),
self.install_chunked_canister(
canister_id,
sender,
mode,
canister_id,
hashes,
wasm_module_hash,
arg,
)
.await
}
Expand Down
73 changes: 70 additions & 3 deletions packages/pocket-ic/tests/tests.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use candid::{decode_one, encode_one, CandidType, Decode, Deserialize, Encode, Principal};
use pocket_ic::management_canister::{
CanisterId, CanisterIdRecord, CanisterSettings, EcdsaPublicKeyResult, HttpRequestResult,
ProvisionalCreateCanisterWithCyclesArgs, SchnorrAlgorithm, SchnorrPublicKeyArgsKeyId,
SchnorrPublicKeyResult,
CanisterId, CanisterIdRecord, CanisterInstallMode, CanisterSettings, EcdsaPublicKeyResult,
HttpRequestResult, ProvisionalCreateCanisterWithCyclesArgs, SchnorrAlgorithm,
SchnorrPublicKeyArgsKeyId, SchnorrPublicKeyResult,
};
use pocket_ic::{
common::rest::{
Expand Down Expand Up @@ -1699,3 +1699,70 @@ fn test_canister_snapshots() {
.unwrap();
pic.take_canister_snapshot(canister_id, None, None).unwrap();
}

#[test]
fn test_wasm_chunk_store() {
let pic = PocketIc::new();

// We create an empty canister and top it up with cycles (WASM chunk store operations cost cycles).
let canister_id = pic.create_canister();
pic.add_cycles(canister_id, INIT_CYCLES);

// There should be no chunks in the WASM chunk store yet.
let stored_chunks = pic.stored_chunks(canister_id, None).unwrap();
assert!(stored_chunks.is_empty());

// Chunk the test canister into two chunks.
let mut first_chunk = test_canister_wasm();
let second_chunk = first_chunk.split_off(first_chunk.len() / 2);
assert!(!first_chunk.is_empty());
assert!(!second_chunk.is_empty());

// We upload a bogus chunk to the WASM chunk store and confirm that the returned hash
// matches the actual hash of the chunk.
let first_chunk_hash = pic
.upload_chunk(canister_id, None, first_chunk.clone())
.unwrap();
let mut hasher = Sha256::new();
hasher.update(first_chunk.clone());
assert_eq!(first_chunk_hash, hasher.finalize().to_vec());

// We upload the same chunk once more and get the same hash back.
let same_chunk_hash = pic
.upload_chunk(canister_id, None, first_chunk.clone())
.unwrap();
assert_eq!(first_chunk_hash, same_chunk_hash);

// We upload a different chunk.
let second_chunk_hash = pic.upload_chunk(canister_id, None, second_chunk).unwrap();

// Now the two chunks should be stored in the WASM chunk store.
let stored_chunks = pic.stored_chunks(canister_id, None).unwrap();
assert_eq!(stored_chunks.len(), 2);
assert!(stored_chunks.contains(&first_chunk_hash));
assert!(stored_chunks.contains(&second_chunk_hash));

// We create a new canister and install it from chunks.
let test_canister = pic.create_canister();
pic.add_cycles(test_canister, INIT_CYCLES);
let mut hasher = Sha256::new();
hasher.update(test_canister_wasm());
let test_canister_wasm_hash = hasher.finalize().to_vec();
pic.install_chunked_canister(
test_canister,
None,
CanisterInstallMode::Install,
canister_id,
vec![first_chunk_hash, second_chunk_hash],
test_canister_wasm_hash,
Encode!(&()).unwrap(),
)
.unwrap();

// We clear the WASM chunk store.
pic.clear_chunk_store(canister_id, None).unwrap();

// There should be no more chunks in the WASM chunk store.
let stored_chunks = pic.stored_chunks(canister_id, None).unwrap();
assert!(stored_chunks.is_empty());
}

0 comments on commit fd4bbd4

Please sign in to comment.