Kernel: Add support for profiling kmalloc()/kfree()

Kernel/Heap/kmalloc.cpp

@@ -18,6 +18,7 @@
 #include <Kernel/Heap/kmalloc.h>
 #include <Kernel/KSyms.h>
 #include <Kernel/Panic.h>
+#include <Kernel/PerformanceManager.h>
 #include <Kernel/Process.h>
 #include <Kernel/Scheduler.h>
 #include <Kernel/SpinLock.h>
@@ -189,6 +190,7 @@ __attribute__((section(".heap"))) static u8 kmalloc_pool_heap[POOL_SIZE];
 static size_t g_kmalloc_bytes_eternal = 0;
 static size_t g_kmalloc_call_count;
 static size_t g_kfree_call_count;
+static size_t g_nested_kfree_calls;
 bool g_dump_kmalloc_stacks;
 
 static u8* s_next_eternal_ptr;
@@ -255,6 +257,12 @@ void* kmalloc(size_t size)
         PANIC("kmalloc: Out of memory (requested size: {})", size);
     }
 
+    Process* current_process = Process::current();
+    if (!current_process && Scheduler::colonel_initialized())
+        current_process = Scheduler::colonel();
+    if (current_process)
+        PerformanceManager::add_kmalloc_perf_event(*current_process, size, (FlatPtr)ptr);
+
     return ptr;
 }
 
@@ -266,8 +274,18 @@ void kfree(void* ptr)
     kmalloc_verify_nospinlock_held();
     ScopedSpinLock lock(s_lock);
     ++g_kfree_call_count;
+    ++g_nested_kfree_calls;
+
+    if (g_nested_kfree_calls == 1) {
+        Process* current_process = Process::current();
+        if (!current_process && Scheduler::colonel_initialized())
+            current_process = Scheduler::colonel();
+        if (current_process)
+            PerformanceManager::add_kfree_perf_event(*current_process, 0, (FlatPtr)ptr);
+    }
 
     g_kmalloc_global->m_heap.deallocate(ptr);
+    --g_nested_kfree_calls;
 }
 
 void* krealloc(void* ptr, size_t new_size)

Kernel/PerformanceEventBuffer.cpp

@@ -114,6 +114,14 @@ KResult PerformanceEventBuffer::append_with_eip_and_ebp(ProcessID pid, ThreadID
         event.data.context_switch.next_pid = arg1;
         event.data.context_switch.next_tid = arg2;
         break;
+    case PERF_EVENT_KMALLOC:
+        event.data.kmalloc.size = arg1;
+        event.data.kmalloc.ptr = arg2;
+        break;
+    case PERF_EVENT_KFREE:
+        event.data.kfree.size = arg1;
+        event.data.kfree.ptr = arg2;
+        break;
     default:
         return EINVAL;
     }
@@ -192,6 +200,16 @@ bool PerformanceEventBuffer::to_json_impl(Serializer& object) const
             event_object.add("next_pid", static_cast<u64>(event.data.context_switch.next_pid));
             event_object.add("next_tid", static_cast<u64>(event.data.context_switch.next_tid));
             break;
+        case PERF_EVENT_KMALLOC:
+            event_object.add("type", "kmalloc");
+            event_object.add("ptr", static_cast<u64>(event.data.kmalloc.ptr));
+            event_object.add("size", static_cast<u64>(event.data.kmalloc.size));
+            break;
+        case PERF_EVENT_KFREE:
+            event_object.add("type", "kfree");
+            event_object.add("ptr", static_cast<u64>(event.data.kfree.ptr));
+            event_object.add("size", static_cast<u64>(event.data.kfree.size));
+            break;
         }
         event_object.add("pid", event.pid);
         event_object.add("tid", event.tid);

Kernel/PerformanceEventBuffer.h

@@ -52,6 +52,16 @@ struct [[gnu::packed]] ContextSwitchPerformanceEvent {
     u32 next_tid;
 };
 
+struct [[gnu::packed]] KMallocPerformanceEvent {
+    size_t size;
+    FlatPtr ptr;
+};
+
+struct [[gnu::packed]] KFreePerformanceEvent {
+    size_t size;
+    FlatPtr ptr;
+};
+
 struct [[gnu::packed]] PerformanceEvent {
     u16 type { 0 };
     u8 stack_size { 0 };
@@ -68,6 +78,8 @@ struct [[gnu::packed]] PerformanceEvent {
         ProcessExecPerformanceEvent process_exec;
         ThreadCreatePerformanceEvent thread_create;
         ContextSwitchPerformanceEvent context_switch;
+        KMallocPerformanceEvent kmalloc;
+        KFreePerformanceEvent kfree;
     } data;
     static constexpr size_t max_stack_frame_count = 64;
     FlatPtr stack[max_stack_frame_count];

Kernel/PerformanceManager.h

@@ -92,6 +92,20 @@ public:
         }
     }
 
