/*
 * PROJECT: Alcyone System Kernel
 * LICENSE: BSD Clause 3
 * PURPOSE: Thread Aware Spin Lock 
 * NT KERNEL: 5.11.9360
 * COPYRIGHT:  2023-2029 Dibymartanda Samanta <>
 */ 

 /*
  Thread Aware Spinlock, designed keeping following things in mind:
  *
  *Uses a ticket system (Ticket and Serving counters) to ensure first-come, first-served order
  *Tracks the current owner of the lock (OwnerThread), ensuring easier debugging of deadlock
  *Reentrant capability: A thread can "acquire" the lock multiple times if it already owns it, preventing self deadlocking
  *DeadLock Prevention: Detects and prevents attempts by a thread to re-acquire a lock it already holds, prevent non owning thread from releasing a lock.
  *
  */

  class KeThreadAwareSpinLockImpl
{
public:
    volatile LONG Ticket;
    volatile LONG Serving;
    PKTHREAD OwnerThread;
    KPRIORITY OriginalPriority;
    BOOLEAN IsBoosted;

    KeThreadAwareSpinLockImpl() 
        : Ticket(0), Serving(0), OwnerThread(nullptr), 
          OriginalPriority(0), IsBoosted(FALSE) {}

    void Initialize()
    {
        Ticket = 0;
        Serving = 0;
        OwnerThread = nullptr;
        OriginalPriority = 0;
        IsBoosted = FALSE;
    }

    void Acquire()
    {
        LONG CurrentTicket = {0};
        PKTHREAD CurrentThread = KeGetCurrentThread();

        if (OwnerThread == CurrentThread)
        {
            /* Already owned return the thread */
            return; 
        }

        CurrentTicket = InterlockedIncrement(&Ticket) - 1;

        while (CurrentTicket != Serving)
        {
            YieldProcessor();
        }

        OwnerThread = CurrentThread;
        BoostPriority();
    }

    void Release()
    {
        PKTHREAD CurrentThread = KeGetCurrentThread();

        if (OwnerThread != CurrentThread)
        {
            ASSERT(FALSE);
            return;
        }

        RestorePriority();
        OwnerThread = nullptr;
        InterlockedIncrement(&Serving);
    }

private:
    void BoostPriority()
{
    if (!IsBoosted)
    {
        OriginalPriority = KeGetPriorityThread(OwnerThread);
        
        /*Boost to the highest priority below real-time range*/
        KPRIORITY BoostPriority = HIGH_PRIORITY;
        
        /* Ensure we don't lower the priority if it's already higher*/
        if (OriginalPriority < BoostPriority)
        {
            KeSetPriorityThread(OwnerThread, BoostPriority);
            IsBoosted = TRUE;
        }
    }

    void RestorePriority()
    {
        if (IsBoosted)
        {
            KeSetPriorityThread(OwnerThread, OriginalPriority);
            IsBoosted = FALSE;
        }
    }
};

typedef KeThreadAwareSpinLockImpl* PKETHREADAWARESPINLOCK;

/* Exported Function */ 
VOID
NTAPI
KeThreadAwareSpinLockInitialize(
    _Out_ PKETHREADAWARESPINLOCK SpinLock)
{
    if (SpinLock)
    {
        SpinLock->Initialize();
    }
}

VOID
NTAPI
KeThreadAwareSpinLockAcquire(
    _Inout_ PKETHREADAWARESPINLOCK SpinLock)
{
    if (SpinLock)
    {
        SpinLock->Acquire();
    }
}

VOID
NTAPI
KeThreadAwareSpinLockRelease(
    _Inout_ PKETHREADAWARESPINLOCK SpinLock)
{
    if (SpinLock)
    {
        SpinLock->Release();
    }
}