spsc_ring.c

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

typedef struct {
    void **buf;
    size_t size;
    atomic_size_t head;
    atomic_size_t tail;
} spsc_ring_t;

static int spsc_init(spsc_ring_t *q, size_t cap)
{
    size_t size = 1;
    while (size < cap) size <<= 1;
    q->buf = calloc(size, sizeof(void *));
    if (!q->buf) return -1;
    q->size = size;
    atomic_init(&q->head, 0);
    atomic_init(&q->tail, 0);
    return 0;
}

static void spsc_destroy(spsc_ring_t *q)
{
    free(q->buf);
}

static int spsc_push(spsc_ring_t *q, void *v)
{
    size_t head = atomic_load_explicit(&q->head, memory_order_relaxed);
    size_t tail = atomic_load_explicit(&q->tail, memory_order_acquire);
    if (((head + 1) & (q->size - 1)) == (tail & (q->size - 1))) return -1;
    q->buf[head & (q->size - 1)] = v;
    atomic_store_explicit(&q->head, head + 1, memory_order_release);
    return 0;
}

static int spsc_pop(spsc_ring_t *q, void **out)
{
    size_t tail = atomic_load_explicit(&q->tail, memory_order_relaxed);
    size_t head = atomic_load_explicit(&q->head, memory_order_acquire);
    if (tail == head) return -1;
    void *v = q->buf[tail & (q->size - 1)];
    atomic_store_explicit(&q->tail, tail + 1, memory_order_release);
    *out = v;
    return 0;
}

typedef struct {
    spsc_ring_t *q;
} arg_t;

static void *producer(void *p)
{
    arg_t *a = p;
    for (int i = 0; i < 100000; i++) {
        int *v = malloc(sizeof(int));
        *v = i;
        while (spsc_push(a->q, v) != 0) {
            sched_yield();
        }
    }
    return NULL;
}

static void *consumer(void *p)
{
    arg_t *a = p;
    int count = 0;
    while (count < 100000) {
        void *v;
        if (spsc_pop(a->q, &v) == 0) {
            int *ip = v;
            count++;
            free(ip);
        } else {
            sched_yield();
        }
    }
    return NULL;
}

int main(void)
{
    spsc_ring_t q;
    if (spsc_init(&q, 1024) < 0) return 1;
    pthread_t pt, ct;
    arg_t a;
    a.q = &q;
    pthread_create(&pt, NULL, producer, &a);
    pthread_create(&ct, NULL, consumer, &a);
    pthread_join(pt, NULL);
    pthread_join(ct, NULL);
    spsc_destroy(&q);
    printf("done\n");
    return 0;
}