coroutine.c

#include "coroutine.h"

static long long get_current_time_ms()
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec * 1000 + tv.tv_usec / 1000;
}

void coroutine_init(coroutine_scheduler_t *scheduler)
{
    memset(scheduler, 0, sizeof(coroutine_scheduler_t));
    scheduler->current_id = -1;
}

int coroutine_create(coroutine_scheduler_t *scheduler, coroutine_func func, void *arg)
{
    for (int i = 0; i < COROUTINE_MAX; ++i)
    {
        if (!scheduler->coroutines[i].is_running)
        {
            coroutine_t *co = &scheduler->coroutines[i];
            co->stack = malloc(COROUTINE_STACK_SIZE);
            if (!co->stack)
            {
                perror("Failed to allocate coroutine stack");
                return -1;
            }
            getcontext(&co->ctx);
            co->ctx.uc_stack.ss_sp = co->stack;
            co->ctx.uc_stack.ss_size = COROUTINE_STACK_SIZE;
            co->ctx.uc_link = &scheduler->main_ctx;
            co->func = func;
            co->arg = arg;
            co->stack_size = COROUTINE_STACK_SIZE;
            co->is_running = 1;
            co->start_time = 0;
            co->elapsed_time = 0;
            co->sleep_time = 0;
            makecontext(&co->ctx, (void (*)())func, 1, arg);
            scheduler->running_count++;
            return i;
        }
    }
    fprintf(stderr, "Too many coroutines!\n");
    return -1;
}

void coroutine_yield(coroutine_scheduler_t *scheduler)
{
    int current_id = scheduler->current_id;
    if (current_id != -1)
    {
        coroutine_t *co = &scheduler->coroutines[current_id];
        co->elapsed_time += get_current_time_ms() - co->start_time;
        scheduler->current_id = -1;
        swapcontext(&co->ctx, &scheduler->main_ctx);
    }
}

void coroutine_resume(coroutine_scheduler_t *scheduler, int id)
{
    coroutine_t *co = &scheduler->coroutines[id];
    scheduler->current_id = id;
    co->start_time = get_current_time_ms();
    swapcontext(&scheduler->main_ctx, &co->ctx);
    if (!co->is_running)
    { // Coroutine has exited
        scheduler->running_count--;
        free(co->stack);
        memset(co, 0, sizeof(coroutine_t));
    }
}

void coroutine_sleep(coroutine_scheduler_t *scheduler, int milliseconds)
{
    int current_id = scheduler->current_id;
    if (current_id != -1)
    {
        coroutine_t *co = &scheduler->coroutines[current_id];
        co->sleep_time = milliseconds;
        coroutine_yield(scheduler);
    }
}

void coroutine_statistics(coroutine_scheduler_t *scheduler)
{
    for (int i = 0; i < COROUTINE_MAX; ++i)
    {
        coroutine_t *co = &scheduler->coroutines[i];
        if (co->is_running)
        {
            printf("Coroutine %d: elapsed time = %lld ms\n", i, co->elapsed_time);
        }
    }
}

void coroutine_destroy(coroutine_scheduler_t *scheduler)
{
    for (int i = 0; i < COROUTINE_MAX; ++i)
    {
        coroutine_t *co = &scheduler->coroutines[i];
        if (co->is_running)
        {
            free(co->stack);
        }
    }
}