Add a generic reference counting

Author: jserv

README.md

@@ -12,6 +12,7 @@ purpose of these programs is to be illustrative and educational.
     - [preempt\_sched](preempt_sched/): A preemptive userspace multitasking based on a SIGALRM signal.
 * Multi-threading Paradigms
     - [tpool](tpool/): A lightweight thread pool.
+    - [refcnt](refcnt/): A generic reference counting.
 * [Producer–consumer problem](https://en.wikipedia.org/wiki/Producer%E2%80%93consumer_problem)
     - [spmc](spmc/): A concurrent single-producer/multiple-consumer queue.
     - [mpsc](mpsc/): An unbounded lockless single-consumer/multiple-producer FIFO queue.

refcnt/Makefile

@@ -0,0 +1,5 @@
+all:
+	gcc -o main -Wall -O2 -fsanitize=thread main.c -lpthread
+
+clean:
+	rm -f main

refcnt/main.c

@@ -0,0 +1,176 @@
+/*
+ * Generic reference counting in C.
+ *
+ * Use refcnt_malloc/refcnt_unref in instead of malloc/free. Use refcnt_ref to
+ * duplicate a reference as necessary. Use refcnt unref to stop using a pointer.
+ *
+ * The resulting string must be released using refcnt_unref since refcnt_strdup
+ * function utilizes refcnt_malloc.
+ *
+ * This implementation is thread-safe with the exception of refcnt_realloc. If
+ * you need to use refcnt_realloc in a multi-threaded environment, you must
+ * synchronize access to the reference.
+ *
+ * If you define REFCNT_CHECK, references passed into refcnt_ref and and
+ * refcnt_unref will be checked that they were created by refcnt_malloc. This is
+ * useful for debugging, but it will slow down your program somewhat.
+ *
+ * If you define REFCNT_TRACE, refcnt_malloc and refcnt_unref will print
+ * the line number and file name where they were called. This is useful for
+ * debugging memory leaks or use after free errors.
+ */
+
+#include <assert.h>
+#include <stdatomic.h>
+#include <stddef.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define REFCNT_CHECK
+
+#define maybe_unused __attribute__((unused))
+
+typedef struct {
+#ifdef REFCNT_CHECK
+    int magic;
+#endif
+    atomic_uint refcount;
+    char data[];
+} refcnt_t;
+
+#ifdef REFCNT_TRACE
+#define _REFCNT_TRACE(call)                                                \
+    ({                                                                     \
+        fprintf(stderr, "%s:%d:(%s) %s", __FILE__, __LINE__, __FUNCTION__, \
+                #call);                                                    \
+        call;                                                              \
+    })
+#define refcnt_malloc refcnt_t_malloc
+#define refcnt_realloc refcnt_t_realloc
+#define refcnt_ref refcnt_t_ref
+#define refcnt_unref refcnt_t_unref
+#define refcnt_strdup refcnt_t_strdup
+#endif
+
+#define REFCNT_MAGIC 0xDEADBEEF
+
+static maybe_unused void *refcnt_malloc(size_t len)
+{
+    refcnt_t *ref = malloc(sizeof(refcnt_t) + len);
+    if (!ref)
+        return NULL;
+#ifdef REFCNT_CHECK
+    ref->magic = REFCNT_MAGIC;
+#endif
+    atomic_init(&ref->refcount, 1);
+    return ref->data;
+}
+
+static maybe_unused void *refcnt_realloc(void *ptr, size_t len)
+{
+    refcnt_t *ref = (void *) (ptr - offsetof(refcnt_t, data));
+#ifdef REFCNT_CHECK
+    assert(ref->magic == REFCNT_MAGIC);
+#endif
+    ref = realloc(ref, sizeof(refcnt_t) + len);
+    if (!ref)
+        return NULL;
+    return ref->data;
+}
+
+static maybe_unused void *refcnt_ref(void *ptr)
+{
+    refcnt_t *ref = (void *) (ptr - offsetof(refcnt_t, data));
+#ifdef REFCNT_CHECK
+    assert(ref->magic == REFCNT_MAGIC && "Invalid refcnt pointer");
+#endif
+    atomic_fetch_add(&ref->refcount, 1);
+    return ref->data;
+}
+
+static maybe_unused void refcnt_unref(void *ptr)
+{
+    refcnt_t *ref = (void *) (ptr - offsetof(refcnt_t, data));
+#ifdef REFCNT_CHECK
+    assert(ref->magic == REFCNT_MAGIC && "Invalid refcnt pointer");
+#endif
+    if (atomic_fetch_sub(&ref->refcount, 1) == 1)
+        free(ref);
+}
+
+static maybe_unused char *refcnt_strdup(char *str)
+{
+    refcnt_t *ref = malloc(sizeof(refcnt_t) + strlen(str) + 1);
+    if (!ref)
+        return NULL;
+#ifdef REFCNT_CHECK
+    ref->magic = REFCNT_MAGIC;
+#endif
+    atomic_init(&ref->refcount, 1);
+    strcpy(ref->data, str);
+    return ref->data;
+}
+
+#ifdef REFCNT_TRACE
+#undef refcnt_malloc
+#undef refcnt_realloc
+#undef refcnt_ref
+#undef refcnt_unref
+#undef refcnt_strdup
+#define refcnt_malloc(len) _REFCNT_TRACE(refcnt_t_malloc(len))
+#define refcnt_realloc(ptr, len) _REFCNT_TRACE(refcnt_t_realloc(ptr, len))
+#define refcnt_ref(ptr) _REFCNT_TRACE(refcnt_t_ref(ptr))
+#define refcnt_unref(ptr) _REFCNT_TRACE(refcnt_t_unref(ptr))
+#define refcnt_strdup(ptr) _REFCNT_TRACE(refcnt_t_strdup(ptr))
+#endif
+
+#include <pthread.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define N_ITERATIONS 100
+
+static void *test_thread(void *arg)
+{
+    char *str = arg;
+    for (int i = 0; i < N_ITERATIONS; i++) {
+        char *str2 = refcnt_ref(str);
+        fprintf(stderr, "Thread %u, %i: %s\n", (unsigned int) pthread_self(), i,
+                str2);
+        refcnt_unref(str2);
+    }
+    refcnt_unref(str);
+    return NULL;
+}
+
+#define N_THREADS 64
+
+int main(int argc, char **argv)
+{
+    /* Create threads */
+    pthread_t threads[N_THREADS];
+
+    /* Create a new string that is count referenced */
+    char *str = refcnt_strdup("Hello, world!");
+
+    /* Start the threads, passing a new counted copy of the referece */
+    for (int i = 0; i < N_THREADS; i++)
+        pthread_create(&threads[i], NULL, test_thread, refcnt_ref(str));
+
+    /* We no longer own the reference */
+    refcnt_unref(str);
+
+    /* Whichever thread finishes last will free the string */
+    for (int i = 0; i < N_THREADS; i++)
+        pthread_join(threads[i], NULL);
+
+    void *ptr = malloc(100);
+    /* This should cause a heap overflow while checking the magic num which the
+     * sanitizer checks.
+     * Leaving commented out for now
+     */
+    // refcnt_ref(ptr);
+
+    free(ptr);
+    return 0;
+}