shmem_spsc.c

#define _POSIX_C_SOURCE 200809L
#include <stdio.h>
#include <stdlib.h>
#include <stdatomic.h>
#include <stdint.h>
#include <unistd.h>
#include <sched.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <string.h>

typedef struct {
    atomic_size_t head;
    atomic_size_t tail;
    size_t size;
    int data[0];
} shm_ring_t;

static shm_ring_t *shm_create(const char *name, size_t cap)
{
    size_t sz = sizeof(shm_ring_t) + cap * sizeof(int);
    int fd = shm_open(name, O_CREAT | O_RDWR, 0600);
    if (fd < 0) return NULL;
    if (ftruncate(fd, sz) < 0) {
        close(fd);
        return NULL;
    }
    void *p = mmap(NULL, sz, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    close(fd);
    if (p == MAP_FAILED) return NULL;
    shm_ring_t *r = p;
    atomic_init(&r->head, 0);
    atomic_init(&r->tail, 0);
    r->size = cap;
    return r;
}

static shm_ring_t *shm_open_existing(const char *name, size_t cap)
{
    size_t sz = sizeof(shm_ring_t) + cap * sizeof(int);
    int fd = shm_open(name, O_RDWR, 0);
    if (fd < 0) return NULL;
    void *p = mmap(NULL, sz, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    close(fd);
    if (p == MAP_FAILED) return NULL;
    return p;
}

static int shm_push(shm_ring_t *r, int v)
{
    size_t head = atomic_load_explicit(&r->head, memory_order_relaxed);
    size_t tail = atomic_load_explicit(&r->tail, memory_order_acquire);
    if (((head + 1) % r->size) == (tail % r->size)) return -1;
    r->data[head % r->size] = v;
    atomic_store_explicit(&r->head, head + 1, memory_order_release);
    return 0;
}

static int shm_pop(shm_ring_t *r, int *out)
{
    size_t tail = atomic_load_explicit(&r->tail, memory_order_relaxed);
    size_t head = atomic_load_explicit(&r->head, memory_order_acquire);
    if (tail == head) return -1;
    *out = r->data[tail % r->size];
    atomic_store_explicit(&r->tail, tail + 1, memory_order_release);
    return 0;
}

int main(void)
{
    const char *name = "/shmem_spsc_example";
    size_t cap = 1024;
    shm_unlink(name);
    shm_ring_t *r = shm_create(name, cap);
    if (!r) return 1;
    pid_t pid = fork();
    if (pid < 0) return 1;
    if (pid == 0) {
        shm_ring_t *rc = shm_open_existing(name, cap);
        if (!rc) return 1;
        int count = 0;
        while (count < 1000) {
            int v;
            if (shm_pop(rc, &v) == 0) {
                printf("consumer got %d\n", v);
                count++;
            } else {
                sched_yield();
            }
        }
        return 0;
    } else {
        for (int i = 0; i < 1000; i++) {
            while (shm_push(r, i) != 0) {
                sched_yield();
            }
        }
        wait(NULL);
        shm_unlink(name);
    }
    return 0;
}