Introduce a sub system that can store io cookies for file descriptors, thread safely so

6 files changed, 446 insertions, 4 deletions

diff --git a/Makefile b/Makefile
index c7105d4..3bf254a 100644
--- a/Makefile
+++ b/Makefile
@@ -68,8 +68,8 @@ CFLAGS+=-I$(LIBOWFAT_HEADERS) -DGIT_VERSION=$(GIT_VERSION) -Wall -pipe -pthread
 LDFLAGS+=-L$(LIBOWFAT_LIBRARY) -lowfat -pthread
 
 BINARY =opentracker
-HEADERS=trackerlogic.h scan_urlencoded_query.h ot_batch.h ot_buf.h ot_hex.h ot_ip.h ot_config.h ot_scan.h ot_socket.h ot_mutex.h ot_random.h ot_stats.h ot_vector.h ot_clean.h ot_udp.h ot_iovec.h ot_fullscrape.h ot_accesslist.h ot_http.h ot_livesync.h ot_rijndael.h
-SOURCES=opentracker.c trackerlogic.c scan_urlencoded_query.c ot_batch.c ot_buf.c ot_hex.c ot_ip.c ot_config.c ot_scan.c ot_socket.c ot_mutex.c ot_random.c ot_stats.c ot_vector.c ot_clean.c ot_udp.c ot_iovec.c ot_fullscrape.c ot_accesslist.c ot_http.c ot_livesync.c ot_rijndael.c
+HEADERS=trackerlogic.h scan_urlencoded_query.h ot_batch.h ot_buf.h ot_fdstate.h ot_hex.h ot_ip.h ot_config.h ot_scan.h ot_socket.h ot_mutex.h ot_random.h ot_stats.h ot_vector.h ot_clean.h ot_udp.h ot_iovec.h ot_fullscrape.h ot_accesslist.h ot_http.h ot_livesync.h ot_rijndael.h
+SOURCES=opentracker.c trackerlogic.c scan_urlencoded_query.c ot_batch.c ot_buf.c ot_fdstate.c ot_hex.c ot_ip.c ot_config.c ot_scan.c ot_socket.c ot_mutex.c ot_random.c ot_stats.c ot_vector.c ot_clean.c ot_udp.c ot_iovec.c ot_fullscrape.c ot_accesslist.c ot_http.c ot_livesync.c ot_rijndael.c
 SOURCES_proxy=proxy.c ot_batch.c ot_vector.c ot_mutex.c ot_random.c ot_scan.c ot_ip.c ot_socket.c
 
 OBJECTS = $(SOURCES:%.c=%.o)
@@ -119,8 +119,8 @@ $(OBJECTS_proxy_debug): $(HEADERS)
 # has to. tests/stubs.c supplies the few symbols those modules reference but
 # the tests never reach.
 TESTBIN=tests/runner
-TEST_SOURCES=tests/runner.c tests/test.c tests/stubs.c tests/t_scan_urlencoded_query.c tests/t_iovec.c tests/t_vector.c tests/t_random.c tests/t_hex.c tests/t_accesslist.c tests/t_scan.c tests/t_buf.c tests/t_ip.c tests/t_socket.c tests/t_mutex.c tests/t_batch.c
-TEST_UNITS=scan_urlencoded_query.c ot_batch.c ot_buf.c ot_ip.c ot_socket.c ot_iovec.c ot_vector.c ot_random.c ot_hex.c ot_scan.c ot_accesslist.c ot_mutex.c
+TEST_SOURCES=tests/runner.c tests/test.c tests/stubs.c tests/t_scan_urlencoded_query.c tests/t_iovec.c tests/t_vector.c tests/t_random.c tests/t_hex.c tests/t_accesslist.c tests/t_scan.c tests/t_buf.c tests/t_ip.c tests/t_socket.c tests/t_mutex.c tests/t_batch.c tests/t_fdstate.c
+TEST_UNITS=scan_urlencoded_query.c ot_batch.c ot_fdstate.c ot_buf.c ot_ip.c ot_socket.c ot_iovec.c ot_vector.c ot_random.c ot_hex.c ot_scan.c ot_accesslist.c ot_mutex.c
 TEST_HEADERS=tests/test.h $(HEADERS)
 # Tests reach code that sits behind feature flags, so the units under test are
 # built with those features enabled. A test binary is not a production binary.
