1use crate::{
13 raw::{
14 BTRFS_BLOCK_GROUP_ITEM_KEY, BTRFS_CHUNK_ITEM_KEY,
15 BTRFS_CHUNK_TREE_OBJECTID, BTRFS_EXTENT_TREE_OBJECTID,
16 BTRFS_FIRST_CHUNK_TREE_OBJECTID,
17 },
18 space::BlockGroupFlags,
19 tree_search::{SearchKey, tree_search},
20};
21use btrfs_disk::items::ChunkItem;
22use std::os::unix::io::BorrowedFd;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct DeviceAllocation {
32 pub devid: u64,
34 pub flags: BlockGroupFlags,
37 pub bytes: u64,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct ChunkEntry {
49 pub devid: u64,
51 pub physical_start: u64,
53 pub logical_start: u64,
55 pub length: u64,
58 pub flags: BlockGroupFlags,
60 pub used: u64,
63}
64
65pub fn device_chunk_allocations(
82 fd: BorrowedFd,
83) -> nix::Result<Vec<DeviceAllocation>> {
84 let mut allocs: Vec<DeviceAllocation> = Vec::new();
85
86 tree_search(
87 fd,
88 SearchKey::for_type(
89 u64::from(BTRFS_CHUNK_TREE_OBJECTID),
90 BTRFS_CHUNK_ITEM_KEY,
91 ),
92 |_hdr, data| {
93 if let Some((stripe_len, flags, stripes)) = parse_chunk(data) {
94 for devid in stripes {
95 accumulate(&mut allocs, devid, flags, stripe_len);
96 }
97 }
98 Ok(())
99 },
100 )?;
101
102 Ok(allocs)
103}
104
105pub fn chunk_list(fd: BorrowedFd) -> nix::Result<Vec<ChunkEntry>> {
119 let mut entries: Vec<ChunkEntry> = Vec::new();
120
121 tree_search(
122 fd,
123 SearchKey::for_objectid_range(
124 u64::from(BTRFS_CHUNK_TREE_OBJECTID),
125 BTRFS_CHUNK_ITEM_KEY,
126 u64::from(BTRFS_FIRST_CHUNK_TREE_OBJECTID),
127 u64::from(BTRFS_FIRST_CHUNK_TREE_OBJECTID),
128 ),
129 |hdr, data| {
130 if let Some(chunk) = ChunkItem::parse(data) {
131 let logical_start = hdr.offset;
132 let flags = BlockGroupFlags::from_bits_truncate(
133 chunk.chunk_type.bits(),
134 );
135 let used = block_group_used(fd, logical_start).unwrap_or(0);
136 for stripe in &chunk.stripes {
137 entries.push(ChunkEntry {
138 devid: stripe.devid,
139 physical_start: stripe.offset,
140 logical_start,
141 length: chunk.length,
142 flags,
143 used,
144 });
145 }
146 }
147 Ok(())
148 },
149 )?;
150
151 Ok(entries)
152}
153
154fn block_group_used(fd: BorrowedFd, logical_start: u64) -> Option<u64> {
159 let mut used: Option<u64> = None;
160 tree_search(
161 fd,
162 SearchKey {
163 tree_id: u64::from(BTRFS_EXTENT_TREE_OBJECTID),
164 min_objectid: logical_start,
165 max_objectid: logical_start,
166 min_type: BTRFS_BLOCK_GROUP_ITEM_KEY,
167 max_type: BTRFS_BLOCK_GROUP_ITEM_KEY,
168 min_offset: 0,
169 max_offset: u64::MAX,
170 min_transid: 0,
171 max_transid: u64::MAX,
172 },
173 |_hdr, data| {
174 if let Some(bg) = btrfs_disk::items::BlockGroupItem::parse(data) {
175 used = Some(bg.used);
176 }
177 Ok(())
178 },
179 )
180 .ok()?;
181 used
182}
183
184fn parse_chunk(data: &[u8]) -> Option<(u64, BlockGroupFlags, Vec<u64>)> {
189 let chunk = ChunkItem::parse(data)?;
190 let flags = BlockGroupFlags::from_bits_truncate(chunk.chunk_type.bits());
191 let devids: Vec<u64> = chunk.stripes.iter().map(|s| s.devid).collect();
192 Some((chunk.stripe_len, flags, devids))
193}
194
195fn accumulate(
198 allocs: &mut Vec<DeviceAllocation>,
199 devid: u64,
200 flags: BlockGroupFlags,
201 bytes: u64,
202) {
203 if let Some(entry) = allocs
204 .iter_mut()
205 .find(|a| a.devid == devid && a.flags == flags)
206 {
207 entry.bytes += bytes;
208 } else {
209 allocs.push(DeviceAllocation {
210 devid,
211 flags,
212 bytes,
213 });
214 }
215}
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220
221 fn build_chunk_buf(
224 length: u64,
225 stripe_len: u64,
226 type_bits: u64,
227 num_stripes: u16,
228 stripes: &[(u64, u64)], ) -> Vec<u8> {
230 let mut buf = Vec::new();
231 buf.extend_from_slice(&length.to_le_bytes());
232 buf.extend_from_slice(&0u64.to_le_bytes()); buf.extend_from_slice(&stripe_len.to_le_bytes());
234 buf.extend_from_slice(&type_bits.to_le_bytes()); buf.extend_from_slice(&4096u32.to_le_bytes()); buf.extend_from_slice(&4096u32.to_le_bytes()); buf.extend_from_slice(&4096u32.to_le_bytes()); buf.extend_from_slice(&num_stripes.to_le_bytes());
239 buf.extend_from_slice(&0u16.to_le_bytes()); for &(devid, offset) in stripes {
241 buf.extend_from_slice(&devid.to_le_bytes());
242 buf.extend_from_slice(&offset.to_le_bytes());
243 buf.extend_from_slice(&[0u8; 16]); }
245 buf
246 }
247
248 #[test]
249 fn parse_chunk_single_stripe() {
250 let data_flags = BlockGroupFlags::DATA.bits();
251 let buf = build_chunk_buf(1024 * 1024, 65536, data_flags, 1, &[(1, 0)]);
252 let (stripe_len, flags, devids) = parse_chunk(&buf).unwrap();
253 assert_eq!(stripe_len, 65536);
254 assert_eq!(flags, BlockGroupFlags::DATA);
255 assert_eq!(devids, vec![1]);
256 }
257
258 #[test]
259 fn parse_chunk_two_stripes() {
260 let flags_bits =
261 (BlockGroupFlags::DATA | BlockGroupFlags::RAID1).bits();
262 let buf = build_chunk_buf(
263 1 << 30,
264 1 << 30,
265 flags_bits,
266 2,
267 &[(1, 0), (2, 4096)],
268 );
269 let (_, flags, devids) = parse_chunk(&buf).unwrap();
270 assert_eq!(flags, BlockGroupFlags::DATA | BlockGroupFlags::RAID1);
271 assert_eq!(devids, vec![1, 2]);
272 }
273
274 #[test]
275 fn parse_chunk_too_short() {
276 let buf = vec![0u8; 10];
277 assert!(parse_chunk(&buf).is_none());
278 }
279
280 #[test]
281 fn parse_chunk_claims_more_stripes_than_fit() {
282 let buf = build_chunk_buf(1024, 1024, 0, 5, &[(1, 0)]);
284 let result = parse_chunk(&buf);
286 assert!(result.is_some());
287 let (_, _, devids) = result.unwrap();
288 assert_eq!(devids.len(), 1);
289 }
290
291 #[test]
294 fn accumulate_new_entry() {
295 let mut allocs = Vec::new();
296 accumulate(&mut allocs, 1, BlockGroupFlags::DATA, 1000);
297 assert_eq!(allocs.len(), 1);
298 assert_eq!(allocs[0].devid, 1);
299 assert_eq!(allocs[0].bytes, 1000);
300 }
301
302 #[test]
303 fn accumulate_merge_same_devid_flags() {
304 let mut allocs = Vec::new();
305 accumulate(&mut allocs, 1, BlockGroupFlags::DATA, 1000);
306 accumulate(&mut allocs, 1, BlockGroupFlags::DATA, 2000);
307 assert_eq!(allocs.len(), 1);
308 assert_eq!(allocs[0].bytes, 3000);
309 }
310
311 #[test]
312 fn accumulate_separate_different_flags() {
313 let mut allocs = Vec::new();
314 accumulate(&mut allocs, 1, BlockGroupFlags::DATA, 1000);
315 accumulate(&mut allocs, 1, BlockGroupFlags::METADATA, 2000);
316 assert_eq!(allocs.len(), 2);
317 }
318
319 #[test]
320 fn accumulate_separate_different_devids() {
321 let mut allocs = Vec::new();
322 accumulate(&mut allocs, 1, BlockGroupFlags::DATA, 1000);
323 accumulate(&mut allocs, 2, BlockGroupFlags::DATA, 2000);
324 assert_eq!(allocs.len(), 2);
325 }
326}