5
\$\begingroup\$

Background Story

I crafted a simple single-threaded coroutine in C, running on Linux x86-64.

Short Technical Explanation

1) Task

There are two types of tasks in this implementation.

  1. Main task. This is the main context that is responsible to enumerate all coroutine tasks. This task uses the main program stack that's given by the kernel during execve().

  2. Coroutine task. This is the sub-context that executes user code. The program can have many coroutine tasks. Each task will execute a function given by the user via a pointer. Every function that is run by the coroutine task is responsible to call schedule() periodically to fairly execute all coroutine tasks. Each coroutine task has its own stack provided by mmap().

2) Scheduling

A function named schedule() is called to yield from a coroutine context. This function gives the control back to the main context. The main context then resumes another runnable task until the schedule() is called again or the user function that's registered for the coroutine task returns.

3) Calling convention

Compatible with System V ABI x86-64. The coroutine task sees the schedule() call just like a function call. The caller is not obliged to be aware of stack switching and %rip jump to another context. The schedule() function preserves all call-saved registers (%r12, %r13, %r14, %r15, %rbp, %rbx, %rsp) w.r.t. the caller.

4) Known limitations

  • The coroutine context is not allowed to block, e.g. calling accept() from a blocking file descriptor, since it will freeze the entire program.

  • Slow tasks enumeration. The task enumerator isn't so good, it enumerates all context slots and checks whether it's runnable or not by NULL checking task->func.

Code

#include <stdio.h>
#include <sys/mman.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stdint.h>
#include <errno.h>

#define ASM_PUSH_CALL_SAVED_REGS        \
        "pushq  %%r12\n\t"              \
        "pushq  %%r13\n\t"              \
        "pushq  %%r14\n\t"              \
        "pushq  %%r15\n\t"              \
        "pushq  %%rbp\n\t"              \
        "pushq  %%rbx\n\t"

#define ASM_POP_CALL_SAVED_REGS         \
        "popq   %%rbx\n\t"              \
        "popq   %%rbp\n\t"              \
        "popq   %%r15\n\t"              \
        "popq   %%r14\n\t"              \
        "popq   %%r13\n\t"              \
        "popq   %%r12\n\t"

#define noinline __attribute__((__noinline__))

#define NR_CORO_TASKS           10
#define UPTR_T(NUM)             ((uintptr_t)(NUM))
#define CORO_STACK_SIZE         (1024ul*1024ul*8ul)
#define PAGE_SIZE               (4096ul)
#define PAGE_ALIGN(PTR)         (void *)(UPTR_T(PTR) & UPTR_T(-4095UL))
#define RESET_RSP(RSP)          PAGE_ALIGN(UPTR_T(RSP) + PAGE_SIZE + CORO_STACK_SIZE)

#ifndef likely
        #define likely(EXPR)    __builtin_expect((EXPR), 1)
#endif
#ifndef unlikely
        #define unlikely(EXPR)  __builtin_expect((EXPR), 0)
#endif

struct coro_task_ctx {
        void            *stack;
        void            *cur_rsp;
        void            *cur_rip;
        void            (*func)(void);
};

/*
 * struct coro_task_ctx fields offset, used for inline ASM.
 */
#define CT_STACK        "0"
#define CT_CUR_RSP      "8"
#define CT_CUR_RIP      "16"
#define CT_FUNC         "24"

static void *g_main_rsp;
static struct coro_task_ctx g_coro_tasks[NR_CORO_TASKS];
static size_t g_coro_cur_task = 0;

static int coro_register_func(void (*func)(void))
{
        size_t i;

        for (i = 0; i < NR_CORO_TASKS; i++) {
                if (!g_coro_tasks[i].func) {
                        g_coro_tasks[i].func = func;
                        return 0;
                }
        }
        return -EAGAIN;
}

static void coro_global_destroy(void)
{
        size_t i;
        void *rsp;

        for (i = 0; i < NR_CORO_TASKS; i++) {
                rsp = g_coro_tasks[i].stack;
                if (!rsp)
                        continue;
                munmap(rsp, CORO_STACK_SIZE);
                memset(&g_coro_tasks[i], 0, sizeof(g_coro_tasks[i]));
        }
}

