test.c

#include <stdlib.h>
#include <stdio.h>
#include <signal.h>
#include <pthread.h>
#include <signal.h>
#include <stdint.h>
#include <time.h>

#include "queue.h"
#include "erl_nif.h"

#define min(a, b) (((a) < (b)) ? (a) : (b))

#define STACK_SIZE 4000000
#define INT_SIGNAL 62

void add_nif_timer(pthread_t thread);
void rem_nif_timer(pthread_t thread);

// We store our registers in this type when switching ctx.
typedef void *exec_ctx[8];

typedef struct {
    // Pointer to our pocket universe
    void *alloc_stack_ptr;
    int alloc_stack_size;
    // This always stores data related to the exec stack.
    exec_ctx ctx;
    // This always stores data related to the scheduler stack.
    exec_ctx return_ctx;
    // Debug counter that stores the amount of reschedules we have experienced.
    int reschedules;
} Exec_state;

// We store the data we need in order to get back into our original
// stack in a thread local.
_Thread_local Exec_state *exec_state = 0;

// This is what is sent to the exec stack when we first jump into it.
// It contains all the usual arguements that you would expect to get
// when a NIF first runs.
// Passed on the the users NIF function.
typedef struct {
    ErlNifEnv *env;
    int argc;
    const ERL_NIF_TERM *argv;
} NifArgs;

// Used to signal what to do when we return from the nif execution stack
// back to the scheduler stack.
// The options are RESCHEDULE if execution has not yet finished, or RETURN
// if we want to return a term. If the enum is RETURN, we also need to
// provide a return_term.
enum ReturnType {RESCHEDULE, RETURN};
typedef struct {
    enum ReturnType type;
    ERL_NIF_TERM return_term;
} ReturnStruct;

// The resource type we use to store our stack :)
ErlNifResourceType *INCOMPLETE_EXEC_ENV;

// List types for the global active timer list.
struct nif_timers_entry {
    struct timespec start;
    pthread_t thread;
    LIST_ENTRY(nif_timers_entry) pointers;
};
typedef struct {
    LIST_HEAD(nif_timers_list, nif_timers_entry) list;
} Timers;

ErlNifMutex *timers_mutex;
Timers timers;

// This will switch from one stack to another, passing message over to the 
// other world.
// Inspired by the context switch code of luajit's Coco.
static inline void *ctx_switch(exec_ctx from, exec_ctx to, void *message) {
    __asm__ __volatile__ (
            "leaq 1f(%%rip), %%rax\n" // 1f refers to label 1 forwards
            "movq %%rax, (%0)\n"      // Move rip to from[0]
            "movq %%rsp, 8(%0)\n"     // Move rsp to from[1]
            "movq %%rbp, 16(%0)\n"    // from[2]
            "movq %%rbx, 24(%0)\n"    // from[3]
            "movq %%r12, 32(%0)\n"    // from[4]
            "movq %%r13, 40(%0)\n"    // from[5]
            "movq %%r14, 48(%0)\n"    // from[6]
            "movq %%r15, 56(%0)\n"    // from[7]
            "movq 56(%1), %%r15\n"    // Restore backwards
            "movq 48(%1), %%r14\n" 
            "movq 40(%1), %%r13\n"
            "movq 32(%1), %%r12\n" 
            "movq 24(%1), %%rbx\n" 
            "movq 16(%1), %%rbp\n"
            "movq 8(%1), %%rsp\n" 
            "jmpq *(%1)\n" // This is where we switch worlds..
            "1:\n"         // and this is where the worlds join back together.
            : "+S" (from), "+D" (to), 
            // We force message into the c register, this will be read by the 
            // receiver in the other execution context.
            "+c" (message) 
            :
            // Clobber registers. This prevents the compiler from using these.
            : "rax", "rcx", "rdx", "r8", "r9", "r10", "r11", "memory", "cc"
            );
    return message;
}

