diff --git a/Cargo.lock b/Cargo.lock index 118b60cea..feb65da67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -597,6 +597,12 @@ dependencies = [ "windows-sys 0.61.0", ] +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + [[package]] name = "endi" version = "1.1.0" @@ -1031,6 +1037,15 @@ version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" @@ -2418,6 +2433,7 @@ name = "vm-migration" version = "0.1.0" dependencies = [ "anyhow", + "itertools", "serde", "serde_json", "thiserror 2.0.17", diff --git a/Cargo.toml b/Cargo.toml index 4b74245de..e1d2f158b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -145,6 +145,7 @@ dirs = "6.0.0" env_logger = "0.11.8" epoll = "4.4.0" flume = "0.11.1" +itertools = "0.14.0" libc = "0.2.177" log = "0.4.28" signal-hook = "0.3.18" diff --git a/vm-migration/Cargo.toml b/vm-migration/Cargo.toml index 7a8c9337b..69d57076f 100644 --- a/vm-migration/Cargo.toml +++ b/vm-migration/Cargo.toml @@ -6,6 +6,7 @@ version = "0.1.0" [dependencies] anyhow = { workspace = true } +itertools = { workspace = true } serde = { workspace = true, features = ["derive", "rc"] } serde_json = { workspace = true } thiserror = { workspace = true } diff --git a/vm-migration/src/protocol.rs b/vm-migration/src/protocol.rs index 2ed782ae3..8c320c68e 100644 --- a/vm-migration/src/protocol.rs +++ b/vm-migration/src/protocol.rs @@ -5,10 +5,12 @@ use std::io::{Read, Write}; +use itertools::Itertools; use serde::{Deserialize, Serialize}; use vm_memory::ByteValued; use crate::MigratableError; +use crate::bitpos_iterator::BitposIteratorExt; // Migration protocol // 1: Source establishes communication with destination (file socket or TCP connection.) @@ -211,38 +213,47 @@ pub struct MemoryRange { pub length: u64, } +impl MemoryRange { + /// Turn an iterator over the dirty bitmap into an iterator of dirty ranges. + pub fn dirty_ranges( + bitmap: impl IntoIterator, + start_addr: u64, + page_size: u64, + ) -> impl Iterator { + bitmap + .into_iter() + .bit_positions() + // Turn them into single-element ranges for coalesce. + .map(|b| b..(b + 1)) + // Merge adjacent ranges. + .coalesce(|prev, curr| { + if prev.end == curr.start { + Ok(prev.start..curr.end) + } else { + Err((prev, curr)) + } + }) + .map(move |r| Self { + gpa: start_addr + r.start * page_size, + length: (r.end - r.start) * page_size, + }) + } +} + #[derive(Clone, Default, Serialize, Deserialize)] pub struct MemoryRangeTable { data: Vec, } impl MemoryRangeTable { - pub fn from_bitmap(bitmap: Vec, start_addr: u64, page_size: u64) -> Self { - let mut table = MemoryRangeTable::default(); - let mut entry: Option = None; - for (i, block) in bitmap.iter().enumerate() { - for j in 0..64 { - let is_page_dirty = ((block >> j) & 1u64) != 0u64; - let page_offset = ((i * 64) + j) as u64 * page_size; - if is_page_dirty { - if let Some(entry) = &mut entry { - entry.length += page_size; - } else { - entry = Some(MemoryRange { - gpa: start_addr + page_offset, - length: page_size, - }); - } - } else if let Some(entry) = entry.take() { - table.push(entry); - } - } + pub fn from_bitmap( + bitmap: impl IntoIterator, + start_addr: u64, + page_size: u64, + ) -> Self { + Self { + data: MemoryRange::dirty_ranges(bitmap, start_addr, page_size).collect(), } - if let Some(entry) = entry.take() { - table.push(entry); - } - - table } pub fn regions(&self) -> &[MemoryRange] {