Kernel: Track KCOVInstance via Process instead of HashMap

While this clutters Process.cpp a tiny bit, I feel that it's worth it:
- 2x speed on the kcov_loop benchmark. Likely more during fuzzing.
- Overall code complexity is going down with this change.
- By reducing the code reachable from __sanitizer_cov_trace_pc code,
  we can now instrument more code.

Kernel/Devices/KCOVDevice.cpp

@@ -16,9 +16,6 @@
 
 namespace Kernel {
 
-HashMap<ProcessID, KCOVInstance*>* KCOVDevice::proc_instance;
-HashMap<ThreadID, KCOVInstance*>* KCOVDevice::thread_instance;
-
 UNMAP_AFTER_INIT NonnullLockRefPtr<KCOVDevice> KCOVDevice::must_create()
 {
     auto kcov_device_or_error = DeviceManagement::try_create_device<KCOVDevice>();
@@ -30,85 +27,45 @@ UNMAP_AFTER_INIT NonnullLockRefPtr<KCOVDevice> KCOVDevice::must_create()
 UNMAP_AFTER_INIT KCOVDevice::KCOVDevice()
     : BlockDevice(30, 0)
 {
-    proc_instance = new HashMap<ProcessID, KCOVInstance*>();
-    thread_instance = new HashMap<ThreadID, KCOVInstance*>();
     dbgln("KCOVDevice created");
 }
 
-void KCOVDevice::free_thread()
-{
-    auto thread = Thread::current();
-    auto tid = thread->tid();
-
-    auto maybe_kcov_instance = thread_instance->get(tid);
-    if (!maybe_kcov_instance.has_value())
-        return;
-
-    auto kcov_instance = maybe_kcov_instance.value();
-    VERIFY(kcov_instance->state() == KCOVInstance::TRACING);
-    kcov_instance->set_state(KCOVInstance::OPENED);
-    thread_instance->remove(tid);
-}
-
-void KCOVDevice::free_process()
-{
-    auto pid = Process::current().pid();
-
-    auto maybe_kcov_instance = proc_instance->get(pid);
-    if (!maybe_kcov_instance.has_value())
-        return;
-
-    auto kcov_instance = maybe_kcov_instance.value();
-    VERIFY(kcov_instance->state() == KCOVInstance::OPENED);
-    kcov_instance->set_state(KCOVInstance::UNUSED);
-    proc_instance->remove(pid);
-    delete kcov_instance;
-}
-
 ErrorOr<NonnullRefPtr<OpenFileDescription>> KCOVDevice::open(int options)
 {
-    auto pid = Process::current().pid();
-    if (proc_instance->get(pid).has_value())
+    auto& proc = Process::current();
+    auto* kcov_instance = proc.kcov_instance();
+    if (kcov_instance != nullptr)
         return EBUSY; // This process already open()ed the kcov device
-    auto kcov_instance = new KCOVInstance(pid);
-    kcov_instance->set_state(KCOVInstance::OPENED);
-    proc_instance->set(pid, kcov_instance);
 
+    proc.set_kcov_instance(new KCOVInstance(proc.pid()));
     return Device::open(options);
 }
 
 ErrorOr<void> KCOVDevice::ioctl(OpenFileDescription&, unsigned request, Userspace<void*> arg)
 {
-    auto thread = Thread::current();
-    auto tid = thread->tid();
-    auto pid = thread->pid();
-    auto maybe_kcov_instance = proc_instance->get(pid);
-    if (!maybe_kcov_instance.has_value())
-        return ENXIO; // This proc hasn't opened the kcov dev yet
-    auto kcov_instance = maybe_kcov_instance.value();
+    auto& proc = Process::current();
+    auto* thread = Thread::current();
+    auto* kcov_instance = proc.kcov_instance();
+    if (kcov_instance == nullptr) {
+        VERIFY(!Process::is_kcov_busy()); // No thread should be tracing at this point.
+        return ENXIO;                     // This proc hasn't opened the kcov dev yet
+    }
 
     SpinlockLocker locker(kcov_instance->spinlock());
     switch (request) {
     case KCOV_SETBUFSIZE:
-        if (kcov_instance->state() >= KCOVInstance::TRACING)
+        if (Process::is_kcov_busy())
+            // Buffer is shared among all proc threads, hence we need to check if any of them
+            // are currently tracing. If so, don't mess with the buffer.
             return EBUSY;
         return kcov_instance->buffer_allocate((FlatPtr)arg.unsafe_userspace_ptr());
     case KCOV_ENABLE:
-        if (kcov_instance->state() >= KCOVInstance::TRACING)
-            return EBUSY;
         if (!kcov_instance->has_buffer())
             return ENOBUFS;
-        VERIFY(kcov_instance->state() == KCOVInstance::OPENED);
-        kcov_instance->set_state(KCOVInstance::TRACING);
-        thread_instance->set(tid, kcov_instance);
+        thread->m_kcov_enabled = true;
         return {};
     case KCOV_DISABLE: {
-        auto maybe_kcov_instance = thread_instance->get(tid);
-        if (!maybe_kcov_instance.has_value())
-            return ENOENT;
-        VERIFY(kcov_instance->state() == KCOVInstance::TRACING);
-        kcov_instance->set_state(KCOVInstance::OPENED);
-        thread_instance->remove(tid);
+        thread->m_kcov_enabled = false;
         return {};
     }
     default:
@@ -118,10 +75,8 @@ ErrorOr<void> KCOVDevice::ioctl(OpenFileDescription&, unsigned request, Userspac
 
 ErrorOr<NonnullLockRefPtr<Memory::VMObject>> KCOVDevice::vmobject_for_mmap(Process& process, Memory::VirtualRange const&, u64&, bool)
 {
-    auto pid = process.pid();
-    auto maybe_kcov_instance = proc_instance->get(pid);
-    VERIFY(maybe_kcov_instance.has_value()); // Should happen on fd open()
-    auto kcov_instance = maybe_kcov_instance.value();
+    auto* kcov_instance = process.kcov_instance();
+    VERIFY(kcov_instance != nullptr); // Should have happened on fd open()
 
     if (!kcov_instance->vmobject())
         return ENOBUFS; // mmaped, before KCOV_SETBUFSIZE

Kernel/Devices/KCOVDevice.h

@@ -14,9 +14,6 @@ class KCOVDevice final : public BlockDevice {
     friend class DeviceManagement;
 
 public:
-    static HashMap<ProcessID, KCOVInstance*>* proc_instance;
-    static HashMap<ThreadID, KCOVInstance*>* thread_instance;
-
     static NonnullLockRefPtr<KCOVDevice> must_create();
     static void free_thread();
     static void free_process();

Kernel/Devices/KCOVInstance.h

@@ -15,15 +15,18 @@ namespace Kernel {
 #define KCOV_MAX_ENTRIES (10 * 1024 * 1024)
 
 /*
- * One KCOVInstance is allocated per process, when the process opens /dev/kcov
- * for the first time. At this point it is in state OPENED. When a thread in
- * the same process then uses the KCOV_ENABLE ioctl on the block device, the
- * instance enters state TRACING.
- *
- * A KCOVInstance in state TRACING can return to state OPENED by either the
- * KCOV_DISABLE ioctl or by killing the thread. A KCOVInstance in state OPENED
- * can return to state UNUSED only when the process dies. At this point
- * KCOVDevice::free_process will delete the KCOVInstance.
+ * 1. When a thread opens /dev/kcov for the first time, a KCOVInstance is
+ * allocated and tracked via an OwnPtr on the Kernel::Process object.
+ * 2. When a thread in the same process then uses the KCOV_SETBUFSIZE ioctl
+ * on the block device, a Memory::Region is allocated and tracked via an
+ * OwnPtr on the KCOVInstance.
+ * 3. When a thread in the same process then uses the KCOV_ENABLE ioctl on
+ * the block device, a flag is set in the Thread object and __sanitizer_cov_trace_pc
+ * will start recording this threads visited code paths .
+ * 3. When the same thread then uses the KCOV_DISABLE ioctl on the block device,
+ * a flag is unset in the Thread object and __sanitizer_cov_trace_pc will
+ * no longer record this threads visited code paths.
+ * 4. When the Process dies, the KCOVInstance and Memory::Region are GCed.
  */
 class KCOVInstance final {
 public:
@@ -33,15 +36,6 @@ public:
     bool has_buffer() const { return m_buffer != nullptr; }
     void buffer_add_pc(u64 pc);
 
-    enum State {
-        UNUSED = 0,
-        OPENED = 1,
-        TRACING = 2,
-    } m_state { UNUSED };
-
-    State state() const { return m_state; }
-    void set_state(State state) { m_state = state; }
-
     Memory::VMObject* vmobject() { return m_vmobject; }
 
     Spinlock<LockRank::None>& spinlock() { return m_lock; }

Kernel/SanCov.cpp

@@ -4,34 +4,34 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <Kernel/Arch/Processor.h>
 #include <Kernel/Devices/KCOVDevice.h>
-#include <Kernel/Tasks/Thread.h>
 
 extern bool g_in_early_boot;
 
-extern "C" {
 
-void __sanitizer_cov_trace_pc(void);
-void __sanitizer_cov_trace_pc(void)
+// Set ENABLE_KERNEL_COVERAGE_COLLECTION=ON via cmake, to inject this function on every program edge.
+// Note: This function is only used by fuzzing builds. When in use, it becomes an ultra hot code path.
+// See https://clang.llvm.org/docs/SanitizerCoverage.html#edge-coverage
+extern "C" void __sanitizer_cov_trace_pc(void);
+extern "C" void __sanitizer_cov_trace_pc(void)
 {
     if (g_in_early_boot) [[unlikely]]
         return;
 
+    auto* thread = Processor::current_thread();
+    auto* kcov_instance = thread->process().kcov_instance();
+    if (kcov_instance == nullptr || !thread->m_kcov_enabled) [[likely]]
+        return; // KCOV device hasn't been opened yet or thread is not traced
+
     if (Processor::current_in_irq()) [[unlikely]] {
-        // Do not trace in interrupts.
+        // Do not collect coverage caused by interrupts. We want the collected coverage to be a function
+        // of the syscalls executed by the fuzzer. Interrupts can occur more or less randomly. Fuzzers
+        // uses coverage to identify function call sequences, which triggered new code paths. If the
+        // coverage is noisy, the fuzzer will waste time on unintersting sequences.
         return;
     }
 
-    auto const* thread = Thread::current();
-    auto tid = thread->tid();
-    auto maybe_kcov_instance = KCOVDevice::thread_instance->get(tid);
-    if (!maybe_kcov_instance.has_value()) [[likely]] {
-        // not traced
-        return;
-    }
-    auto* kcov_instance = maybe_kcov_instance.value();
-    if (kcov_instance->state() < KCOVInstance::TRACING) [[likely]]
-        return;
     kcov_instance->buffer_add_pc((u64)__builtin_return_address(0));
-}
+    return;
 }

Kernel/Tasks/Process.cpp

@@ -9,23 +9,18 @@
 #include <AK/StdLibExtras.h>
 #include <AK/Time.h>
 #include <AK/Types.h>
-#include <Kernel/API/Syscall.h>
-#include <Kernel/Debug.h>
-#include <Kernel/Devices/DeviceManagement.h>
-#include <Kernel/Interrupts/InterruptDisabler.h>
-#include <Kernel/Security/Credentials.h>
-#include <Kernel/Tasks/Coredump.h>
-#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
-#    include <Kernel/Devices/KCOVDevice.h>
-#endif
 #include <Kernel/API/POSIX/errno.h>
 #include <Kernel/API/POSIX/sys/limits.h>
+#include <Kernel/API/Syscall.h>
 #include <Kernel/Arch/PageDirectory.h>
+#include <Kernel/Debug.h>
+#include <Kernel/Devices/DeviceManagement.h>
 #include <Kernel/Devices/Generic/NullDevice.h>
 #include <Kernel/Devices/TTY/TTY.h>
 #include <Kernel/FileSystem/Custody.h>
 #include <Kernel/FileSystem/OpenFileDescription.h>
 #include <Kernel/FileSystem/VirtualFileSystem.h>
+#include <Kernel/Interrupts/InterruptDisabler.h>
 #include <Kernel/KSyms.h>
 #include <Kernel/Library/KBufferBuilder.h>
 #include <Kernel/Library/Panic.h>
@@ -33,6 +28,8 @@
 #include <Kernel/Memory/AnonymousVMObject.h>
 #include <Kernel/Memory/SharedInodeVMObject.h>
 #include <Kernel/Sections.h>
+#include <Kernel/Security/Credentials.h>
+#include <Kernel/Tasks/Coredump.h>
 #include <Kernel/Tasks/PerformanceEventBuffer.h>
 #include <Kernel/Tasks/PerformanceManager.h>
 #include <Kernel/Tasks/Process.h>
@@ -561,6 +558,21 @@ void Process::crash(int signal, Optional<RegisterState const&> regs, bool out_of
     VERIFY_NOT_REACHED();
 }
 
+#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
+bool Process::is_kcov_busy()
+{
+    bool is_busy = false;
+    Kernel::Process::current().for_each_thread([&](auto& thread) {
+        if (thread.m_kcov_enabled) {
+            is_busy = true;
+            return IterationDecision::Break;
+        }
+        return IterationDecision::Continue;
+    });
+    return is_busy;
+}
+#endif
+
 RefPtr<Process> Process::from_pid_in_same_jail(ProcessID pid)
 {
     return Process::current().m_jail_process_list.with([&](auto const& list_ptr) -> RefPtr<Process> {
@@ -964,9 +976,6 @@ void Process::die()
     });
 
     kill_all_threads();
-#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
-    KCOVDevice::free_process();
-#endif
 }
 
 void Process::terminate_due_to_signal(u8 signal)

Kernel/Tasks/Process.h

@@ -18,6 +18,9 @@
 #include <Kernel/API/POSIX/select.h>
 #include <Kernel/API/POSIX/sys/resource.h>
 #include <Kernel/API/Syscall.h>
+#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
+#    include <Kernel/Devices/KCOVInstance.h>
+#endif
 #include <Kernel/FileSystem/InodeMetadata.h>
 #include <Kernel/FileSystem/OpenFileDescription.h>
 #include <Kernel/FileSystem/UnveilNode.h>
@@ -222,7 +225,19 @@ public:
     bool is_profiling() const { return m_profiling; }
     void set_profiling(bool profiling) { m_profiling = profiling; }
 
-    bool should_generate_coredump() const { return m_should_generate_coredump; }
+#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
+    NO_SANITIZE_COVERAGE KCOVInstance* kcov_instance()
+    {
+        return m_kcov_instance;
+    }
+    void set_kcov_instance(KCOVInstance* kcov_instance) { m_kcov_instance = kcov_instance; }
+    static bool is_kcov_busy();
+#endif
+
+    bool should_generate_coredump() const
+    {
+        return m_should_generate_coredump;
+    }
     void set_should_generate_coredump(bool b) { m_should_generate_coredump = b; }
 
     bool is_dying() const { return m_state.load(AK::MemoryOrder::memory_order_acquire) != State::Running; }
@@ -915,6 +930,10 @@ private:
     Atomic<bool, AK::MemoryOrder::memory_order_relaxed> m_is_stopped { false };
     bool m_should_generate_coredump { false };
 
+#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
+    KCOVInstance* m_kcov_instance { nullptr };
+#endif
+
     SpinlockProtected<RefPtr<Custody>, LockRank::None> m_executable;
 
     SpinlockProtected<RefPtr<Custody>, LockRank::None> m_current_directory;

Kernel/Tasks/Thread.cpp

@@ -14,7 +14,6 @@
 #include <Kernel/Arch/SmapDisabler.h>
 #include <Kernel/Arch/TrapFrame.h>
 #include <Kernel/Debug.h>
-#include <Kernel/Devices/KCOVDevice.h>
 #include <Kernel/FileSystem/OpenFileDescription.h>
 #include <Kernel/Interrupts/InterruptDisabler.h>
 #include <Kernel/KSyms.h>
@@ -441,9 +440,6 @@ void Thread::exit(void* exit_value)
             space->deallocate_region(*region);
         });
     }
-#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
-    KCOVDevice::free_thread();
-#endif
     die_if_needed();
 }
 

Kernel/Tasks/Thread.h

@@ -1259,6 +1259,11 @@ public:
     using GlobalList = IntrusiveList<&Thread::m_global_thread_list_node>;
 
     static SpinlockProtected<GlobalList, LockRank::None>& all_instances();
+
+#ifdef ENABLE_KERNEL_COVERAGE_COLLECTION
+    // Used by __sanitizer_cov_trace_pc to identify traced threads.
+    bool m_kcov_enabled { false };
+#endif
 };
 
 AK_ENUM_BITWISE_OPERATORS(Thread::FileBlocker::BlockFlags);