greenies.c

// Compile with clang or MSVC (WINDOWS ONLY RN)
//
// Implementing a POC green threads system using safepoints to show how cheap and simple it can
// be done, all you need to do is call SAFEPOINT_POLL in your own language at the top of every
// loop and function body (you can loosen up on this depending on the latency of pausing you're
// willing to pay). Safepoint polling is made cheap because it's a load without a use site
// which means it doesn't introduce a stall and pays a sub-cycle cost because of it (wastes resources
// sure but doesn't block up the rest of execution).
//
//  # safepoint poll
//  test al, byte [poll_site]
//
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>
#include <process.h>
#include <stdbool.h>

#define WIN32_LEAN_AND_MEAN
#include <windows.h>

#pragma comment(lib, "user32.lib")

static int green_count;
static CONTEXT greens[5];
static int curr;

#ifdef __clang__
static _Alignas(4096) volatile char poll_site[4096];
#define SAFEPOINT_POLL() asm(".align 1\n\ttest %%al, poll_site(%%rip)" :::)
#else
static __declspec(align(4096)) volatile char poll_site[4096];
#define SAFEPOINT_POLL() (poll_site[0])
#endif

static LONG its_so_over(EXCEPTION_POINTERS* e) {
    // read from PAUSE_ADDR means we hit a safepoint during a pause
    if (e->ExceptionRecord->ExceptionCode == EXCEPTION_GUARD_PAGE &&
        e->ExceptionRecord->ExceptionInformation[1] == (uintptr_t) poll_site) {
        // save context
        greens[curr] = *e->ContextRecord;

        // run next green thread
        curr += 1;
        if (curr >= green_count) {
            curr = 0;
        }

        *e->ContextRecord = greens[curr];
        return EXCEPTION_CONTINUE_EXECUTION;
    } else {
        return EXCEPTION_CONTINUE_SEARCH;
    }
}

static CONTEXT create_green(CONTEXT* base, void fn(void*), void* param, size_t stack_size) {
    if (stack_size == 0) {
        stack_size = 32*1024; // default green stack is 32KiB
    }

    void* stack = VirtualAlloc(NULL, stack_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    CONTEXT c = *base;
    c.Rip = (uint64_t) fn;
    c.Rsp = ((uint64_t) stack) + (stack_size - 40);
    c.Rcx = (uint64_t) param;
    return c;
}

static void printer(void* param);

static void worker(void* param) {
    CONTEXT base;
    RtlCaptureContext(&base);

    green_count = 5;
    greens[0] = create_green(&base, printer, "Jeepers!", 0);
    greens[1] = create_green(&base, printer, "Jinkies!", 0);
    greens[2] = create_green(&base, printer, "Ruh-roh!", 0);
    greens[3] = create_green(&base, printer, "Zoinks!",  0);
    greens[4] = create_green(&base, printer, "Fuck!",    0);

    // hook guard handler
    AddVectoredExceptionHandler(1, its_so_over);

    // jump into first green thread
    RtlRestoreContext(&greens[0], NULL);
}

static void printer(void* param) {
    for (;;) {
        SAFEPOINT_POLL();
        printf("%s\n", (char*) param);
    }
}

int main() {
    int thread = _beginthread(worker, 8*1024, NULL);
    for (;;) {
        // set pause state
        DWORD old;
        if (!VirtualProtect((void*) poll_site, 4096, PAGE_GUARD | PAGE_READONLY, &old)) {
            fprintf(stderr, "error: could not reset guard page!\n");
            abort();
        }

        Sleep(10);
    }
    return 0;
}