commit b70bdb5052a3d92e5b24cb58e4e17e6ce9b0bd59
author Alice Wang <aliceywang@google.com> Mon Jun 12 08:17:58 2023 +0000
committer Alice Wang <aliceywang@google.com> Mon Jun 12 13:17:07 2023 +0000
tree 699be4c615d885d6e85889aa302f2ba084affb9d
parent e9ae5c44e599bfd3c24c6258dacfb3e872752856

[rialto] Set up and tear down the global memory tracker in rialto

Bug: 284462758
Test: atest rialto_test
Change-Id: Ie69341845d07b7888d646fed3fe279b30df0f75c

rialto/src/main.rs

diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 11ca570..44e83ee 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -31,7 +31,7 @@
 use vmbase::{
     layout::{self, crosvm},
     main,
-    memory::{PageTable, PAGE_SIZE},
+    memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE},
     power::reboot,
 };
 
@@ -50,7 +50,7 @@
     }
 }
 
-fn init_page_table() -> Result<()> {
+fn new_page_table() -> Result<PageTable> {
     let mut page_table = PageTable::default();
 
     page_table.map_device(&crosvm::MMIO_RANGE)?;
@@ -60,11 +60,7 @@
     page_table.map_rodata(&layout::rodata_range())?;
     page_table.map_device(&layout::console_uart_range())?;
 
-    // SAFETY: It is safe to activate the page table by setting `TTBR0_EL1` to point to
-    // it as this is the first time we activate the page table.
-    unsafe { page_table.activate() }
-    info!("Activated kernel page table.");
-    Ok(())
+    Ok(page_table)
 }
 
 fn try_init_logger() -> Result<bool> {
@@ -94,20 +90,26 @@
     let pci_info = PciInfo::from_fdt(fdt)?;
     debug!("PCI: {:#x?}", pci_info);
 
-    init_page_table()?;
+    let page_table = new_page_table()?;
+
+    MEMORY.lock().replace(MemoryTracker::new(
+        page_table,
+        crosvm::MEM_START..layout::MAX_VIRT_ADDR,
+        crosvm::MMIO_RANGE,
+        None, // Rialto doesn't have any payload for now.
+    ));
     Ok(())
 }
 
 fn try_unshare_all_memory(mmio_guard_supported: bool) -> Result<()> {
-    if !mmio_guard_supported {
-        return Ok(());
-    }
     info!("Starting unsharing memory...");
 
-    // TODO(b/284462758): Unshare all the memory here.
-
     // No logging after unmapping UART.
-    get_hypervisor().mmio_guard_unmap(vmbase::console::BASE_ADDRESS)?;
+    if mmio_guard_supported {
+        get_hypervisor().mmio_guard_unmap(vmbase::console::BASE_ADDRESS)?;
+    }
+    // Unshares all memory and deactivates page table.
+    drop(MEMORY.lock().take());
     Ok(())
 }