timer_wheel.c

#define _POSIX_C_SOURCE 200809L
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>

typedef void (*timer_cb)(void *);

typedef struct timer_entry {
    uint64_t expiry_tick;
    timer_cb cb;
    void *arg;
    struct timer_entry *next;
} timer_entry_t;

typedef struct {
    timer_entry_t **buckets;
    size_t wheel_size;
    uint64_t tick;
    uint64_t tick_ms;
    int stop;
    pthread_mutex_t m;
    pthread_cond_t cv;
} timer_wheel_t;

static uint64_t now_ms(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (uint64_t)ts.tv_sec * 1000 + ts.tv_nsec / 1000000;
}

static int tw_init(timer_wheel_t *tw, size_t wheel_size, uint64_t tick_ms)
{
    tw->buckets = calloc(wheel_size, sizeof(timer_entry_t *));
    if (!tw->buckets) return -1;
    tw->wheel_size = wheel_size;
    tw->tick = 0;
    tw->tick_ms = tick_ms;
    tw->stop = 0;
    pthread_mutex_init(&tw->m, NULL);
    pthread_cond_init(&tw->cv, NULL);
    return 0;
}

static void tw_destroy(timer_wheel_t *tw)
{
    for (size_t i = 0; i < tw->wheel_size; i++) {
        timer_entry_t *e = tw->buckets[i];
        while (e) {
            timer_entry_t *n = e->next;
            free(e);
            e = n;
        }
    }
    free(tw->buckets);
    pthread_mutex_destroy(&tw->m);
    pthread_cond_destroy(&tw->cv);
}

static void tw_add(timer_wheel_t *tw, uint64_t delay_ms, timer_cb cb, void *arg)
{
    pthread_mutex_lock(&tw->m);
    uint64_t ticks = delay_ms / tw->tick_ms;
    uint64_t expiry = tw->tick + ticks;
    size_t idx = expiry % tw->wheel_size;
    timer_entry_t *e = malloc(sizeof(timer_entry_t));
    e->expiry_tick = expiry;
    e->cb = cb;
    e->arg = arg;
    e->next = tw->buckets[idx];
    tw->buckets[idx] = e;
    pthread_cond_signal(&tw->cv);
    pthread_mutex_unlock(&tw->m);
}

static void tw_stop(timer_wheel_t *tw)
{
    pthread_mutex_lock(&tw->m);
    tw->stop = 1;
    pthread_cond_broadcast(&tw->cv);
    pthread_mutex_unlock(&tw->m);
}

static void *tw_worker(void *p)
{
    timer_wheel_t *tw = p;
    uint64_t base = now_ms();
    for (;;) {
        pthread_mutex_lock(&tw->m);
        if (tw->stop) {
            pthread_mutex_unlock(&tw->m);
            break;
        }
        uint64_t now = now_ms();
        uint64_t elapsed = now - base;
        uint64_t target_tick = elapsed / tw->tick_ms;
        while (tw->tick <= target_tick) {
            size_t idx = tw->tick % tw->wheel_size;
            timer_entry_t *e = tw->buckets[idx];
            tw->buckets[idx] = NULL;
            while (e) {
                timer_entry_t *n = e->next;
                if (e->expiry_tick <= tw->tick) {
                    timer_cb cb = e->cb;
                    void *arg = e->arg;
                    pthread_mutex_unlock(&tw->m);
                    cb(arg);
                    pthread_mutex_lock(&tw->m);
                } else {
                    size_t nidx = e->expiry_tick % tw->wheel_size;
                    e->next = tw->buckets[nidx];
                    tw->buckets[nidx] = e;
                    e = n;
                    continue;
                }
                free(e);
                e = n;
            }
            tw->tick++;
        }
        struct timespec ts;
        ts.tv_sec = 0;
        ts.tv_nsec = tw->tick_ms * 1000000;
        pthread_cond_timedwait(&tw->cv, &tw->m, &ts);
        pthread_mutex_unlock(&tw->m);
    }
    return NULL;
}

static void say_hello(void *arg)
{
    const char *s = arg;
    printf("timer: %s\n", s);
}

int main(void)
{
    timer_wheel_t tw;
    if (tw_init(&tw, 64, 100) < 0) return 1;
    pthread_t th;
    pthread_create(&th, NULL, tw_worker, &tw);
    tw_add(&tw, 500, say_hello, "0.5s");
    tw_add(&tw, 1500, say_hello, "1.5s");
    tw_add(&tw, 2500, say_hello, "2.5s");
    tw_add(&tw, 750, say_hello, "0.75s");
    sleep(3);
    tw_stop(&tw);
    pthread_join(th, NULL);
    tw_destroy(&tw);
    return 0;
}