+    inline static void add_kmalloc_perf_event(Process& current_process, size_t size, FlatPtr ptr)
+    {
+        if (auto* event_buffer = current_process.current_perf_events_buffer()) {
+            [[maybe_unused]] auto res = event_buffer->append(PERF_EVENT_KMALLOC, size, ptr, nullptr);
+        }
+    }
+
+    inline static void add_kfree_perf_event(Process& current_process, size_t size, FlatPtr ptr)
+    {
+        if (auto* event_buffer = current_process.current_perf_events_buffer()) {
+            [[maybe_unused]] auto res = event_buffer->append(PERF_EVENT_KFREE, size, ptr, nullptr);
+        }
+    }
+
     inline static void timer_tick(RegisterState const& regs)
     {
         static Time last_wakeup;

Kernel/Scheduler.cpp

@@ -442,6 +442,11 @@ void Scheduler::prepare_for_idle_loop()
     scheduler_data.m_in_scheduler = true;
 }
 
+bool Scheduler::colonel_initialized()
+{
+    return !!s_colonel_process;
+}
+
 Process* Scheduler::colonel()
 {
     VERIFY(s_colonel_process);

Kernel/Scheduler.h

@@ -43,6 +43,7 @@ public:
     static void leave_on_first_switch(u32 flags);
     static void prepare_after_exec();
     static void prepare_for_idle_loop();
+    static bool colonel_initialized();
     static Process* colonel();
     static void idle_loop(void*);
     static void invoke_async();

Kernel/Syscalls/execve.cpp

@@ -530,7 +530,13 @@ KResult Process::do_exec(NonnullRefPtr<FileDescription> main_program_description
 
     set_dumpable(!executable_is_setid);
 
-    m_space = load_result.space.release_nonnull();
+    {
+        // We must disable global profiling (especially kfree tracing) here because
+        // we might otherwise end up walking the stack into the process' space that
+        // is about to be destroyed.
+        TemporaryChange global_profiling_disabler(g_profiling_all_threads, false);
+        m_space = load_result.space.release_nonnull();
+    }
     MemoryManager::enter_space(*m_space);
 
     auto signal_trampoline_region = m_space->allocate_region_with_vmobject(signal_trampoline_range.value(), g_signal_trampoline_region->vmobject(), 0, "Signal trampoline", PROT_READ | PROT_EXEC, true);

Kernel/UnixTypes.h

@@ -58,6 +58,8 @@ enum {
     PERF_EVENT_THREAD_CREATE = 256,
     PERF_EVENT_THREAD_EXIT = 512,
     PERF_EVENT_CONTEXT_SWITCH = 1024,
+    PERF_EVENT_KMALLOC = 2048,
+    PERF_EVENT_KFREE = 4096,
 };
 
 #define WNOHANG 1

Userland/Libraries/LibC/serenity.h

@@ -87,6 +87,8 @@ enum {
     PERF_EVENT_THREAD_CREATE = 256,
     PERF_EVENT_THREAD_EXIT = 512,
     PERF_EVENT_CONTEXT_SWITCH = 1024,
+    PERF_EVENT_KMALLOC = 2048,
+    PERF_EVENT_KFREE = 4096,
 };
 
 #define PERF_EVENT_MASK_ALL (~0ull)

Userland/Utilities/profile.cpp

@@ -40,6 +40,10 @@ int main(int argc, char** argv)
                 event_mask |= PERF_EVENT_SAMPLE;
             else if (event_type == "context_switch")
                 event_mask |= PERF_EVENT_CONTEXT_SWITCH;
+            else if (event_type == "kmalloc")
+                event_mask |= PERF_EVENT_KMALLOC;
+            else if (event_type == "kfree")
+                event_mask |= PERF_EVENT_KFREE;
             else {
                 warnln("Unknown event type '{}' specified.", event_type);
                 exit(1);
@@ -49,7 +53,7 @@ int main(int argc, char** argv)
 
     auto print_types = [] {
         outln();
-        outln("Event type can be one of: sample and context_switch.");
+        outln("Event type can be one of: sample, context_switch, kmalloc and kfree.");
     };
 
     if (!args_parser.parse(argc, argv, false)) {