nt-lite/code/kernel/ke_timer.c

// Kernel waitable timer objects. This file owns timer deadlines and timer wait signaling.

#include <stdint.h>
#include <stdbool.h>
#include <kernel/ke_list.h>
#include <string.h>
#include <kernel/ke_timer.h>
#include <kernel/ke_time.h>
#include <kernel/ob_object.h>

#define TIMER_SIGNAL_BATCH 32u

static LIST_ENTRY gTimerList;
static SPINLOCK gTimerLock;
static bool gTimerInitialized;

// Initializes the global timer list on first use.
static void KiTimerEnsureInitialized(void) {
    if (gTimerInitialized)
        return;

    gTimerLock = (SPINLOCK)SPINLOCK_INIT;
    LIST_INIT(&gTimerList);
    gTimerInitialized = 1;
}

// Converts a millisecond interval to microseconds with saturation.
static uint64_t KiTimerMillisecondsToMicroseconds(uint64_t milliseconds) {
    if (milliseconds > UINT64_MAX / 1000ull)
        return UINT64_MAX;
    return milliseconds * 1000ull;
}

// Adds a monotonic interval to the current time with saturation.
static uint64_t KiTimerDeadlineFromNow(uint64_t intervalMicroseconds) {
    uint64_t now = KeQueryMonotonicTicks();
    if (intervalMicroseconds > UINT64_MAX - now)
        return UINT64_MAX;
    return now + intervalMicroseconds;
}

// Removes an active timer from the global deadline list while the timer lock is held.
static uint8_t KiTimerRemoveLocked(KTIMER* timer) {
    if (!timer || !timer->Linked)
        return 0;

    KeListRemoveChecked(&timer->Link);
    timer->Linked = 0;
    timer->Active = 0;
    timer->DueMicroseconds = 0;
    timer->PeriodMicroseconds = 0;
    return 1;
}

// Implements the timer Destroy operation.
static void KiTimerDestroy(OBJECT* obj) {
    KTIMER* timer = (KTIMER*)obj;
    if (!timer)
        return;

    KiTimerEnsureInitialized();
    uintptr_t flags = KeAcquireSpinlock(&gTimerLock);
    (void)KiTimerRemoveLocked(timer);
    KeReleaseSpinlock(&gTimerLock, flags);

    if (timer->Event) {
        ObDereferenceObject(&timer->Event->Header);
        timer->Event = (EVENT*)0;
    }
}

// Creates an unset waitable timer object.
KSTATUS KeTimerCreate(KTIMER** out) {
    if (!out)
        return NTLSTATUS_INVALID_PARAMETER;
    *out = (KTIMER*)0;

    KiTimerEnsureInitialized();

    KTIMER* timer = (KTIMER*)ObAllocateObject(ObjectTypeTimer, sizeof(KTIMER));
    if (!timer)
        return NTLSTATUS_NO_MEMORY;

    timer->Header.Destroy = KiTimerDestroy;
    LIST_INIT(&timer->Link);

    KSTATUS s = KeEventCreate(EventAutoReset, 0, &timer->Event);
    if (!KSUCCESS(s)) {
        ObDereferenceObject(&timer->Header);
        return s;
    }

    *out = timer;
    return NTLSTATUS_SUCCESS;
}

