tests: fix MMIO gaps in memory monitor tool

The memory monitor was only assuming a single MMIO gap on x86_64 when
calculating the memory regions that corresponded to guest memory. Now we
need to account for two MMIO gaps in x86 and one in ARM.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

tests/host_tools/memory.py

@@ -8,15 +8,17 @@
 
 import psutil
 
+from framework.properties import global_props
+
 
 class MemoryUsageExceededError(Exception):
     """A custom exception containing details on excessive memory usage."""
 
     def __init__(self, usage, threshold, *args):
         """Compose the error message containing the memory consumption."""
         super().__init__(
-            f"Memory usage ({usage / 2**20:.2f} MiB) exceeded maximum threshold "
-            f"({threshold / 2**20} MiB)",
+            f"Memory usage ({usage / 1 << 20:.2f} MiB) exceeded maximum threshold "
+            f"({threshold / 1 << 20} MiB)",
             *args,
         )
 
@@ -28,10 +30,20 @@ class MemoryMonitor(Thread):
     VMM memory usage.
     """
 
-    # If guest memory is >3328MB, it is split in a 2nd region
-    X86_MEMORY_GAP_START = 3328 * 2**20
-
-    def __init__(self, vm, threshold=5 * 2**20, period_s=0.05):
+    # If guest memory is >3GiB, it is split in a 2nd region
+    # Gap starts at 3GiBs and is 1GiB long
+    X86_32BIT_MEMORY_GAP_START = 3 << 30
+    X86_32BIT_MEMORY_GAP_SIZE = 1 << 30
+    # If guest memory is >255GiB, it is split in a 3rd region
+    # Gap starts at 256 GiB and is 256GiB long
+    X86_64BIT_MEMORY_GAP_START = 256 << 30
+    # On ARM64 we just have a single gap, but memory starts at an offset
+    # Gap starts at 256 GiB and is GiB long
+    # Memory starts at 2GiB
+    ARM64_64BIT_MEMORY_GAP_START = 256 << 30
+    ARM64_MEMORY_START = 2 << 30
+
+    def __init__(self, vm, threshold=5 << 20, period_s=0.01):
         """Initialize monitor attributes."""
         Thread.__init__(self)
         self._vm = vm
@@ -72,7 +84,9 @@ def run(self):
             mem_total = 0
             for mmap in mmaps:
                 if self.is_guest_mem(mmap.size, guest_mem_bytes):
+                    print(f"Region {mmap} is guest memory")
                     continue
+
                 mem_total += mmap.rss
             self._current_rss = mem_total
             if mem_total > self.threshold:
@@ -81,24 +95,55 @@ def run(self):
 
             time.sleep(self._period_s)
 
-    def is_guest_mem(self, size, guest_mem_bytes):
+    def is_guest_mem_x86(self, size, guest_mem_bytes):
         """
-        If the address is recognised as a guest memory region,
-        return True, otherwise return False.
+        Checks if a region is a guest memory region based on
+        x86_64 physical memory layout
         """
+        return size in (
+            # memory fits before the first gap
+            guest_mem_bytes,
+            # guest memory spans at least two regions & memory fits before the second gap
+            self.X86_32BIT_MEMORY_GAP_START,
+            # guest memory spans exactly two regions
+            guest_mem_bytes - self.X86_32BIT_MEMORY_GAP_START,
+            # guest memory fills the space between the two gaps
+            self.X86_64BIT_MEMORY_GAP_START
+            - self.X86_32BIT_MEMORY_GAP_START
+            - self.X86_32BIT_MEMORY_GAP_SIZE,
+            # guest memory spans 3 regions, this is what remains past the second gap
+            guest_mem_bytes
+            - self.X86_64BIT_MEMORY_GAP_START
+            + self.X86_32BIT_MEMORY_GAP_SIZE,
+        )
 
-        # If x86_64 guest memory exceeds 3328M, it will be split
-        # in 2 regions: 3328M and the rest. We have 3 cases here
-        # to recognise a guest memory region:
-        #  - its size matches the guest memory exactly
-        #  - its size is 3328M
-        #  - its size is guest memory minus 3328M.
+    def is_guest_mem_arch64(self, size, guest_mem_bytes):
+        """
+        Checks if a region is a guest memory region based on
+        ARM64 physical memory layout
+        """
         return size in (
+            # guest memory fits before the gap
             guest_mem_bytes,
-            self.X86_MEMORY_GAP_START,
-            guest_mem_bytes - self.X86_MEMORY_GAP_START,
+            # guest memory fills the space before the gap
+            self.ARM64_64BIT_MEMORY_GAP_START - self.ARM64_MEMORY_START,
+            # guest memory spans 2 regions, this is what remains past the gap
+            guest_mem_bytes
+            - self.ARM64_64BIT_MEMORY_GAP_START
+            + self.ARM64_MEMORY_START,
         )
 
+    def is_guest_mem(self, size, guest_mem_bytes):
+        """
+        If the address is recognised as a guest memory region,
+        return True, otherwise return False.
+        """
+
+        if global_props.cpu_architecture == "x86_64":
+            return self.is_guest_mem_x86(size, guest_mem_bytes)
+
+        return self.is_guest_mem_arch64(size, guest_mem_bytes)
+
     def check_samples(self):
         """Check that there are no samples over the threshold."""
         if self._exceeded is not None: