Skip to main content

arcbox_hypervisor/linux/vm/
dirty.rs

1use std::sync::atomic::Ordering;
2
3use crate::{
4    error::HypervisorError,
5    memory::{GuestAddress, PAGE_SIZE},
6    types::DirtyPageInfo,
7};
8
9use super::{KvmUserspaceMemoryRegion, KvmVm, MemorySlotInfo, ffi};
10
11impl KvmVm {
12    /// Adds an additional memory region to the VM.
13    pub fn add_memory_region(
14        &self,
15        guest_addr: GuestAddress,
16        host_addr: *mut u8,
17        size: u64,
18        read_only: bool,
19    ) -> Result<u32, HypervisorError> {
20        let slot = self.next_slot.fetch_add(1, Ordering::SeqCst);
21
22        let mut base_flags = 0u32;
23        if read_only {
24            base_flags |= ffi::KVM_MEM_READONLY;
25        }
26
27        let mut region_flags = base_flags;
28        if self.dirty_tracking_enabled.load(Ordering::SeqCst) {
29            region_flags |= ffi::KVM_MEM_LOG_DIRTY_PAGES;
30        }
31
32        let region = KvmUserspaceMemoryRegion {
33            slot,
34            flags: region_flags,
35            guest_phys_addr: guest_addr.raw(),
36            memory_size: size,
37            userspace_addr: host_addr as u64,
38        };
39
40        self.vm_fd.set_user_memory_region(&region).map_err(|e| {
41            HypervisorError::MemoryError(format!("Failed to add memory region: {}", e))
42        })?;
43
44        {
45            let mut slots = self
46                .memory_slots
47                .write()
48                .map_err(|_| HypervisorError::SnapshotError("Lock poisoned".to_string()))?;
49            slots.push(MemorySlotInfo {
50                slot,
51                guest_phys_addr: guest_addr.raw(),
52                size,
53                userspace_addr: host_addr as u64,
54            });
55        }
56
57        self.memory
58            .register_slot(slot, guest_addr.raw(), size, host_addr as u64, base_flags)?;
59
60        tracing::debug!(
61            "Added memory region {} at {}: {}MB, read_only={}",
62            slot,
63            guest_addr,
64            size / (1024 * 1024),
65            read_only
66        );
67
68        Ok(slot)
69    }
70
71    /// Removes a memory region from the VM.
72    pub fn remove_memory_region(&self, slot: u32) -> Result<(), HypervisorError> {
73        let region = KvmUserspaceMemoryRegion {
74            slot,
75            flags: 0,
76            guest_phys_addr: 0,
77            memory_size: 0,
78            userspace_addr: 0,
79        };
80
81        self.vm_fd.set_user_memory_region(&region).map_err(|e| {
82            HypervisorError::MemoryError(format!("Failed to remove memory region: {}", e))
83        })?;
84
85        {
86            let mut slots = self
87                .memory_slots
88                .write()
89                .map_err(|_| HypervisorError::SnapshotError("Lock poisoned".to_string()))?;
90            slots.retain(|entry| entry.slot != slot);
91        }
92
93        self.memory.unregister_slot(slot)?;
94
95        tracing::debug!("Removed memory region {}", slot);
96
97        Ok(())
98    }
99    /// Enables dirty page tracking for all memory regions.
100    ///
101    /// When enabled, KVM tracks which pages have been written to by the guest.
102    /// Use `get_dirty_pages` to retrieve and clear the dirty page bitmap.
103    ///
104    /// This is useful for:
105    /// - Live migration: Only transfer modified pages
106    /// - Snapshotting: Track incremental changes
107    ///
108    /// # Errors
109    ///
110    /// Returns an error if dirty logging cannot be enabled.
111    pub fn enable_dirty_tracking(&self) -> Result<(), HypervisorError> {
112        if self.dirty_tracking_enabled.load(Ordering::SeqCst) {
113            // Already enabled.
114            return Ok(());
115        }
116
117        let slots = self
118            .memory_slots
119            .read()
120            .map_err(|_| HypervisorError::SnapshotError("Lock poisoned".to_string()))?;
121
122        // Enable dirty logging for all memory slots.
123        for slot in slots.iter() {
124            self.vm_fd
125                .enable_dirty_logging(
126                    slot.slot,
127                    slot.guest_phys_addr,
128                    slot.size,
129                    slot.userspace_addr,
130                )
131                .map_err(|e| {
132                    HypervisorError::SnapshotError(format!(
133                        "Failed to enable dirty logging for slot {}: {}",
134                        slot.slot, e
135                    ))
136                })?;
137
138            tracing::debug!(
139                "Enabled dirty logging for slot {}: guest={:#x}, size={}MB",
140                slot.slot,
141                slot.guest_phys_addr,
142                slot.size / (1024 * 1024)
143            );
144        }
145
146        self.dirty_tracking_enabled.store(true, Ordering::SeqCst);
147        self.memory.set_dirty_tracking_enabled(true);
148        tracing::info!("Dirty page tracking enabled for VM {}", self.id);
149
150        Ok(())
151    }
152
153    /// Disables dirty page tracking for all memory regions.
154    ///
155    /// # Errors
156    ///
157    /// Returns an error if dirty logging cannot be disabled.
158    pub fn disable_dirty_tracking(&self) -> Result<(), HypervisorError> {
159        if !self.dirty_tracking_enabled.load(Ordering::SeqCst) {
160            // Already disabled.
161            return Ok(());
162        }
163
164        let slots = self
165            .memory_slots
166            .read()
167            .map_err(|_| HypervisorError::SnapshotError("Lock poisoned".to_string()))?;
168
169        // Disable dirty logging for all memory slots.
170        for slot in slots.iter() {
171            self.vm_fd
172                .disable_dirty_logging(
173                    slot.slot,
174                    slot.guest_phys_addr,
175                    slot.size,
176                    slot.userspace_addr,
177                )
178                .map_err(|e| {
179                    HypervisorError::SnapshotError(format!(
180                        "Failed to disable dirty logging for slot {}: {}",
181                        slot.slot, e
182                    ))
183                })?;
184        }
185
186        self.dirty_tracking_enabled.store(false, Ordering::SeqCst);
187        self.memory.set_dirty_tracking_enabled(false);
188        tracing::info!("Dirty page tracking disabled for VM {}", self.id);
189
190        Ok(())
191    }
192
193    /// Returns whether dirty page tracking is enabled.
194    #[must_use]
195    pub fn is_dirty_tracking_enabled(&self) -> bool {
196        self.dirty_tracking_enabled.load(Ordering::SeqCst)
197    }
198
199    /// Gets the list of dirty pages across all memory regions.
200    ///
201    /// This retrieves and clears the dirty page bitmap from KVM.
202    /// Each call returns pages that were written since the last call.
203    ///
204    /// # Errors
205    ///
206    /// Returns an error if dirty tracking is not enabled or if the
207    /// dirty log cannot be retrieved.
208    pub fn get_dirty_pages(&self) -> Result<Vec<DirtyPageInfo>, HypervisorError> {
209        if !self.dirty_tracking_enabled.load(Ordering::SeqCst) {
210            return Err(HypervisorError::SnapshotError(
211                "Dirty tracking not enabled".to_string(),
212            ));
213        }
214
215        let slots = self
216            .memory_slots
217            .read()
218            .map_err(|_| HypervisorError::SnapshotError("Lock poisoned".to_string()))?;
219
220        let mut dirty_pages = Vec::new();
221
222        for slot in slots.iter() {
223            // Get the dirty bitmap for this slot.
224            let bitmap = self
225                .vm_fd
226                .get_dirty_log(slot.slot, slot.size, PAGE_SIZE)
227                .map_err(|e| {
228                    HypervisorError::SnapshotError(format!(
229                        "Failed to get dirty log for slot {}: {}",
230                        slot.slot, e
231                    ))
232                })?;
233
234            // Parse the bitmap to extract dirty page addresses.
235            let pages = Self::parse_dirty_bitmap(&bitmap, slot.guest_phys_addr, slot.size);
236
237            tracing::debug!(
238                "Slot {}: {} dirty pages out of {} total",
239                slot.slot,
240                pages.len(),
241                slot.size / PAGE_SIZE
242            );
243
244            dirty_pages.extend(pages);
245        }
246
247        tracing::debug!(
248            "get_dirty_pages: found {} dirty pages total",
249            dirty_pages.len()
250        );
251
252        Ok(dirty_pages)
253    }
254
255    /// Parses a dirty bitmap to extract individual dirty page addresses.
256    ///
257    /// # Arguments
258    /// * `bitmap` - The bitmap from KVM_GET_DIRTY_LOG
259    /// * `base_addr` - The guest physical address of the region start
260    /// * `size` - Total size of the region
261    ///
262    /// # Returns
263    /// A vector of DirtyPageInfo for each dirty page.
264    pub(super) fn parse_dirty_bitmap(
265        bitmap: &[u64],
266        base_addr: u64,
267        size: u64,
268    ) -> Vec<DirtyPageInfo> {
269        let mut pages = Vec::new();
270        let num_pages = size / PAGE_SIZE;
271
272        for (word_idx, &word) in bitmap.iter().enumerate() {
273            if word == 0 {
274                // Skip words with no dirty pages.
275                continue;
276            }
277
278            // Check each bit in the word.
279            for bit_idx in 0..64 {
280                if (word >> bit_idx) & 1 != 0 {
281                    let page_num = (word_idx as u64 * 64) + bit_idx as u64;
282                    if page_num < num_pages {
283                        pages.push(DirtyPageInfo {
284                            guest_addr: base_addr + page_num * PAGE_SIZE,
285                            size: PAGE_SIZE,
286                        });
287                    }
288                }
289            }
290        }
291
292        pages
293    }
294}