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};
14
15type SharedImage = Arc<FormatAccess<Box<dyn DynStorage>>>;
20
21#[derive(Clone, Debug)]
23pub struct CompactLayer {
24 pub path: PathBuf,
26 pub qcow2: bool,
28}
29
30#[derive(Clone, Debug)]
32pub struct CompactMaterialization {
33 pub virtual_size: u64,
35 pub materialized_bytes: u64,
37}
38
39async fn open_chain(layers: &[CompactLayer]) -> io::Result<SharedImage> {
45 open_chain_access(layers, false).await
46}
47
48pub(crate) async fn open_writable_chain(layers: &[CompactLayer]) -> io::Result<SharedImage> {
50 open_chain_access(layers, true).await
51}
52
53async fn open_chain_access(
54 layers: &[CompactLayer],
55 writable_head: bool,
56) -> io::Result<SharedImage> {
57 if layers.is_empty() || layers.iter().skip(1).any(|layer| !layer.qcow2) {
58 return Err(io::Error::new(
59 io::ErrorKind::InvalidInput,
60 "invalid raw/qcow2 compaction prefix",
61 ));
62 }
63 let mut backing: Option<SharedImage> = None;
64 for (index, layer) in layers.iter().enumerate() {
65 let writable = writable_head && index + 1 == layers.len();
66 let storage: Box<dyn DynStorage> = Box::new(ImagoFile::try_from(
67 std::fs::OpenOptions::new()
68 .read(true)
69 .write(writable)
70 .open(&layer.path)?,
71 )?);
72 backing = Some(if layer.qcow2 {
73 let image = Qcow2::<Box<dyn DynStorage>, SharedImage>::builder(storage)
74 .write(writable)
75 .backing(backing)
76 .data_file(None)
77 .open(DenyImplicitOpenGate::default())
78 .await?;
79 if image.requires_external_data_file() {
80 return Err(io::Error::new(
81 io::ErrorKind::Unsupported,
82 "external qcow2 data files are not supported",
83 ));
84 }
85 Arc::new(FormatAccess::new(image))
86 } else {
87 Arc::new(FormatAccess::new(
88 Raw::<Box<dyn DynStorage>>::builder(storage)
89 .write(writable)
90 .open(DenyImplicitOpenGate::default())
91 .await?,
92 ))
93 });
94 }
95 Ok(backing.expect("nonempty chain checked above"))
96}
97
98pub async fn materialize_compact_prefix(
104 layers: &[CompactLayer],
105 destination: &Path,
106) -> io::Result<CompactMaterialization> {
107 let source = open_chain(layers).await?;
108 let virtual_size = source.size();
109 if virtual_size == 0 || !virtual_size.is_multiple_of(512) {
110 return Err(io::Error::new(
111 io::ErrorKind::InvalidData,
112 "invalid compaction disk capacity",
113 ));
114 }
115 let file = std::fs::OpenOptions::new()
116 .read(true)
117 .write(true)
118 .create_new(true)
119 .open(destination)?;
120 Qcow2::<ImagoFile>::create_builder(ImagoFile::try_from(file)?)
121 .size(virtual_size)
122 .cluster_size(65536)
123 .create()
124 .await?;
125 let file = std::fs::OpenOptions::new()
126 .read(true)
127 .write(true)
128 .open(destination)?;
129 let target = FormatAccess::new(
130 Qcow2::<ImagoFile>::builder(ImagoFile::try_from(file)?)
131 .write(true)
132 .backing(None)
133 .data_file(None)
134 .open(DenyImplicitOpenGate::default())
135 .await?,
136 );
137 let mut buffer = vec![0u8; 1024 * 1024];
138 let mut offset = 0;
139 let mut materialized_bytes = 0;
140 while offset < virtual_size {
141 let (mapping, length) = source.get_mapping(offset, virtual_size - offset).await?;
142 if length == 0 {
143 return Err(io::Error::new(
144 io::ErrorKind::InvalidData,
145 "compaction mapping made no progress",
146 ));
147 }
148 if matches!(mapping, Mapping::Zero { .. }) {
149 offset += length;
150 continue;
151 }
152 let count = length.min(buffer.len() as u64) as usize;
153 source.read(&mut buffer[..count], offset).await?;
154 if buffer[..count].iter().any(|byte| *byte != 0) {
157 target.write(&buffer[..count], offset).await?;
158 materialized_bytes += count as u64;
159 }
160 offset += count as u64;
161 }
162 target.flush().await?;
163 target.sync().await?;
164 Ok(CompactMaterialization {
165 virtual_size,
166 materialized_bytes,
167 })
168}
169
170pub async fn compact_layer_capacity(layer: CompactLayer) -> io::Result<u64> {
172 Ok(open_chain(&[layer]).await?.size())
173}
174
175pub fn layer_capacities(layers: Vec<CompactLayer>) -> io::Result<Vec<u64>> {
178 std::thread::spawn(move || {
179 tokio::runtime::Builder::new_current_thread()
180 .enable_all()
181 .build()?
182 .block_on(async {
183 let mut capacities = Vec::with_capacity(layers.len());
184 for layer in layers {
185 capacities.push(compact_layer_capacity(layer).await?);
186 }
187 Ok(capacities)
188 })
189 })
190 .join()
191 .map_err(|_| io::Error::other("layer-capacity worker panicked"))?
192}
193
194pub async fn validate_compact_chain(layers: &[CompactLayer]) -> io::Result<()> {
198 open_chain(layers).await?;
199 Ok(())
200}
201
202#[cfg(test)]
207mod tests {
208 use super::*;
209 use crate::checkpoint::create_qcow2_overlay;
210
211 #[tokio::test]
212 async fn compacted_prefix_preserves_overwrites_zeroes_and_grown_tail() {
213 let dir = tempfile::tempdir().unwrap();
214 let raw = dir.path().join("base.raw");
215 std::fs::write(&raw, vec![17u8; 131072]).unwrap();
216 let overlay = dir.path().join("layer.qcow2");
217 create_qcow2_overlay(&overlay, 262144, &raw, "raw")
218 .await
219 .unwrap();
220 let image = FormatAccess::new(
221 Qcow2::<ImagoFile>::builder(
222 ImagoFile::try_from(
223 std::fs::OpenOptions::new()
224 .read(true)
225 .write(true)
226 .open(&overlay)
227 .unwrap(),
228 )
229 .unwrap(),
230 )
231 .write(true)
232 .backing(None)
233 .data_file(None)
234 .open(DenyImplicitOpenGate::default())
235 .await
236 .unwrap(),
237 );
238 image.write(&vec![29u8; 65536][..], 0).await.unwrap();
239 image.write_zeroes(65536, 65536).await.unwrap();
240 image.flush().await.unwrap();
241 image.sync().await.unwrap();
242 drop(image);
243 let layers = vec![
244 CompactLayer {
245 path: raw.clone(),
246 qcow2: false,
247 },
248 CompactLayer {
249 path: overlay.clone(),
250 qcow2: true,
251 },
252 ];
253 let original = std::fs::read(&overlay).unwrap();
254 let destination = dir.path().join("compact.qcow2");
255 let result = materialize_compact_prefix(&layers, &destination)
256 .await
257 .unwrap();
258 assert_eq!(result.virtual_size, 262144);
259 let input = open_chain(&layers).await.unwrap();
260 let output = open_chain(&[CompactLayer {
261 path: destination.clone(),
262 qcow2: true,
263 }])
264 .await
265 .unwrap();
266 let mut before = vec![0; 262144];
267 let mut after = before.clone();
268 input.read(&mut before[..], 0).await.unwrap();
269 output.read(&mut after[..], 0).await.unwrap();
270 assert!(before == after, "compaction changed guest bytes");
271 assert!(after[..65536].iter().all(|byte| *byte == 29));
272 assert!(after[65536..].iter().all(|byte| *byte == 0));
273 assert_eq!(std::fs::read(overlay).unwrap(), original);
274 assert!(
275 materialize_compact_prefix(&layers, &destination)
276 .await
277 .is_err()
278 );
279 }
280}