// Used to call a function on a new stack
static inline void exec_stack_launchpad(void) {
    void *func;
    NifArgs *message;
    // When we jump to the launchpad, this assembly runs directly
    // after ctx_switch executes jmpq.
    __asm__ __volatile__ (
            "movq %%r12, %0\n"
            "movq %%rcx, %1\n"
            : "=m" (func), "=m" (message)
            );
    ReturnStruct ret;
    ret.type = RETURN;
    
    pthread_t thread = pthread_self();
    add_nif_timer(thread);
    ret.return_term = ((ERL_NIF_TERM (*)(ErlNifEnv*, int, const ERL_NIF_TERM[]))func)
        (message->env, message->argc, message->argv);
    rem_nif_timer(thread);

    free(message);
    ctx_switch(exec_state->ctx, exec_state->return_ctx, &ret);
}

// ==========
// Timers
// ==========

//void timespec_diff(struct timespec *start, struct timespec *stop,
//        struct timespec *result) {
//    if ((stop->tv_nsec - start->tv_nsec) < 0) {
//        result->tv_sec = stop->tv_sec - start->tv_sec - 1;
//        result->tv_nsec = stop->tv_nsec - start->tv_nsec + 1000000000;
//    } else {
//        result->tv_sec = stop->tv_sec - start->tv_sec;
//        result->tv_nsec = stop->tv_nsec - start->tv_nsec;
//    }
//
//    return;
//}
void timespec_diff(struct timespec *start, struct timespec *end, struct timespec *result) {
	if ((end->tv_nsec-start->tv_nsec)<0) {
		result->tv_sec = end->tv_sec-start->tv_sec-1;
		result->tv_nsec = 1000000000+end->tv_nsec-start->tv_nsec;
	} else {
		result->tv_sec = end->tv_sec-start->tv_sec;
		result->tv_nsec = end->tv_nsec-start->tv_nsec;
	}
}

void add_nif_timer(pthread_t thread) {
    enif_mutex_lock(timers_mutex);
    struct nif_timers_entry *timer = malloc(sizeof(struct nif_timers_entry));
    clock_gettime(CLOCK_MONOTONIC_RAW, &timer->start);
    timer->thread = thread;
    LIST_INSERT_HEAD(&timers.list, timer, pointers);
    enif_mutex_unlock(timers_mutex);
}

void rem_nif_timer(pthread_t thread) {
    enif_mutex_lock(timers_mutex);
    struct nif_timers_entry *ct, *ct_temp;
    LIST_FOREACH_SAFE(ct, &timers.list, pointers, ct_temp) {
        LIST_REMOVE(ct, pointers);
        free(ct);
    }
    enif_mutex_unlock(timers_mutex);
}

// Go through the timer thread looking for expired timers.
// Interrupt any threads that have been running too long by sending
// our interrupt signal.
//
// This will cause the handle_thread_int signal handler to run,
// which will switch back to the scheduler stack and reschedule.
void send_nif_timer_signals() {
    enif_mutex_lock(timers_mutex);

    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC_RAW, &now);
    
    struct nif_timers_entry *ct, *ct_temp;
    struct timespec current;
    LIST_FOREACH_SAFE(ct, &timers.list, pointers, ct_temp) {
        timespec_diff(&ct->start, &now, &current);
        long diff_int = (current.tv_sec * 1000000000) + current.tv_nsec;
        if (diff_int > 1000000) { // 1ms
            pthread_kill(ct->thread, INT_SIGNAL);
            LIST_REMOVE(ct, pointers);
            free(ct);
        }
    }

    enif_mutex_unlock(timers_mutex);
}

// Main timer loop.
// TODO: Make sleep smarter
void *timer_thread_main() {
    while (1) {
        send_nif_timer_signals();
        usleep(500);
    }
}

// Signal handler that reschedules the nif.
// Registered in nif_init:
void handle_thread_int(int signum) {
    ReturnStruct ret;
    ret.type = RESCHEDULE;
    ctx_switch(exec_state->ctx, exec_state->return_ctx, &ret);
}

int nif_load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info) {
    // Timer data
    timers_mutex = enif_mutex_create("timers_mutex");
    LIST_INIT(&timers.list);

    // Register interrupt signal handler
    signal(INT_SIGNAL, handle_thread_int);

    // Stack resource term
    INCOMPLETE_EXEC_ENV = enif_open_resource_type(env, NULL, "incomplete_exec", NULL, ERL_NIF_RT_CREATE, NULL);

    // Go timer thread!
    pthread_t thread;
    pthread_create(&thread, NULL, timer_thread_main, NULL);

    return 0;
}