diff --git a/ot_fdstate.c b/ot_fdstate.c
new file mode 100644
index 0000000..6ddbea3
--- /dev/null
+++ b/ot_fdstate.c
@@ -0,0 +1,164 @@
+/* This software was written by Dirk Engling <erdgeist@erdgeist.org>
+   It is considered beerware. Prost. Skol. Cheers or whatever.
+
+   $id$ */
+
+/* System */
+#include <stdatomic.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <sys/mman.h>
+#include <sys/resource.h>
+#include <sys/types.h>
+
+/* Opentracker */
+#include "ot_fdstate.h"
+
+enum { FD_FREE, FD_IDLE, FD_BUSY, FD_CLOSING };
+
+/* Generation and state share one word so that taking the state is also a
+   check that the generation has not moved */
+#define FD_GEN(w)       ((uint32_t)((w) >> 32))
+#define FD_ST(w)        ((uint32_t)(((w) >> 24) & 0xff))
+#define FD_WORD(g, s)   (((uint64_t)(g) << 32) | ((uint64_t)(s) << 24))
+#define FD_SET_ST(w, s) (((w) & ~(uint64_t)0xff000000) | ((uint64_t)(s) << 24))
+
+typedef struct {
+  _Atomic uint64_t state;
+  _Atomic(void *)  cookie;
+  _Atomic ot_time  expiry;
+} ot_fdslot;
+
+/* A descriptor the kernel can hand out is below the limit it was handed out
+   under, but the limit can be raised after we sized the table, so an out of
+   range descriptor is rejected */
+#define FD_SLOTS_MAX (1u << 20)
+
+static ot_fdslot *g_slots;
+static size_t     g_slot_count;
+
+int fdstate_init(void) {
+  struct rlimit limit;
+  size_t        count = FD_SLOTS_MAX;
+
+  if (getrlimit(RLIMIT_NOFILE, &limit) == 0 && limit.rlim_cur != RLIM_INFINITY && limit.rlim_cur < FD_SLOTS_MAX)
+    count = (size_t)limit.rlim_cur;
+
+  g_slots = mmap(NULL, count * sizeof(ot_fdslot), PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  if (g_slots == MAP_FAILED) {
+    g_slots = NULL;
+    return -1;
+  }
+
+  g_slot_count = count;
+  return 0;
+}
+
+void fdstate_deinit(void) {
+  if (g_slots)
+    munmap(g_slots, g_slot_count * sizeof(ot_fdslot));
+  g_slots      = NULL;
+  g_slot_count = 0;
+}
+
+size_t fdstate_slots(void) { return g_slot_count; }
+
+static ot_fdslot *fdstate_slot(int fd) {
+  if (fd < 0 || (size_t)fd >= g_slot_count)
+    return NULL;
+  return g_slots + fd;
+}
+
+uint32_t fdstate_open(int fd, void *cookie, ot_time expiry) {
+  ot_fdslot *slot = fdstate_slot(fd);
+  uint64_t   word;
+
+  if (!slot)
+    return 0;
+
+  word = atomic_load_explicit(&slot->state, memory_order_relaxed);
+  atomic_store_explicit(&slot->cookie, cookie, memory_order_relaxed);
+  atomic_store_explicit(&slot->expiry, expiry, memory_order_relaxed);
+  /* Publishes the two above to whoever claims the slot next */
+  atomic_store_explicit(&slot->state, FD_WORD(FD_GEN(word), FD_IDLE), memory_order_release);
+
+  return FD_GEN(word);
+}
+
+int fdstate_claim(int fd, uint32_t generation) {
+  ot_fdslot *slot = fdstate_slot(fd);
+  uint64_t   word;
+
+  if (!slot)
+    return 0;
+
+  word = atomic_load_explicit(&slot->state, memory_order_acquire);
+  for (;;) {
+    if (FD_GEN(word) != generation || FD_ST(word) != FD_IDLE)
+      return 0;
+    if (atomic_compare_exchange_weak_explicit(&slot->state, &word, FD_SET_ST(word, FD_BUSY), memory_order_acquire,
+                                              memory_order_relaxed))
+      return 1;
+    /* word now holds what was really there, so the test above runs again */
+  }
+}
+
+void fdstate_release(int fd) {
+  ot_fdslot *slot = fdstate_slot(fd);
+  uint64_t   word;
+
+  if (!slot)
+    return;
+
+  word = atomic_load_explicit(&slot->state, memory_order_relaxed);
+  atomic_store_explicit(&slot->state, FD_SET_ST(word, FD_IDLE), memory_order_release);
+}
+
+int fdstate_claim_expired(int fd, ot_time now) {
+  ot_fdslot *slot = fdstate_slot(fd);
+  uint64_t   word;
+  ot_time    expiry;
+
+  if (!slot)
+    return 0;
+
+  /* Reading before asking for the slot keeps the common case, a descriptor
+     that has not expired, off the bus entirely */
+  word = atomic_load_explicit(&slot->state, memory_order_relaxed);
+  if (FD_ST(word) != FD_IDLE)
+    return 0;
+  expiry = atomic_load_explicit(&slot->expiry, memory_order_relaxed);
+  if (!expiry || expiry > now)
+    return 0;
+
+  /* Nothing can move the expiry while the slot sits at FD_IDLE, so winning
+     the slot is the last word */
+  return atomic_compare_exchange_strong_explicit(&slot->state, &word, FD_SET_ST(word, FD_CLOSING), memory_order_acquire,
+                                                 memory_order_relaxed);
+}
+
+void fdstate_close(int fd) {
+  ot_fdslot *slot = fdstate_slot(fd);
+  uint64_t   word;
+
+  if (!slot)
+    return;
+
+  word = atomic_load_explicit(&slot->state, memory_order_relaxed);
+  atomic_store_explicit(&slot->cookie, NULL, memory_order_relaxed);
+  atomic_store_explicit(&slot->expiry, 0, memory_order_relaxed);
+  atomic_store_explicit(&slot->state, FD_WORD(FD_GEN(word) + 1, FD_FREE), memory_order_release);
+}
+
+void *fdstate_cookie(int fd) {
+  ot_fdslot *slot = fdstate_slot(fd);
+
+  return slot ? atomic_load_explicit(&slot->cookie, memory_order_relaxed) : NULL;
+}
+
+void fdstate_set_expiry(int fd, ot_time expiry) {
+  ot_fdslot *slot = fdstate_slot(fd);
+
+  if (slot)
+    atomic_store_explicit(&slot->expiry, expiry, memory_order_relaxed);
+}
diff --git a/ot_fdstate.h b/ot_fdstate.h
new file mode 100644
index 0000000..2693a3b
--- /dev/null
+++ b/ot_fdstate.h
@@ -0,0 +1,52 @@
+/* This software was written by Dirk Engling <erdgeist@erdgeist.org>
+   It is considered beerware. Prost. Skol. Cheers or whatever.
+
+   $id$ */
+
+#ifndef OT_FDSTATE_H__
+#define OT_FDSTATE_H__
+
+#include <stdint.h>
+
+#include "trackerlogic.h"
+
+/* What the tracker knows about a file descriptor: who owns it, when it
+   should be given up on, and whether anyone is working on it right now.
+
+   A descriptor number says nothing about identity, since the kernel hands
+   the same number out again as soon as it is closed. Every slot therefore
+   carries a generation that is bumped when the descriptor is given up, and
+   every caller that was told about a descriptor has to say which generation
+   it was told about. A caller holding a stale generation is turned away
+   instead of being handed somebody else's connection.
+
+   The table is sized once from RLIMIT_NOFILE and never moves, so looking a
+   descriptor up needs no synchronisation at all. */
+
+int      fdstate_init(void);
+void     fdstate_deinit(void);
+size_t   fdstate_slots(void);
+
+/* Take a freshly accepted descriptor into the table and return the
+   generation to quote when asking for it back. expiry of 0 never expires. */
+uint32_t fdstate_open(int fd, void *cookie, ot_time expiry);
+
+/* Ask for exclusive use. Returns 0 if the descriptor has been recycled since
+   generation was handed out, or if somebody else is using it. */
+int      fdstate_claim(int fd, uint32_t generation);
+void     fdstate_release(int fd);
+
+/* The cleaner's side of the same thing: takes the descriptor only if nobody
+   is working on it and it really has expired. Never waits. */
+int      fdstate_claim_expired(int fd, ot_time now);
+
+/* Give the descriptor up. The generation is bumped here, which has to happen
+   before close(2) so that the kernel cannot reissue the number while events
+   quoting the old generation are still in flight. */
+void     fdstate_close(int fd);
+
+/* Both only valid while the caller holds the descriptor. */
+void    *fdstate_cookie(int fd);
+void     fdstate_set_expiry(int fd, ot_time expiry);
+
+#endif
diff --git a/tests/runner.c b/tests/runner.c
index 148b928..335d83d 100644
--- a/tests/runner.c
+++ b/tests/runner.c
@@ -25,6 +25,7 @@ static const struct module modules[] = {{"scan_urlencoded_query", test_scan_urle
                                         {"ot_socket", test_socket},
                                         {"ot_mutex", test_mutex},
                                         {"ot_batch", test_batch},
+                                        {"ot_fdstate", test_fdstate},
                                         {"ot_accesslist", test_accesslist},
                                         {NULL, NULL}};
 
diff --git a/tests/t_fdstate.c b/tests/t_fdstate.c
new file mode 100644
index 0000000..bf8724a
--- /dev/null
+++ b/tests/t_fdstate.c
@@ -0,0 +1,224 @@
+/* This software was written by Dirk Engling <erdgeist@erdgeist.org>
+   It is considered beerware. Prost. Skol. Cheers or whatever.
+
+   $id$ */
+
+/* System */
+#include <pthread.h>
+#include <stdatomic.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+/* Opentracker */
+#include "ot_fdstate.h"
+#include "test.h"
+
+/* ------------------------------------------------------- one thread first */
+
+static void case_open_claim_release(void) {
+  char     cookie[4];
+  uint32_t generation;
+
+  test_case("a descriptor taken into the table can be claimed and gives back its cookie");
+  generation = fdstate_open(7, cookie, 0);
+  CHECK_INT(fdstate_claim(7, generation), 1);
+  CHECK_PTR(fdstate_cookie(7), cookie);
+  fdstate_release(7);
+
+  test_case("and can be claimed again once released");
+  CHECK_INT(fdstate_claim(7, generation), 1);
+  fdstate_release(7);
+  fdstate_close(7);
+}
+
+static void case_one_holder_at_a_time(void) {
+  uint32_t generation = fdstate_open(8, NULL, 0);
+
+  test_case("a descriptor somebody holds cannot be claimed again");
+  CHECK_INT(fdstate_claim(8, generation), 1);
+  CHECK_INT(fdstate_claim(8, generation), 0);
+  fdstate_release(8);
+  CHECK_INT(fdstate_claim(8, generation), 1);
+  fdstate_release(8);
+  fdstate_close(8);
+}
+
+static void case_stale_generation_is_refused(void) {
+  char     first[4], second[4];
+  uint32_t old, new;
+
+  test_case("a caller quoting the generation of a closed descriptor is turned away");
+  old = fdstate_open(9, first, 0);
+  fdstate_close(9);
+  new = fdstate_open(9, second, 0);
+  CHECK(old != new);
+  CHECK_INT(fdstate_claim(9, old), 0);
+
+  test_case("while the one quoting the new generation gets the new cookie");
+  CHECK_INT(fdstate_claim(9, new), 1);
+  CHECK_PTR(fdstate_cookie(9), second);
+  fdstate_release(9);
+  fdstate_close(9);
+}
+
+static void case_expiry(void) {
+  uint32_t generation = fdstate_open(10, NULL, 100);
+
+  test_case("a descriptor past its expiry is collected");
+  CHECK_INT(fdstate_claim_expired(10, 99), 0);
+  CHECK_INT(fdstate_claim_expired(10, 100), 1);
+  fdstate_close(10);
+
+  test_case("one that never expires is never collected");
+  generation = fdstate_open(10, NULL, 0);
+  CHECK_INT(fdstate_claim_expired(10, 1 << 30), 0);
+
+  test_case("and one somebody is working on is left alone");
+  fdstate_set_expiry(10, 1);
+  CHECK_INT(fdstate_claim(10, generation), 1);
+  CHECK_INT(fdstate_claim_expired(10, 1 << 30), 0);
+  fdstate_release(10);
+  CHECK_INT(fdstate_claim_expired(10, 1 << 30), 1);
+  fdstate_close(10);
+}
+
+static void case_out_of_range(void) {
+  test_case("a descriptor outside the table is refused rather than trusted");
+  CHECK_INT(fdstate_claim(-1, 0), 0);
+  CHECK_INT(fdstate_claim((int)fdstate_slots(), 0), 0);
+  CHECK_INT(fdstate_claim_expired(-1, 0), 0);
+  CHECK_PTR(fdstate_cookie(-1), NULL);
+}
+
+/* --------------------------------------------------------- then in anger */
+
+/* The race this is about only exists while a descriptor is being recycled,
+   so recycling is driven by work done rather than by a clock: every slot is
+   given up after RECYCLE_AFTER claims, whatever the machine's speed. */
+
+#define STRESS_SLOTS   64
+#define RECYCLE_AFTER  32
+#define STRESS_CLAIMS  200000
+#define STRESS_GIVEUP  (STRESS_CLAIMS * 100)
+
+static _Atomic long     stress_claims, stress_recycles, stress_violations, stress_attempts;
+static _Atomic int      stress_lives[STRESS_SLOTS];
+static _Atomic int      stress_live[STRESS_SLOTS];
+
+/* Generation and cookie are published in one word so a worker can never read
+   a cookie that belongs to a different generation than the one it claimed
+   with. Each connection gets a cookie nobody else ever gets. */
+static _Atomic uint64_t stress_conn[STRESS_SLOTS];
+static _Atomic uint64_t stress_next_cookie;
+
+#define STRESS_CONN(cookie, gen) (((uint64_t)(cookie) << 32) | (uint32_t)(gen))
+#define STRESS_CONN_COOKIE(w)    ((void *)(uintptr_t)((w) >> 32))
+#define STRESS_CONN_GEN(w)       ((uint32_t)(w))
+
+static int stress_running(void) {
+  return atomic_load(&stress_claims) < STRESS_CLAIMS && atomic_load(&stress_attempts) < STRESS_GIVEUP;
+}
+
+static void *stress_worker(void *arg) {
+  unsigned seed = (unsigned)(uintptr_t)arg;
+
+  while (stress_running()) {
+    int      fd   = (int)(rand_r(&seed) % STRESS_SLOTS);
+    uint64_t conn = atomic_load(&stress_conn[fd]);
+
+    atomic_fetch_add(&stress_attempts, 1);
+    if (!fdstate_claim(fd, STRESS_CONN_GEN(conn)))
+      continue;
+
+    if (fdstate_cookie(fd) != STRESS_CONN_COOKIE(conn))
+      atomic_fetch_add(&stress_violations, 1);
+    if (atomic_fetch_sub(&stress_lives[fd], 1) <= 1)
+      fdstate_set_expiry(fd, 1); /* retired by its holder, as a handler would */
+    if (fdstate_cookie(fd) != STRESS_CONN_COOKIE(conn))
+      atomic_fetch_add(&stress_violations, 1);
+
+    fdstate_release(fd);
+    atomic_fetch_add(&stress_claims, 1);
+  }
+  return NULL;
+}
+
+static void *stress_keeper(void *arg) {
+  (void)arg;
+  while (stress_running()) {
+    int fd;
+
+    for (fd = 0; fd < STRESS_SLOTS; ++fd) {
+      if (atomic_load(&stress_live[fd])) {
+        if (!fdstate_claim_expired(fd, 1))
+          continue;                            /* a worker has it, come back */
+        fdstate_close(fd);
+        atomic_store(&stress_live[fd], 0);
+        atomic_fetch_add(&stress_recycles, 1);
+      }
+      {
+        void    *cookie = (void *)(uintptr_t)(atomic_fetch_add(&stress_next_cookie, 1) + 1);
+        uint32_t gen    = fdstate_open(fd, cookie, 0);
+
+        atomic_store(&stress_lives[fd], RECYCLE_AFTER);
+        atomic_store(&stress_conn[fd], STRESS_CONN((uintptr_t)cookie, gen));
+        atomic_store(&stress_live[fd], 1);
+      }
+    }
+  }
+  return NULL;
+}
+
+static void case_recycling_under_load(void) {
+  pthread_t worker[4], keeper;
+  size_t    i;
+
+  test_case("no worker is ever handed a connection from another generation");
+  atomic_store(&stress_claims, 0);
+  atomic_store(&stress_recycles, 0);
+  atomic_store(&stress_violations, 0);
+  atomic_store(&stress_attempts, 0);
+  atomic_store(&stress_next_cookie, 0);
+  /* Every slot is opened and published before any worker runs: a worker only
+     ever learns of a descriptor from whoever registered it, and never
+     invents a generation of its own. */
+  for (i = 0; i < STRESS_SLOTS; ++i) {
+    void    *cookie = (void *)(uintptr_t)(atomic_fetch_add(&stress_next_cookie, 1) + 1);
+    uint32_t gen    = fdstate_open((int)i, cookie, 0);
+
+    atomic_store(&stress_lives[i], RECYCLE_AFTER);
+    atomic_store(&stress_conn[i], STRESS_CONN((uintptr_t)cookie, gen));
+    atomic_store(&stress_live[i], 1);
+  }
+
+  pthread_create(&keeper, NULL, stress_keeper, NULL);
+  for (i = 0; i < sizeof(worker) / sizeof(worker[0]); ++i)
+    pthread_create(&worker[i], NULL, stress_worker, (void *)(uintptr_t)(i + 1));
+  for (i = 0; i < sizeof(worker) / sizeof(worker[0]); ++i)
+    pthread_join(worker[i], NULL);
+  pthread_join(keeper, NULL);
+
+  CHECK_INT(atomic_load(&stress_violations), 0);
+
+  test_case("and the run actually recycled descriptors, or it proved nothing");
+  CHECK(atomic_load(&stress_claims) >= STRESS_CLAIMS);
+  CHECK(atomic_load(&stress_recycles) > STRESS_CLAIMS / (RECYCLE_AFTER * 4));
+
+  for (i = 0; i < STRESS_SLOTS; ++i)
+    fdstate_close((int)i);
+}
+
+void test_fdstate(void) {
+  CHECK_INT(fdstate_init(), 0);
+  CHECK(fdstate_slots() > STRESS_SLOTS);
+
+  case_open_claim_release();
+  case_one_holder_at_a_time();
+  case_stale_generation_is_refused();
+  case_expiry();
+  case_out_of_range();
+  case_recycling_under_load();
+
+  fdstate_deinit();
+}
diff --git a/tests/test.h b/tests/test.h
index fd574c2..bc8c3bd 100644
--- a/tests/test.h
+++ b/tests/test.h
@@ -88,5 +88,6 @@ void test_ip(void);
 void test_socket(void);
 void test_mutex(void);
 void test_batch(void);
+void test_fdstate(void);
 
 #endif