nt-lite/code/kernel/ke_queue.c

// Kernel message queue primitive. This file owns bounded queue storage and queue timeout behavior.

#include <stdint.h>
#include <string.h>
#include <kernel/ke_queue.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_irql.h>
#include <kernel/ob_object.h>

#define QUEUE_TIMEOUT_INFINITE UINT64_MAX

// Destroys a queue and releases its private state events.
static void KiQueueDestroy(OBJECT* obj) {
    QUEUE* queue = (QUEUE*)obj;
    if (!queue)
        return;
    if (queue->NotEmptyEvent) {
        ObDereferenceObject(&queue->NotEmptyEvent->Header);
        queue->NotEmptyEvent = (EVENT*)0;
    }
    if (queue->NotFullEvent) {
        ObDereferenceObject(&queue->NotFullEvent->Header);
        queue->NotFullEvent = (EVENT*)0;
    }
}

// Updates waitable state after one message has been inserted.
static void KiQueueNoteInsertedLocked(QUEUE* queue, uint32_t previousCount) {
    if (!queue || !queue->NotEmptyEvent || !queue->NotFullEvent)
        return;

    if (previousCount == 0)
        KeSignalEvent(queue->NotEmptyEvent);
    if (queue->Count == QUEUE_DEPTH)
        KeClearEvent(queue->NotFullEvent);
}

// Updates waitable state after one message has been removed.
static void KiQueueNoteRemovedLocked(QUEUE* queue, uint32_t previousCount) {
    if (!queue || !queue->NotEmptyEvent || !queue->NotFullEvent)
        return;

    if (previousCount == QUEUE_DEPTH)
        KeSignalEvent(queue->NotFullEvent);
    if (queue->Count == 0)
        KeClearEvent(queue->NotEmptyEvent);
}

// Returns the scheduler tick value used by queue finite-timeout accounting.
static uint64_t KiQueueReadSchedulerTick(void) {
    KE_SCHED_STATS stats;
    KeGetSchedulerStats(&stats);
    return stats.Ticks;
}

// Converts a millisecond timeout to at least one scheduler tick.
static uint64_t KiQueueTimeoutTicks(uint64_t timeoutMilliseconds) {
    uint64_t ticks = timeoutMilliseconds / SCHED_TICK_MS;
    if ((timeoutMilliseconds % SCHED_TICK_MS) != 0)
        ticks++;
    return ticks ? ticks : 1;
}

// Computes an absolute scheduler-tick deadline for a finite timeout.
static uint64_t KiQueueTimeoutDeadline(uint64_t timeoutMilliseconds) {
    uint64_t now = KiQueueReadSchedulerTick();
    uint64_t ticks = KiQueueTimeoutTicks(timeoutMilliseconds);
    uint64_t deadline = now + ticks;
    return deadline < now ? UINT64_MAX : deadline;
}

// Computes the remaining wait time for a finite timeout deadline.
static uint64_t KiQueueRemainingMilliseconds(uint64_t deadlineTick) {
    uint64_t now = KiQueueReadSchedulerTick();
    if (deadlineTick <= now)
        return 0;

    uint64_t ticks = deadlineTick - now;
    if (ticks > UINT64_MAX / SCHED_TICK_MS)
        return UINT64_MAX;
    return ticks * SCHED_TICK_MS;
}

// Waits for a queue state event using either an infinite or finite timeout.
static KSTATUS KiQueueWaitForEvent(QUEUE* queue, EVENT* event, uint64_t timeoutMilliseconds, uint64_t deadlineTick) {
    if (!queue || !event)
        return STATUS_INVALID_PARAMETER;

    uint64_t waitMilliseconds = timeoutMilliseconds;
    if (timeoutMilliseconds != QUEUE_TIMEOUT_INFINITE) {
        waitMilliseconds = KiQueueRemainingMilliseconds(deadlineTick);
        if (waitMilliseconds == 0)
            return STATUS_TIMEOUT;
    }

    EVENT* events[1] = { event };
    return KeWaitForMultipleEvents(events, 1, 0, waitMilliseconds, (uint32_t*)0);
}