// This is called by the user nif function when he wants to stop rescheduling.
//
// NOTE: This should ALWAYS be called before you start creating terms!!
// If you start creating terms before calling this, you risk your terms getting
// garbage collected in a reschedule before you get to return them!
void stop_reschedule() {
    rem_nif_timer(pthread_self());
}
void start_reschedule() {
    add_nif_timer(pthread_self());
}

ERL_NIF_TERM inner_test(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    ErlNifBinary bin, outbin;
    unsigned char byte;
    unsigned val, i;

    if (argc != 2 || !enif_inspect_binary(env, argv[0], &bin) ||
            !enif_get_uint(env, argv[1], &val) || val > 255)
        return enif_make_badarg(env);
    if (bin.size == 0)
        return argv[0];
    byte = (unsigned char)val;
    enif_alloc_binary(bin.size, &outbin);
    for (i = 0; i < bin.size; i++)
        outbin.data[i] = bin.data[i] ^ byte;

    stop_reschedule();
    return enif_make_tuple2(env,
            enif_make_binary(env, &outbin),
            enif_make_int(env, 0));

}

static ERL_NIF_TERM nif_resume(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    enif_get_resource(env, argv[0], INCOMPLETE_EXEC_ENV, (void *)exec_state);
    exec_state->reschedules += 1;

    add_nif_timer(pthread_self());
    ReturnStruct *ret = (ReturnStruct *)ctx_switch(exec_state->return_ctx, exec_state->ctx, NULL);

    switch (ret->type) {
        case RETURN:
            {
                printf("Finished nif exec with %d reschedules\n", exec_state->reschedules);
                return ret->return_term;
            }
        case RESCHEDULE: 
            {
                return enif_schedule_nif(env, "nif_resume", 0, nif_resume, 1, argv);
            }
    }
}

static ERL_NIF_TERM nif_test(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
    exec_state = enif_alloc_resource(INCOMPLETE_EXEC_ENV, sizeof(Exec_state));
    exec_state->alloc_stack_ptr = malloc(STACK_SIZE);
    exec_state->alloc_stack_size = STACK_SIZE;
    exec_state->reschedules = 0;

    // Find the start of our stack, the uppermost address
    size_t *stack_start = (size_t *)(exec_state->alloc_stack_ptr + STACK_SIZE);
    // Set the bottom of the stack to a value easy to spot in a debugger
    stack_start[-1] = 0xdeaddeaddeaddead;

    // Initialize the context we are going to jump to in a second
    exec_state->ctx[0] = (void *)(exec_stack_launchpad); // PC address
    exec_state->ctx[1] = (void *)(&stack_start[-1]);     // SP address
    exec_state->ctx[2] = (void *)0;
    exec_state->ctx[3] = (void *)0;
    exec_state->ctx[4] = (void *)(inner_test);           // Argument for the launchpad
    exec_state->ctx[5] = (void *)0;
    exec_state->ctx[6] = (void *)0;
    exec_state->ctx[7] = (void *)0;

    // Nif args for the launchpad
    NifArgs *args = malloc(sizeof(NifArgs));
    args->env = env;
    args->argc = argc;
    args->argv = argv;

    // Go!
    ReturnStruct *ret = (ReturnStruct *)ctx_switch(exec_state->return_ctx, exec_state->ctx, args);

    switch (ret->type) {
        case RETURN:
            {
                printf("Finished nif exec with %d reschedules\n", exec_state->reschedules);
                return ret->return_term;
            }
        case RESCHEDULE: 
            {
                ERL_NIF_TERM term = enif_make_resource(env, exec_state);
                enif_release_resource(exec_state);
                ERL_NIF_TERM args[] = {term};
                return enif_schedule_nif(env, "nif_resume", 0, nif_resume, 1, args);
            }
    }
}

static ErlNifFunc nif_funcs[] = {
    {"test", 2, nif_test}
};

ERL_NIF_INIT(Elixir.InterruptNif.Native, nif_funcs, nif_load, NULL, NULL, NULL);