Extract generic MemoryUsageHandler to use as base class for ust&kernel

tracetools_analysis/tracetools_analysis/processor/memory_usage.py

@@ -24,8 +24,33 @@ from ..data_model.memory_usage import MemoryUsageDataModel
 
 
 class MemoryUsageHandler(EventHandler):
+    """Generic handler for memory usage."""
+
+    def __init__(
+        self,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self._data_model = MemoryUsageDataModel()
+
+    @property
+    def data(self) -> MemoryUsageDataModel:
+        return self._data_model
+
+    def _update(
+        self,
+        timestamp: int,
+        tid: int,
+        memory_difference: int,
+    ) -> None:
+        # Add to data model
+        self.data.add_memory_difference(timestamp, tid, memory_difference)
+
+
+class UserspaceMemoryUsageHandler(MemoryUsageHandler):
     """
-    Handler that extracts data for memory usage.
+    Handler that extracts userspace memory usage data.
 
     It uses the following events:
         * lttng_ust_libc:malloc
@@ -66,17 +91,11 @@ class MemoryUsageHandler(EventHandler):
             **kwargs,
         )
 
-        self._data_model = MemoryUsageDataModel()
-
         # Temporary buffers
         # pointer -> current memory size
         # (used to know keep track of the memory size allocated at a given pointer)
         self._memory: Dict[int, int] = {}
 
-    @property
-    def data(self) -> MemoryUsageDataModel:
-        return self._data_model
-
     def _handle_malloc(
         self, event: Dict, metadata: EventMetadata
     ) -> None:
@@ -148,5 +167,59 @@ class MemoryUsageHandler(EventHandler):
             if allocated_memory is not None:
                 memory_difference = -allocated_memory
 
-        # Add to data model
+        self._update(timestamp, tid, memory_difference)
-        self.data.add_memory_difference(timestamp, tid, memory_difference)
+
+
+class KernelMemoryUsageHandler(MemoryUsageHandler):
+    """
+    Handler that extracts userspace memory usage data.
+
+    It uses the following events:
+        * kmem_mm_page_alloc
+        * kmem_mm_page_free
+
+    Implementation inspired by Trace Compass' implementation:
+    https://git.eclipse.org/c/tracecompass/org.eclipse.tracecompass.git/tree/analysis/org.eclipse.tracecompass.analysis.os.linux.core/src/org/eclipse/tracecompass/analysis/os/linux/core/kernelmemoryusage/KernelMemoryStateProvider.java#n84
+    """
+
+    PAGE_SIZE = 4096
+
+    def __init__(
+        self,
+        **kwargs,
+    ) -> None:
+        # Link event to handling method
+        handler_map = {
+            'kmem_mm_page_alloc':
+                self._handle_malloc,
+            'kmem_mm_page_free':
+                self._handle_free,
+        }
+        super().__init__(
+            handler_map=handler_map,
+            **kwargs,
+        )
+
+    def _handle_malloc(
+        self, event: Dict, metadata: EventMetadata
+    ) -> None:
+        self._handle(event, metadata, self.PAGE_SIZE)
+
+    def _handle_free(
+        self, event: Dict, metadata: EventMetadata
+    ) -> None:
+        self._handle(event, metadata, -self.PAGE_SIZE)
+
+    def _handle(
+        self,
+        event: Dict,
+        metadata: EventMetadata,
+        inc: int,
+    ) -> None:
+        order = get_field(event, 'order')
+        inc <<= order
+
+        timestamp = metadata.timestamp
+        tid = metadata.tid
+
+        self._update(timestamp, tid, inc)

tracetools_analysis/tracetools_analysis/scripts/memory_usage.py

@@ -18,7 +18,7 @@ import pandas as pd
 
 from tracetools_analysis.loading import load_file
 from tracetools_analysis.processor import Processor
-from tracetools_analysis.processor.memory_usage import MemoryUsageHandler
+from tracetools_analysis.processor.memory_usage import UserspaceMemoryUsageHandler
 from tracetools_analysis.processor.ros2 import Ros2Handler
 from tracetools_analysis.utils.memory_usage import MemoryUsageDataModelUtil
 from tracetools_analysis.utils.ros2 import Ros2DataModelUtil
@@ -43,14 +43,14 @@ def main():
     file_path = sys.argv[1]
 
     events = load_file(file_path)
-    memory_handler = MemoryUsageHandler()
+    ust_memory_handler = UserspaceMemoryUsageHandler()
     ros2_handler = Ros2Handler()
-    Processor(memory_handler, ros2_handler).process(events)
+    Processor(ust_memory_handler, ros2_handler).process(events)
 
-    memory_data_util = MemoryUsageDataModelUtil(memory_handler.data)
+    ust_memory_data_util = MemoryUsageDataModelUtil(ust_memory_handler.data)
     ros2_data_util = Ros2DataModelUtil(ros2_handler.data)
 
-    memory_usage_dfs = memory_data_util.get_absolute_memory_usage_by_tid()
+    ust_memory_usage_dfs = ust_memory_data_util.get_absolute_userspace_memory_usage_by_tid()
     tids = ros2_data_util.get_tids()
 
     data = [
@@ -59,7 +59,7 @@ def main():
             ros2_data_util.get_node_names_from_tid(tid),
             format_memory_size(memory_usage['memory_usage'].max(), precision=1),
         ]
-        for tid, memory_usage in memory_usage_dfs.items()
+        for tid, memory_usage in ust_memory_usage_dfs.items()
         if tid in tids
     ]
 

tracetools_analysis/tracetools_analysis/utils/memory_usage.py

@@ -28,21 +28,47 @@ class MemoryUsageDataModelUtil(DataModelUtil):
 
     def __init__(
         self,
-        data_model: MemoryUsageDataModel,
+        *,
+        data_model_userspace: MemoryUsageDataModel = None,
+        data_model_kernel: MemoryUsageDataModel = None,
     ) -> None:
         """
         Create a MemoryUsageDataModelUtil.
 
-        :param data_model: the data model object to use
+        At least one non-`None` `MemoryUsageDataModel` must be given
-        """
-        super().__init__(data_model)
 
-    def get_absolute_memory_usage_by_tid(self) -> Dict[int, DataFrame]:
+        :param data_model_userspace: the userspace data model object to use
+        :param data_model_kernel: the kernel data model object to use
         """
-        Get absolute memory usage over time per tid.
+        # Not giving any model to the base class; we'll own them ourselves
+        super().__init__(None)
+
+        if data_model_userspace is None and data_model_kernel is None:
+            raise RuntimeError('must provide at least one (userspace or kernel) data model!')
+
+        self.data_ust = data_model_userspace
+        self.data_kernel = data_model_kernel
+
+    def get_absolute_userspace_memory_usage_by_tid(self) -> Dict[int, DataFrame]:
+        """
+        Get absolute userspace memory usage over time per tid.
 
         :return (tid -> DataFrame of absolute memory usage over time)
         """
+        return self._get_absolute_memory_usage_by_tid(self.data_ust)
+
+    def get_absolute_kernel_memory_usage_by_tid(self) -> Dict[int, DataFrame]:
+        """
+        Get absolute kernel memory usage over time per tid.
+
+        :return (tid -> DataFrame of absolute memory usage over time)
+        """
+        return self._get_absolute_memory_usage_by_tid(self.data_kernel)
+
+    def _get_absolute_memory_usage_by_tid(
+        self,
+        data: MemoryUsageDataModel,
+    ) -> Dict[int, DataFrame]:
         previous = defaultdict(int)
         data = defaultdict(list)
         for index, row in self.data.memory_diff.iterrows():