// Creates a fixed-depth queue with waitable occupancy and capacity events.
KSTATUS KeQueueCreate(QUEUE **out) {
    if (!out) return STATUS_INVALID_PARAMETER;
    *out = (QUEUE*)0;

    QUEUE* q = (QUEUE *)ObAllocateObject(ObjectTypeQueue, sizeof(QUEUE));
    if (!q)  return STATUS_NO_MEMORY;
    q->Header.Destroy = KiQueueDestroy;
    q->Lock  = (SPINLOCK)SPINLOCK_INIT;
    q->Head  = q->Tail = q->Count = 0;

    KSTATUS s = KeEventCreate(EventManualReset, 0, &q->NotEmptyEvent);
    if (!KSUCCESS(s)) {
        ObDereferenceObject(&q->Header);
        return s;
    }

    s = KeEventCreate(EventManualReset, 1, &q->NotFullEvent);
    if (!KSUCCESS(s)) {
        ObDereferenceObject(&q->Header);
        return s;
    }

    *out = q;
    return STATUS_SUCCESS;
}

// Sends one fixed-size queue message, waiting according to the timeout contract.
KSTATUS KeSendMessage(QUEUE* queue, const QUEUE_MESSAGE* msg, uint64_t timeoutMilliseconds) {
    if (!queue || !msg) return STATUS_INVALID_PARAMETER;
    if (timeoutMilliseconds != 0 && !KeCanWait()) return STATUS_INVALID_IRQL;

    uint64_t deadlineTick = 0;
    if (timeoutMilliseconds != 0 && timeoutMilliseconds != QUEUE_TIMEOUT_INFINITE)
        deadlineTick = KiQueueTimeoutDeadline(timeoutMilliseconds);

    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&queue->Lock);
        if (queue->Count < QUEUE_DEPTH) {
            uint32_t previousCount = queue->Count;
            queue->Ring[queue->Tail] = *msg;
            queue->Tail = (queue->Tail + 1) % QUEUE_DEPTH;
            queue->Count++;
            KiQueueNoteInsertedLocked(queue, previousCount);
            KeReleaseSpinlock(&queue->Lock, flags);
            return STATUS_SUCCESS;
        }
        KeClearEvent(queue->NotFullEvent);
        KeReleaseSpinlock(&queue->Lock, flags);

        if (timeoutMilliseconds == 0)
            return STATUS_WOULD_BLOCK;

        KSTATUS s = KiQueueWaitForEvent(queue, queue->NotFullEvent, timeoutMilliseconds, deadlineTick);
        if (s == STATUS_TIMEOUT || KFAILED(s))
            return s;
    }
}

// Receives one fixed-size queue message, waiting according to the timeout contract.
KSTATUS KeReceiveMessage(QUEUE* queue, QUEUE_MESSAGE* msg, uint64_t timeoutMilliseconds) {
    if (!queue || !msg) return STATUS_INVALID_PARAMETER;
    if (timeoutMilliseconds != 0 && !KeCanWait()) return STATUS_INVALID_IRQL;

    uint64_t deadlineTick = 0;
    if (timeoutMilliseconds != 0 && timeoutMilliseconds != QUEUE_TIMEOUT_INFINITE)
        deadlineTick = KiQueueTimeoutDeadline(timeoutMilliseconds);

    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&queue->Lock);
        if (queue->Count > 0) {
            uint32_t previousCount = queue->Count;
            *msg = queue->Ring[queue->Head];
            queue->Head = (queue->Head + 1) % QUEUE_DEPTH;
            queue->Count--;
            KiQueueNoteRemovedLocked(queue, previousCount);
            KeReleaseSpinlock(&queue->Lock, flags);
            return STATUS_SUCCESS;
        }
        KeClearEvent(queue->NotEmptyEvent);
        KeReleaseSpinlock(&queue->Lock, flags);

        if (timeoutMilliseconds == 0)
            return STATUS_WOULD_BLOCK;

        KSTATUS s = KiQueueWaitForEvent(queue, queue->NotEmptyEvent, timeoutMilliseconds, deadlineTick);
        if (s == STATUS_TIMEOUT || KFAILED(s))
            return s;
    }
}