static int coro_global_init(void)
{
        void *rsp;
        size_t i;
        int ret;

        for (i = 0; i < NR_CORO_TASKS; i++) {
                rsp = mmap(NULL, CORO_STACK_SIZE + 4096UL, PROT_READ|PROT_WRITE,
                           MAP_PRIVATE|MAP_ANONYMOUS|MAP_STACK|MAP_GROWSDOWN,
                           -1, 0);
                if (rsp == MAP_FAILED) {
                        ret = -errno;
                        coro_global_destroy();
                        return ret;
                }

                g_coro_tasks[i].stack = rsp;
                g_coro_tasks[i].cur_rsp = RESET_RSP(rsp);
        }
        return 0;
}

/**
 * This is the entry point for all coroutine tasks.
 */
extern void __coro_task_entry(struct coro_task_ctx *task);
void __coro_task_entry(struct coro_task_ctx *task)
{
        void *rsp;

        /*
         * The @task->func() should call schedule() periodically to
         * fairly execute all coroutine tasks.
         */
        task->func();


        /*
         * The coroutine task has finished, reset its state and give
         * the control back to the main task.
         */
        task->func = NULL;
        task->cur_rip = NULL;
        rsp = task->stack;
        task->cur_rsp = RESET_RSP(rsp);
        __asm__ volatile (
                /*
                 * Load the main task's %rsp.
                 */
                "movq   %[g_main_rsp], %%rsp\n\t"

                /*
                 * Give the control back to the main task.
                 */
                "retq"

                :
                : [g_main_rsp]"m"(g_main_rsp)

                /*
                 * We don't care about registers clobbering here,
                 * because this task has finished, all of them are
                 * unused after the "retq".
                 */
                : "memory", "cc"
        );
        __builtin_unreachable();
}

/*
 * Suspend the coroutine task execution and give the control
 * back to the main task. Must only be called from a coroutine
 * task context.
 */
noinline static void schedule(void)
{
        struct coro_task_ctx *task = &g_coro_tasks[g_coro_cur_task];
        __asm__ volatile (
                /*
                 * Save the coroutine task's regs and return address.
                 */
                ASM_PUSH_CALL_SAVED_REGS
                "leaq   .Lschedule_ret_addr(%%rip), %%rax\n\t"
                "pushq  %%rax\n\t"

                /*
                 * Save the coroutine task's %rsp and return address.
                 */
                "movq   %%rsp, " CT_CUR_RSP "(%[task])\n\t"
                "movq   %%rax, " CT_CUR_RIP "(%[task])\n\t"

                /*
                 * Load the main task's %rsp.
                 */
                "movq   %[g_main_rsp], %%rsp\n\t"

                /*
                 * Give the control back to the main task.
                 */
                "retq\n"

                /*
                 * When the main task calls coro_switch_to(task), it will
                 * jump back here!
                 *
                 * Then we will return to the coroutine task to resume
                 * the execution.
                 */
        ".Lschedule_ret_addr:\n\t"
                ASM_POP_CALL_SAVED_REGS

                : [task]"+D"(task)
                : [g_main_rsp]"m"(g_main_rsp)

                /*
                 * The task sees this as a function call.
                 * Clobber all call-clobbered registers here!
                 */
                : "rax", "rsi", "rdx", "rcx", "r8", "r9", "r10", "r11",
                  "memory", "cc"
        );
}

/*
 * Start the coroutine @task. Must only be called for
 * @task that hasn't started (IOW, @task->cur_rip == NULL).
 */
static void coro_start_task(struct coro_task_ctx *task)
{
        __asm__ volatile (
                /*
                 * Save the main task's call-saved registers and
                 * return address.
                 */
                ASM_PUSH_CALL_SAVED_REGS
                "leaq   .Lcoro_start_task_ret_addr(%%rip), %%rax\n\t"
                "pushq  %%rax\n\t"

                /*
                 * Save the main task's %rsp.
                 */
                "movq   %%rsp, %[g_main_rsp]\n\t"

                /*
                 * Load the coroutine task's %rsp.
                 * At this point, %rsp mod 4096 == 0.
                 */
                "movq   " CT_CUR_RSP "(%[task]), %%rsp\n\t"

                /*
                 * Zero the frame pointer for good backtrace.
                 *
                 * The original %rbp has been saved by
                 * ASM_PUSH_CALL_SAVED_REGS. Will be restored
                 * when the coroutine task calls schedule().
                 */
                "xorl   %%ebp, %%ebp\n\t"

                /*
                 * The System V ABI x86-64 mandates:
                 * On function entry, we must have %rsp mod 16 == 8.
                 */
                "subq   $8, %%rsp\n\t"
                "jmp    __coro_task_entry\n"

                /*
                 * When the coroutine task calls the first schedule(),
                 * it will jump back here!
                 */
        ".Lcoro_start_task_ret_addr:\n\t"
                ASM_POP_CALL_SAVED_REGS

                : [task]"+D"(task),
                  [g_main_rsp]"=m"(g_main_rsp)
                :

                /*
                 * The task sees this as a function call.
                 * Clobber all call-clobbered registers here!
                 *
                 * Call-saved registers have already been
                 * preserved by ASM_{PUSH,POP}_CALL_SAVED_REGS.
                 */
                : "rax", "rsi", "rdx", "rcx", "r8", "r9", "r10", "r11",
                  "memory", "cc"
        );
}

