microsandbox_image/checkpoint/
compact.rs1use std::io;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6
7use imago::file::File as ImagoFile;
8use imago::qcow2::Qcow2;
9use imago::raw::Raw;
10use imago::{
11 DenyImplicitOpenGate, DynStorage, FormatAccess, FormatCreateBuilder, FormatDriverBuilder,
12 Mapping,
13};
14use tokio::io::{AsyncSeekExt, AsyncWriteExt};
15
16type SharedImage = Arc<FormatAccess<Box<dyn DynStorage>>>;
21
22#[derive(Clone, Debug)]
24pub struct CompactLayer {
25 pub path: PathBuf,
27 pub qcow2: bool,
29}
30
31#[derive(Clone, Debug)]
33pub struct CompactMaterialization {
34 pub virtual_size: u64,
36 pub materialized_bytes: u64,
38}
39
40async fn open_chain(layers: &[CompactLayer]) -> io::Result<SharedImage> {
46 open_chain_access(layers, false).await
47}
48
49pub(crate) async fn open_writable_chain(layers: &[CompactLayer]) -> io::Result<SharedImage> {
51 open_chain_access(layers, true).await
52}
53
54async fn open_chain_access(
55 layers: &[CompactLayer],
56 writable_head: bool,
57) -> io::Result<SharedImage> {
58 if layers.is_empty() || layers.iter().skip(1).any(|layer| !layer.qcow2) {
59 return Err(io::Error::new(
60 io::ErrorKind::InvalidInput,
61 "invalid raw/qcow2 compaction prefix",
62 ));
63 }
64 let mut backing: Option<SharedImage> = None;
65 for (index, layer) in layers.iter().enumerate() {
66 let writable = writable_head && index + 1 == layers.len();
67 let storage: Box<dyn DynStorage> = Box::new(ImagoFile::try_from(
68 std::fs::OpenOptions::new()
69 .read(true)
70 .write(writable)
71 .open(&layer.path)?,
72 )?);
73 backing = Some(if layer.qcow2 {
74 let image = Qcow2::<Box<dyn DynStorage>, SharedImage>::builder(storage)
75 .write(writable)
76 .backing(backing)
77 .data_file(None)
78 .open(DenyImplicitOpenGate::default())
79 .await?;
80 if image.requires_external_data_file() {
81 return Err(io::Error::new(
82 io::ErrorKind::Unsupported,
83 "external qcow2 data files are not supported",
84 ));
85 }
86 Arc::new(FormatAccess::new(image))
87 } else {
88 Arc::new(FormatAccess::new(
89 Raw::<Box<dyn DynStorage>>::builder(storage)
90 .write(writable)
91 .open(DenyImplicitOpenGate::default())
92 .await?,
93 ))
94 });
95 }
96 Ok(backing.expect("nonempty chain checked above"))
97}
98
99pub async fn materialize_compact_prefix(
105 layers: &[CompactLayer],
106 destination: &Path,
107) -> io::Result<CompactMaterialization> {
108 let source = open_chain(layers).await?;
109 let virtual_size = source.size();
110 if virtual_size == 0 || !virtual_size.is_multiple_of(512) {
111 return Err(io::Error::new(
112 io::ErrorKind::InvalidData,
113 "invalid compaction disk capacity",
114 ));
115 }
116 let file = std::fs::OpenOptions::new()
117 .read(true)
118 .write(true)
119 .create_new(true)
120 .open(destination)?;
121 Qcow2::<ImagoFile>::create_builder(ImagoFile::try_from(file)?)
122 .size(virtual_size)
123 .cluster_size(65536)
124 .create()
125 .await?;
126 let file = std::fs::OpenOptions::new()
127 .read(true)
128 .write(true)
129 .open(destination)?;
130 let target = FormatAccess::new(
131 Qcow2::<ImagoFile>::builder(ImagoFile::try_from(file)?)
132 .write(true)
133 .backing(None)
134 .data_file(None)
135 .open(DenyImplicitOpenGate::default())
136 .await?,
137 );
138 let mut buffer = vec![0u8; 1024 * 1024];
139 let mut offset = 0;
140 let mut materialized_bytes = 0;
141 while offset < virtual_size {
142 let (mapping, length) = source.get_mapping(offset, virtual_size - offset).await?;
143 if length == 0 {
144 return Err(io::Error::new(
145 io::ErrorKind::InvalidData,
146 "compaction mapping made no progress",
147 ));
148 }
149 if matches!(mapping, Mapping::Zero { .. }) {
150 offset += length;
151 continue;
152 }
153 let count = length.min(buffer.len() as u64) as usize;
154 source.read(&mut buffer[..count], offset).await?;
155 if buffer[..count].iter().any(|byte| *byte != 0) {
158 target.write(&buffer[..count], offset).await?;
159 materialized_bytes += count as u64;
160 }
161 offset += count as u64;
162 }
163 target.flush().await?;
164 target.sync().await?;
165 Ok(CompactMaterialization {
166 virtual_size,
167 materialized_bytes,
168 })
169}
170
171pub async fn materialize_raw_prefix(
177 layers: &[CompactLayer],
178 destination: &Path,
179) -> io::Result<CompactMaterialization> {
180 let source = open_chain(layers).await?;
181 let virtual_size = source.size();
182 if virtual_size == 0 || !virtual_size.is_multiple_of(512) {
183 return Err(io::Error::new(
184 io::ErrorKind::InvalidData,
185 "invalid raw disk capacity",
186 ));
187 }
188 let mut target = tokio::fs::OpenOptions::new()
189 .write(true)
190 .create_new(true)
191 .open(destination)
192 .await?;
193 target.set_len(virtual_size).await?;
194 let mut buffer = vec![0u8; 1024 * 1024];
195 let mut offset = 0;
196 let mut materialized_bytes = 0;
197 while offset < virtual_size {
198 let (mapping, length) = source.get_mapping(offset, virtual_size - offset).await?;
199 if length == 0 {
200 return Err(io::Error::new(
201 io::ErrorKind::InvalidData,
202 "raw disk mapping made no progress",
203 ));
204 }
205 if matches!(mapping, Mapping::Zero { .. }) {
206 offset += length;
207 continue;
208 }
209 let count = length.min(buffer.len() as u64) as usize;
210 source.read(&mut buffer[..count], offset).await?;
211 if buffer[..count].iter().any(|byte| *byte != 0) {
212 target.seek(io::SeekFrom::Start(offset)).await?;
213 target.write_all(&buffer[..count]).await?;
214 materialized_bytes += count as u64;
215 }
216 offset += count as u64;
217 }
218 target.sync_all().await?;
219 Ok(CompactMaterialization {
220 virtual_size,
221 materialized_bytes,
222 })
223}
224
225pub async fn compact_layer_capacity(layer: CompactLayer) -> io::Result<u64> {
227 Ok(open_chain(&[layer]).await?.size())
228}
229
230pub fn layer_capacities(layers: Vec<CompactLayer>) -> io::Result<Vec<u64>> {
233 std::thread::spawn(move || {
234 tokio::runtime::Builder::new_current_thread()
235 .enable_all()
236 .build()?
237 .block_on(async {
238 let mut capacities = Vec::with_capacity(layers.len());
239 for layer in layers {
240 capacities.push(compact_layer_capacity(layer).await?);
241 }
242 Ok(capacities)
243 })
244 })
245 .join()
246 .map_err(|_| io::Error::other("layer-capacity worker panicked"))?
247}
248
249pub async fn validate_compact_chain(layers: &[CompactLayer]) -> io::Result<()> {
253 open_chain(layers).await?;
254 Ok(())
255}
256
257#[cfg(test)]
262mod tests {
263 use super::*;
264 use crate::checkpoint::create_qcow2_overlay;
265
266 #[tokio::test]
267 async fn compacted_prefix_preserves_overwrites_zeroes_and_grown_tail() {
268 let dir = tempfile::tempdir().unwrap();
269 let raw = dir.path().join("base.raw");
270 std::fs::write(&raw, vec![17u8; 131072]).unwrap();
271 let overlay = dir.path().join("layer.qcow2");
272 create_qcow2_overlay(&overlay, 262144, &raw, "raw")
273 .await
274 .unwrap();
275 let image = FormatAccess::new(
276 Qcow2::<ImagoFile>::builder(
277 ImagoFile::try_from(
278 std::fs::OpenOptions::new()
279 .read(true)
280 .write(true)
281 .open(&overlay)
282 .unwrap(),
283 )
284 .unwrap(),
285 )
286 .write(true)
287 .backing(None)
288 .data_file(None)
289 .open(DenyImplicitOpenGate::default())
290 .await
291 .unwrap(),
292 );
293 image.write(&vec![29u8; 65536][..], 0).await.unwrap();
294 image.write_zeroes(65536, 65536).await.unwrap();
295 image.flush().await.unwrap();
296 image.sync().await.unwrap();
297 drop(image);
298 let layers = vec![
299 CompactLayer {
300 path: raw.clone(),
301 qcow2: false,
302 },
303 CompactLayer {
304 path: overlay.clone(),
305 qcow2: true,
306 },
307 ];
308 let original = std::fs::read(&overlay).unwrap();
309 let destination = dir.path().join("compact.qcow2");
310 let result = materialize_compact_prefix(&layers, &destination)
311 .await
312 .unwrap();
313 assert_eq!(result.virtual_size, 262144);
314 let input = open_chain(&layers).await.unwrap();
315 let output = open_chain(&[CompactLayer {
316 path: destination.clone(),
317 qcow2: true,
318 }])
319 .await
320 .unwrap();
321 let mut before = vec![0; 262144];
322 let mut after = before.clone();
323 input.read(&mut before[..], 0).await.unwrap();
324 output.read(&mut after[..], 0).await.unwrap();
325 assert!(before == after, "compaction changed guest bytes");
326 assert!(after[..65536].iter().all(|byte| *byte == 29));
327 assert!(after[65536..].iter().all(|byte| *byte == 0));
328 let raw_destination = dir.path().join("flattened.raw");
329 let raw_result = materialize_raw_prefix(&layers, &raw_destination)
330 .await
331 .unwrap();
332 assert_eq!(raw_result.virtual_size, 262144);
333 assert_eq!(std::fs::read(&raw_destination).unwrap(), before);
334 assert!(
335 materialize_raw_prefix(&layers, &raw_destination)
336 .await
337 .is_err()
338 );
339 assert_eq!(std::fs::read(overlay).unwrap(), original);
340 assert!(
341 materialize_compact_prefix(&layers, &destination)
342 .await
343 .is_err()
344 );
345 }
346}