// Arms or rearms a timer with a relative due time and optional period.
KSTATUS KeTimerSet(KTIMER* timer, uint64_t dueTimeMilliseconds, uint64_t periodMilliseconds) {
    if (!timer || !timer->Event)
        return NTLSTATUS_INVALID_PARAMETER;

    KiTimerEnsureInitialized();

    uint64_t dueMicroseconds = KiTimerMillisecondsToMicroseconds(dueTimeMilliseconds);
    uint64_t periodMicroseconds = KiTimerMillisecondsToMicroseconds(periodMilliseconds);
    uint8_t dropListReference = 0;
    uint8_t addListReference = 0;
    uint8_t signalNow = dueMicroseconds == 0;
    EVENT* event = timer->Event;
    ObReferenceObject(&event->Header);

    uintptr_t flags = KeAcquireSpinlock(&gTimerLock);

    if (timer->Linked) {
        if (signalNow && periodMicroseconds == 0) {
            dropListReference = KiTimerRemoveLocked(timer);
        } else {
            timer->DueMicroseconds = KiTimerDeadlineFromNow(signalNow ? periodMicroseconds : dueMicroseconds);
            timer->PeriodMicroseconds = periodMicroseconds;
            timer->Active = 1;
        }
    } else if (!signalNow || periodMicroseconds != 0) {
        timer->DueMicroseconds = KiTimerDeadlineFromNow(signalNow ? periodMicroseconds : dueMicroseconds);
        timer->PeriodMicroseconds = periodMicroseconds;
        timer->Active = 1;
        timer->Linked = 1;
        KeListInsertTailChecked(&gTimerList, &timer->Link);
        addListReference = 1;
    }

    KeReleaseSpinlock(&gTimerLock, flags);

    if (addListReference)
        ObReferenceObject(&timer->Header);
    if (dropListReference)
        ObDereferenceObject(&timer->Header);

    KeClearEvent(event);
    if (signalNow)
        KeSignalEvent(event);
    ObDereferenceObject(&event->Header);
    return NTLSTATUS_SUCCESS;
}

// Cancels a timer and clears any retained signal.
KSTATUS KeTimerCancel(KTIMER* timer) {
    if (!timer || !timer->Event)
        return NTLSTATUS_INVALID_PARAMETER;

    KiTimerEnsureInitialized();

    uintptr_t flags = KeAcquireSpinlock(&gTimerLock);
    uint8_t hadListReference = KiTimerRemoveLocked(timer);
    KeReleaseSpinlock(&gTimerLock, flags);

    if (hadListReference)
        ObDereferenceObject(&timer->Header);
    KeClearEvent(timer->Event);
    return NTLSTATUS_SUCCESS;
}

// Returns the event that backs timer waits.
EVENT* KeTimerWaitEvent(KTIMER* timer) {
    if (!timer)
        return (EVENT*)0;
    return timer->Event;
}

// Services expired timers from the scheduler tick path.
void KeTimerTick(void) {
    KiTimerEnsureInitialized();

    EVENT* events[TIMER_SIGNAL_BATCH];
    KTIMER* listReferencesToDrop[TIMER_SIGNAL_BATCH];
    uint32_t count = 0;
    uint32_t dropCount = 0;
    uint64_t now = KeQueryMonotonicTicks();

    uintptr_t flags = KeAcquireSpinlock(&gTimerLock);
    LIST_ENTRY* node = gTimerList.Fwd;
    while (node != &gTimerList && count < TIMER_SIGNAL_BATCH) {
        LIST_ENTRY* next = node->Fwd;
        KTIMER* timer = CONTAINING_RECORD(node, KTIMER, Link);

        if (timer->Active && timer->DueMicroseconds <= now) {
            if (timer->Event) {
                ObReferenceObject(&timer->Event->Header);
                events[count++] = timer->Event;
            }

            if (timer->PeriodMicroseconds != 0) {
                uint64_t due = now + timer->PeriodMicroseconds;
                timer->DueMicroseconds = due < now ? UINT64_MAX : due;
            } else {
                if (KiTimerRemoveLocked(timer))
                    listReferencesToDrop[dropCount++] = timer;
            }
        }

        node = next;
    }
    KeReleaseSpinlock(&gTimerLock, flags);

    for (uint32_t i = 0; i < dropCount; i++)
        ObDereferenceObject(&listReferencesToDrop[i]->Header);

    for (uint32_t i = 0; i < count; i++) {
        KeSignalEvent(events[i]);
        ObDereferenceObject(&events[i]->Header);
    }
}