/*
 * Resume a coroutine @task that has been suspended by a
 * schedule() call. Must only be called for already started
 * @task (IOW, @task->cur_rip != NULL).
 */
static void coro_resume_task(struct coro_task_ctx *task)
{
        __asm__ volatile (
                /*
                 * Save the main task's call-saved registers and
                 * return address.
                 */
                ASM_PUSH_CALL_SAVED_REGS
                "leaq   .Lcoro_resume_task_ret_addr(%%rip), %%rax\n\t"
                "pushq  %%rax\n\t"

                /*
                 * Save the main task's %rsp.
                 */
                "movq   %%rsp, %[g_main_rsp]\n\t"

                /*
                 * Load the coroutine task's %rsp.
                 */
                "movq   " CT_CUR_RSP "(%[task]), %%rsp\n\t"

                /*
                 * Give the control to the coroutine task.
                 */
                "retq\n"

                /*
                 * When the coroutine task calls schedule(),
                 * it will jump back here!
                 */
        ".Lcoro_resume_task_ret_addr:\n\t"
                ASM_POP_CALL_SAVED_REGS

                : [task]"+D"(task),
                  [g_main_rsp]"=m"(g_main_rsp)
                :

                /*
                 * The task sees this as a function call.
                 * Clobber all call-clobbered registers here!
                 *
                 * Call-saved registers have already been
                 * preserved by ASM_{PUSH,POP}_CALL_SAVED_REGS.
                 */
                : "rax", "rsi", "rdx", "rcx", "r8", "r9", "r10", "r11",
                  "memory", "cc"
        );
}

static void coro_switch_to(struct coro_task_ctx *task)
{
        if (unlikely(!task->cur_rip)) {
                /*
                 * This @task hasn't started, start it!
                 */
                coro_start_task(task);
        } else {
                /*
                 * This @task has started, but it's suspended
                 * by a schedule() call, resume it!
                 */
                coro_resume_task(task);
        }
}

static void coro_run(void)
{
        struct coro_task_ctx *task, *tasks = g_coro_tasks;
        bool all_clear;
        size_t i;

        do {
                all_clear = true;
                for (i = 0; i < NR_CORO_TASKS; i++) {
                        task = &tasks[i];
                        if (!task->func)
                                continue;

                        all_clear = false;

                        /*
                         * Let the schedule() know which task is 
                         * currently running.
                         */
                        g_coro_cur_task = i;
                        coro_switch_to(task);
                }
        } while (!all_clear);
}

static void func_a(void)
{
        size_t i;

        for (i = 0; i < 5; i++) {
                usleep(100000);
                puts(__func__);
                schedule();
        }
}

static void func_b(void)
{
        size_t i;

        for (i = 0; i < 5; i++) {
                usleep(100000);
                puts(__func__);
                schedule();
        }
}

static void func_c(void)
{
        size_t i;

        for (i = 0; i < 5; i++) {
                usleep(100000);
                puts(__func__);
                schedule();
        }
}

int main(void)
{
        int ret;

        ret = coro_global_init();
        if (ret) {
                errno = -ret;
                perror("coro_global_init");
                return -ret;
        }
        coro_register_func(func_a);
        coro_register_func(func_b);
        coro_register_func(func_c);
        coro_run();

        coro_register_func(func_a);
        coro_register_func(func_b);
        coro_register_func(func_c);
        coro_run();

        coro_global_destroy();
        return 0;
}

I want to get suggestions for context switching design improvement, what is probably missing, what situation that may break with this design, and other improvements.

\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.