|
| 1 | +use byte_unit::{Byte, UnitType}; |
| 2 | +use clap::Args; |
| 3 | +use icp::context::Context; |
| 4 | +use icp::prelude::*; |
| 5 | + |
| 6 | +use super::SnapshotId; |
| 7 | +use crate::commands::args; |
| 8 | +use crate::operations::misc::format_timestamp; |
| 9 | +use crate::operations::snapshot_transfer::{ |
| 10 | + BlobType, SnapshotPaths, SnapshotTransferError, create_transfer_progress_bar, |
| 11 | + delete_download_progress, download_blob_to_file, download_wasm_chunk, load_download_progress, |
| 12 | + load_metadata, read_snapshot_metadata, save_metadata, |
| 13 | +}; |
| 14 | + |
| 15 | +#[derive(Debug, Args)] |
| 16 | +pub(crate) struct DownloadArgs { |
| 17 | + #[command(flatten)] |
| 18 | + pub(crate) cmd_args: args::CanisterCommandArgs, |
| 19 | + |
| 20 | + /// The snapshot ID to download (hex-encoded) |
| 21 | + snapshot_id: SnapshotId, |
| 22 | + |
| 23 | + /// Output directory for the snapshot files |
| 24 | + #[arg(long, short = 'o')] |
| 25 | + output: PathBuf, |
| 26 | + |
| 27 | + /// Resume a previously interrupted download |
| 28 | + #[arg(long)] |
| 29 | + resume: bool, |
| 30 | +} |
| 31 | + |
| 32 | +pub(crate) async fn exec(ctx: &Context, args: &DownloadArgs) -> Result<(), anyhow::Error> { |
| 33 | + let selections = args.cmd_args.selections(); |
| 34 | + |
| 35 | + let agent = ctx |
| 36 | + .get_agent( |
| 37 | + &selections.identity, |
| 38 | + &selections.network, |
| 39 | + &selections.environment, |
| 40 | + ) |
| 41 | + .await?; |
| 42 | + let cid = ctx |
| 43 | + .get_canister_id( |
| 44 | + &selections.canister, |
| 45 | + &selections.network, |
| 46 | + &selections.environment, |
| 47 | + ) |
| 48 | + .await?; |
| 49 | + |
| 50 | + let name = &args.cmd_args.canister; |
| 51 | + let snapshot_id = &args.snapshot_id.0; |
| 52 | + |
| 53 | + // Open or create the snapshot directory with a lock |
| 54 | + let snapshot_dir = SnapshotPaths::new(args.output.clone())?; |
| 55 | + |
| 56 | + snapshot_dir |
| 57 | + .with_write(async |paths| { |
| 58 | + // Ensure directories exist |
| 59 | + paths.ensure_dirs()?; |
| 60 | + |
| 61 | + // Check if we should resume or start fresh |
| 62 | + let metadata = if args.resume && paths.metadata_path().exists() { |
| 63 | + ctx.term.write_line("Resuming previous download...")?; |
| 64 | + load_metadata(paths)? |
| 65 | + } else if !args.resume { |
| 66 | + // Check if directory has existing files (besides lock) |
| 67 | + let has_files = paths.metadata_path().exists() |
| 68 | + || paths.wasm_module_path().exists() |
| 69 | + || paths.wasm_memory_path().exists() |
| 70 | + || paths.stable_memory_path().exists(); |
| 71 | + |
| 72 | + if has_files { |
| 73 | + return Err(SnapshotTransferError::DirectoryNotEmpty { |
| 74 | + path: args.output.clone(), |
| 75 | + } |
| 76 | + .into()); |
| 77 | + } |
| 78 | + |
| 79 | + // Fetch metadata from canister |
| 80 | + ctx.term.write_line(&format!( |
| 81 | + "Downloading snapshot {id} from canister {name} ({cid})", |
| 82 | + id = hex::encode(snapshot_id), |
| 83 | + ))?; |
| 84 | + |
| 85 | + let metadata = read_snapshot_metadata(&agent, cid, snapshot_id).await?; |
| 86 | + |
| 87 | + ctx.term.write_line(&format!( |
| 88 | + " Timestamp: {}", |
| 89 | + format_timestamp(metadata.taken_at_timestamp) |
| 90 | + ))?; |
| 91 | + |
| 92 | + let total_size = metadata.wasm_module_size |
| 93 | + + metadata.wasm_memory_size |
| 94 | + + metadata.stable_memory_size; |
| 95 | + ctx.term.write_line(&format!( |
| 96 | + " Total size: {}", |
| 97 | + Byte::from_u64(total_size).get_appropriate_unit(UnitType::Binary) |
| 98 | + ))?; |
| 99 | + |
| 100 | + // Save metadata |
| 101 | + save_metadata(&metadata, paths)?; |
| 102 | + |
| 103 | + metadata |
| 104 | + } else { |
| 105 | + return Err(SnapshotTransferError::NoExistingDownload { |
| 106 | + path: args.output.clone(), |
| 107 | + } |
| 108 | + .into()); |
| 109 | + }; |
| 110 | + |
| 111 | + // Load download progress (handles gaps from previous interrupted downloads) |
| 112 | + let mut progress = load_download_progress(paths)?; |
| 113 | + |
| 114 | + // Download WASM module |
| 115 | + if metadata.wasm_module_size > 0 { |
| 116 | + if !progress.wasm_module.is_complete(metadata.wasm_module_size) { |
| 117 | + let pb = create_transfer_progress_bar(metadata.wasm_module_size, "WASM module"); |
| 118 | + download_blob_to_file( |
| 119 | + &agent, |
| 120 | + cid, |
| 121 | + snapshot_id, |
| 122 | + BlobType::WasmModule, |
| 123 | + metadata.wasm_module_size, |
| 124 | + paths, |
| 125 | + &mut progress, |
| 126 | + &pb, |
| 127 | + ) |
| 128 | + .await?; |
| 129 | + pb.finish_with_message("done"); |
| 130 | + } else { |
| 131 | + ctx.term.write_line("WASM module: already complete")?; |
| 132 | + } |
| 133 | + } |
| 134 | + |
| 135 | + // Download WASM memory |
| 136 | + if metadata.wasm_memory_size > 0 { |
| 137 | + if !progress.wasm_memory.is_complete(metadata.wasm_memory_size) { |
| 138 | + let pb = create_transfer_progress_bar(metadata.wasm_memory_size, "WASM memory"); |
| 139 | + download_blob_to_file( |
| 140 | + &agent, |
| 141 | + cid, |
| 142 | + snapshot_id, |
| 143 | + BlobType::WasmMemory, |
| 144 | + metadata.wasm_memory_size, |
| 145 | + paths, |
| 146 | + &mut progress, |
| 147 | + &pb, |
| 148 | + ) |
| 149 | + .await?; |
| 150 | + pb.finish_with_message("done"); |
| 151 | + } else { |
| 152 | + ctx.term.write_line("WASM memory: already complete")?; |
| 153 | + } |
| 154 | + } |
| 155 | + |
| 156 | + // Download stable memory |
| 157 | + if metadata.stable_memory_size > 0 { |
| 158 | + if !progress |
| 159 | + .stable_memory |
| 160 | + .is_complete(metadata.stable_memory_size) |
| 161 | + { |
| 162 | + let pb = |
| 163 | + create_transfer_progress_bar(metadata.stable_memory_size, "Stable memory"); |
| 164 | + download_blob_to_file( |
| 165 | + &agent, |
| 166 | + cid, |
| 167 | + snapshot_id, |
| 168 | + BlobType::StableMemory, |
| 169 | + metadata.stable_memory_size, |
| 170 | + paths, |
| 171 | + &mut progress, |
| 172 | + &pb, |
| 173 | + ) |
| 174 | + .await?; |
| 175 | + pb.finish_with_message("done"); |
| 176 | + } else { |
| 177 | + ctx.term.write_line("Stable memory: already complete")?; |
| 178 | + } |
| 179 | + } else { |
| 180 | + // Create empty stable memory file |
| 181 | + icp::fs::write(&paths.stable_memory_path(), &[])?; |
| 182 | + } |
| 183 | + |
| 184 | + // Download WASM chunk store |
| 185 | + if !metadata.wasm_chunk_store.is_empty() { |
| 186 | + ctx.term.write_line(&format!( |
| 187 | + "Downloading {} WASM chunks...", |
| 188 | + metadata.wasm_chunk_store.len() |
| 189 | + ))?; |
| 190 | + |
| 191 | + for chunk_hash in &metadata.wasm_chunk_store { |
| 192 | + let chunk_path = paths.wasm_chunk_path(&chunk_hash.hash); |
| 193 | + if !chunk_path.exists() { |
| 194 | + download_wasm_chunk(&agent, cid, snapshot_id, chunk_hash, paths).await?; |
| 195 | + } |
| 196 | + } |
| 197 | + ctx.term.write_line("WASM chunks: done")?; |
| 198 | + } |
| 199 | + |
| 200 | + // Clean up progress file on success |
| 201 | + delete_download_progress(paths)?; |
| 202 | + |
| 203 | + ctx.term |
| 204 | + .write_line(&format!("Snapshot downloaded to {}", args.output))?; |
| 205 | + |
| 206 | + Ok::<_, anyhow::Error>(()) |
| 207 | + }) |
| 208 | + .await??; |
| 209 | + |
| 210 | + Ok(()) |
| 211 | +} |
0 commit comments