From 88008ac470f9986b138252024fffbca44da2b2a5 Mon Sep 17 00:00:00 2001
From: Alex <enderice2@users.noreply.github.com>
Date: Mon, 6 Feb 2023 19:32:20 +0200
Subject: [PATCH] Added timeout lock

---
 Core/Lock.cpp    | 54 +++++++++++++++++++++++++++++++++++++++++++++---
 include/lock.hpp | 30 ++++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/Core/Lock.cpp b/Core/Lock.cpp
index e485a74..88c7784 100644
--- a/Core/Lock.cpp
+++ b/Core/Lock.cpp
@@ -2,6 +2,7 @@
 
 #include <debug.h>
 #include <smp.hpp>
+#include <atomic.hpp>
 
 #include "../kernel.h"
 
@@ -22,9 +23,7 @@ void LockClass::DeadLock(SpinLockData Lock)
 
     if (Config.UnlockDeadLock && this->DeadLocks > 10)
     {
-        warn("Unlocking lock '%s' held by '%s'! %ld locks in queue. Core %ld is being held by %ld.",
-             Lock.AttemptingToGet, Lock.CurrentHolder,
-             Lock.Count, CCore, Lock.Core);
+        warn("Unlocking lock '%s' to prevent deadlock. (this is enabled in the kernel config)", Lock.AttemptingToGet);
         this->DeadLocks = 0;
         this->Unlock();
     }
@@ -85,3 +84,52 @@ int LockClass::Unlock()
 
     return 0;
 }
+
+void LockClass::TimeoutDeadLock(SpinLockData Lock, uint64_t Timeout)
+{
+    CPUData *CoreData = GetCurrentCPU();
+    long CCore = 0xdead;
+    if (CoreData != nullptr)
+        CCore = CoreData->ID;
+    warn("Potential deadlock in lock '%s' held by '%s'! %ld locks in queue. Core %ld is being held by %ld. Timeout in %ld",
+         Lock.AttemptingToGet, Lock.CurrentHolder,
+         Lock.Count, CCore, Lock.Core,
+         Timeout);
+
+    // TODO: Print on screen too.
+
+    uint64_t Counter = TimeManager->GetCounter();
+    if (Timeout < Counter)
+    {
+        warn("Unlocking lock '%s' because of timeout. (%ld < %ld)", Lock.AttemptingToGet, Timeout, Counter);
+        this->Unlock();
+    }
+
+    if (TaskManager)
+        TaskManager->Schedule();
+}
+
+int LockClass::TimeoutLock(const char *FunctionName, uint64_t Timeout)
+{
+    LockData.AttemptingToGet = FunctionName;
+    Atomic<uint64_t> Target = 0;
+Retry:
+    unsigned int i = 0;
+    while (__atomic_exchange_n(&IsLocked, true, __ATOMIC_ACQUIRE) && ++i < 0x10000000)
+        CPU::Pause();
+    if (i >= 0x10000000)
+    {
+        if (Target.Load() == 0)
+            Target = TimeManager->CalculateTarget(Timeout);
+        TimeoutDeadLock(LockData, Target.Load());
+        goto Retry;
+    }
+    LockData.Count++;
+    LockData.CurrentHolder = FunctionName;
+    CPUData *CoreData = GetCurrentCPU();
+    if (CoreData != nullptr)
+        LockData.Core = CoreData->ID;
+    __sync_synchronize();
+
+    return 0;
+}
diff --git a/include/lock.hpp b/include/lock.hpp
index ed23deb..6a36297 100644
--- a/include/lock.hpp
+++ b/include/lock.hpp
@@ -5,6 +5,7 @@
 #include <cpu.hpp>
 
 #ifdef __cplusplus
+
 /** @brief Please use this macro to create a new lock. */
 class LockClass
 {
@@ -16,18 +17,23 @@ class LockClass
         size_t Count = 0;
         long Core = 0;
     };
-    void DeadLock(SpinLockData Lock);
 
 private:
     SpinLockData LockData;
     bool IsLocked = false;
     unsigned long DeadLocks = 0;
 
+    void DeadLock(SpinLockData Lock);
+    void TimeoutDeadLock(SpinLockData Lock, uint64_t Timeout);
+
 public:
     SpinLockData *GetLockData() { return &LockData; }
     int Lock(const char *FunctionName);
     int Unlock();
+
+    int TimeoutLock(const char *FunctionName, uint64_t Timeout);
 };
+
 /** @brief Please use this macro to create a new smart lock. */
 class SmartLockClass
 {
@@ -42,6 +48,21 @@ public:
     }
     ~SmartLockClass() { this->LockPointer->Unlock(); }
 };
+
+class SmartTimeoutLockClass
+{
+private:
+    LockClass *LockPointer = nullptr;
+
+public:
+    SmartTimeoutLockClass(LockClass &Lock, const char *FunctionName, uint64_t Timeout)
+    {
+        this->LockPointer = &Lock;
+        this->LockPointer->TimeoutLock(FunctionName, Timeout);
+    }
+    ~SmartTimeoutLockClass() { this->LockPointer->Unlock(); }
+};
+
 /** @brief Please use this macro to create a new smart critical section lock. */
 class SmartLockCriticalSectionClass
 {
@@ -65,6 +86,7 @@ public:
             CPU::Interrupts(CPU::Enable);
     }
 };
+
 /** @brief Please use this macro to create a new critical section. */
 class SmartCriticalSectionClass
 {
@@ -89,10 +111,16 @@ public:
 
 /** @brief Create a new lock (can be used with SmartCriticalSection). */
 #define NewLock(Name) LockClass Name
+
 /** @brief Simple lock that is automatically released when the scope ends. */
 #define SmartLock(LockClassName) SmartLockClass CONCAT(lock##_, __COUNTER__)(LockClassName, __FUNCTION__)
+
+/** @brief Simple lock with timeout that is automatically released when the scope ends. */
+#define SmartTimeoutLock(LockClassName, Timeout) SmartTimeoutLockClass CONCAT(lock##_, __COUNTER__)(LockClassName, __FUNCTION__, Timeout)
+
 /** @brief Simple critical section that is automatically released when the scope ends and interrupts are restored if they were enabled. */
 #define SmartCriticalSection(LockClassName) SmartLockCriticalSectionClass CONCAT(lock##_, __COUNTER__)(LockClassName, __FUNCTION__)
+
 /** @brief Automatically disable interrupts and restore them when the scope ends. */
 #define CriticalSection SmartCriticalSectionClass