A concurrent single-producer/multiple-consumer queue

Author: jserv

README.md

@@ -8,6 +8,7 @@
 - [httpd](httpd/): A multi-threaded web server.
 - [ringbuffer](ringbuffer/): A lock-less ring buffer.
 - [mbus](mbus/): A concurrent message bus.
+- [spmc](spmc/): A concurrent single-producer/multiple-consumer queue
 
 ## License
 

spmc/Makefile

@@ -0,0 +1,8 @@
+all:
+	gcc -Wall -Wextra -o spmc spmc.c -lpthread
+
+clean:
+	rm -f spmc
+
+indent:
+	clang-format -i spmc.c

spmc/spmc.c

@@ -0,0 +1,225 @@
+/* A concurrent single-producer, multiple-consumer (SPMC) queue using C11
+ * Atomics. It is lock-free and atomic, allowing one enqueue-caller/producer,
+ * arbitrary amount of dequeue-callers/consumers.
+ *
+ * Known issue: if one has multiple consumers, some of them will be swapped
+ * off the CPU after grabbing curr_dequeue, and will have dequeued an element
+ * from a different node, if that node ends up having free space.
+ */
+
+#include <assert.h>
+#include <limits.h>
+#include <stdatomic.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+typedef struct __spmc_node {
+    size_t cap; /* One more than the number of slots available */
+    _Atomic size_t front, back;
+    struct __spmc_node *_Atomic next;
+    uintptr_t buf[];
+} spmc_node_t;
+
+typedef void (*spmc_destructor_t)(uintptr_t);
+struct spmc_base {
+    /* current node which enqueues/dequeues */
+    spmc_node_t *_Atomic curr_enqueue, *_Atomic curr_dequeue;
+    uint8_t last_power;
+    spmc_destructor_t destructor;
+};
+typedef struct spmc_base *spmc_ref_t;
+
+#define DEFAULT_INITIAL_POWER 6 /* Initial capacity: 64, as a power of two */
+
+#define SIZE_FROM_CAP(cap, offset) ((cap) * sizeof(uintptr_t) + (offset))
+
+#define MODULO(lhs, rhs) ((lhs) & (rhs - 1)) /* Requires rhs is power of 2 */
+#define INDEX_OF(idx, node) (MODULO((idx), (node)->cap))
+#define IS_READABLE(idx, node) ((node)->back - (idx) != 0)
+#define IS_WRITABLE(idx, node) ((idx) - (node)->front < (node)->cap)
+
+/* The head of the spmc resides contiguously after the spmc_base struct itself.
+ * Here, two objects are stored in the same block of memory, but are accessed
+ * separately.
+ */
+#define HEAD_OF(spmc) ((spmc_node_t *) (void *) ((spmc_ref_t)(spmc) + 1))
+
+static void init_node(spmc_node_t *node, spmc_node_t *next, size_t cap)
+{
+    node->cap = cap;
+    atomic_init(&node->front, 0), atomic_init(&node->back, 0);
+    atomic_init(&node->next, next);
+}
+
+/* In the event initial_cap is 0, the spmc will select a default capacity.
+ * Takes capacities as powers of two. i.e., initial_cap argument of 4 =>
+ * an allocation of ~16 machine words.
+ */
+spmc_ref_t spmc_new(size_t initial_cap, spmc_destructor_t destructor)
+{
+    assert(initial_cap < sizeof(size_t) * CHAR_BIT);
+    const uint8_t power = initial_cap ? initial_cap : DEFAULT_INITIAL_POWER;
+    const size_t cap = 1 << power;
+
+    /* Allocate spmc_base and head spmc_node in the same underlying buffer */
+    spmc_ref_t spmc = malloc(
+        SIZE_FROM_CAP(cap, sizeof(struct spmc_base) + sizeof(spmc_node_t)));
+    spmc_node_t *const head = HEAD_OF(spmc);
+    init_node(head, head, cap);
+
+    atomic_init(&spmc->curr_enqueue, head);
+    atomic_init(&spmc->curr_dequeue, head);
+    spmc->destructor = destructor;
+    spmc->last_power = power;
+
+    return spmc;
+}
+
+/* Destroy the SPMC, freeing all nodes/elements now assoicated with it.
+ * Assume all users of the channel are done with it.
+ */
+void spmc_delete(spmc_ref_t spmc)
+{
+    const spmc_node_t *const head = HEAD_OF(spmc);
+    spmc_node_t *prev;
+    if (spmc->destructor) {
+        for (spmc_node_t *node = head->next; node != head;
+             prev = node, node = node->next, free(prev))
+            for (size_t i = node->front; IS_READABLE(i, node); ++i)
+                spmc->destructor(node->buf[i]);
+    } else {
+        for (spmc_node_t *node = head->next; node != head;
+             prev = node, node = node->next, free(prev))
+            ;
+    }
+    /* Also frees the head; it resides reside in the same buffer. */
+    free(spmc);
+}
+
+/* Send (enqueue) an item onto the SPMC */
+bool spmc_enqueue(spmc_ref_t spmc, uintptr_t element)
+{
+    spmc_node_t *node =
+        atomic_load_explicit(&spmc->curr_enqueue, memory_order_relaxed);
+    size_t idx;
+retry:
+    idx = atomic_load_explicit(&node->back, memory_order_consume);
+    if (!IS_WRITABLE(idx, node)) {
+        spmc_node_t *const next =
+            atomic_load_explicit(&node->next, memory_order_relaxed);
+        /* Never move to write on top of the node that is currently being read;
+         * In that case, items would be read out of order they were enqueued.
+         */
+        if (next !=
+            atomic_load_explicit(&spmc->curr_dequeue, memory_order_relaxed)) {
+            node = next;
+            goto retry;
+        }
+
+        const uint8_t power = ++spmc->last_power;
+        assert(power < sizeof(size_t) * CHAR_BIT);
+        const size_t cap = 1 << power;
+        spmc_node_t *new_node = malloc(SIZE_FROM_CAP(cap, sizeof(spmc_node_t)));
+        if (!new_node)
+            return false;
+
+        init_node(new_node, next, cap);
+        atomic_store_explicit(&node->next, new_node, memory_order_release);
+        idx = 0;
+        node = new_node;
+    }
+    node->buf[INDEX_OF(idx, node)] = element;
+    atomic_store_explicit(&spmc->curr_enqueue, node, memory_order_relaxed);
+    atomic_fetch_add_explicit(&node->back, 1, memory_order_release);
+    return true;
+}
+
+/* Recieve (dequeue) an item from the SPMC */
+bool spmc_dequeue(spmc_ref_t spmc, uintptr_t *slot)
+{
+    spmc_node_t *node =
+        atomic_load_explicit(&spmc->curr_dequeue, memory_order_consume);
+    size_t idx;
+no_increment:
+    do {
+        idx = atomic_load_explicit(&node->front, memory_order_consume);
+        if (!IS_READABLE(idx, node)) {
+            if (node != spmc->curr_enqueue)
+                atomic_compare_exchange_strong(
+                    &spmc->curr_dequeue, &node,
+                    atomic_load_explicit(&node->next, memory_order_relaxed));
+            goto no_increment;
+        } else
+            *slot = node->buf[INDEX_OF(idx, node)];
+    } while (
+        !atomic_compare_exchange_weak(&node->front, &(size_t){idx}, idx + 1));
+    return true;
+}
+
+#include <pthread.h>
+#include <stdio.h>
+#include <string.h>
+
+#define N_ITEMS (1024UL * 8)
+static void *producer_thread(void *arg)
+{
+    spmc_ref_t spmc = arg;
+    for (uintptr_t i = 0; i < N_ITEMS; ++i) {
+        if (!spmc_enqueue(spmc, i))
+            fprintf(stderr, "Failed to enqueue on %zu.\n", (size_t) i);
+    }
+    return NULL;
+}
+
+#define N_MC_ITEMS (1024UL * 8)
+static _Atomic size_t observed_count[N_MC_ITEMS + 1];
+
+static void *mc_thread(void *arg)
+{
+    spmc_ref_t spmc = arg;
+    uintptr_t element = 0, greatest = 0;
+
+    for (;;) {
+        greatest = (greatest > element) ? greatest : element;
+        if (!spmc_dequeue(spmc, &element))
+            fprintf(stderr, "Failed to dequeue in mc_thread.\n");
+        else if (observed_count[element]++)
+            fprintf(stderr, "Consumed twice!\n");
+        else if (element < greatest)
+            fprintf(stderr, "%zu after %zu; bad order!\n", (size_t) element,
+                    (size_t) greatest);
+        printf("Observed %zu.\n", (size_t) element);
+
+        /* Test for sentinel signalling termination */
+        if (element >= (N_MC_ITEMS - 1)) {
+            spmc_enqueue(spmc, element + 1); /* notify other threads */
+            break;
+        }
+    }
+    return NULL;
+}
+
+#define N_MC_THREADS 16
+int main()
+{
+    spmc_ref_t spmc = spmc_new(0, NULL);
+    pthread_t mc[N_MC_THREADS], producer;
+
+    pthread_create(&producer, NULL, producer_thread, spmc);
+    for (int i = 0; i < N_MC_THREADS; i++)
+        pthread_create(&mc[i], NULL, mc_thread, spmc);
+
+    pthread_join(producer, NULL);
+    for (int i = 0; i < N_MC_THREADS; i++)
+        pthread_join(mc[i], NULL);
+
+    for (size_t i = 0; i < N_MC_ITEMS; ++i) {
+        if (observed_count[i] == 1)
+            continue;
+        fprintf(stderr, "An item seen %zu times: %zu.\n", observed_count[i], i);
+    }
+    spmc_delete(spmc);
+    return 0;
+}