/* Drepper's three mutexes, against the real futex system call.
 *
 * The Futex page counts atomic operations and system calls by enumerating a
 * model. This is the same three programs, transcribed from "Futexes Are
 * Tricky" (version 1.6, sections 4 to 6), run on a real Linux kernel, so the
 * model's counts can be held against something that is not the model.
 *
 * Every lock() and unlock() counts its own atomic operations and futex calls
 * as it makes them, and the totals are printed as JSON. strace -c counts the
 * same futex calls from outside the process, which is the second route: the
 * program cannot miscount a call strace also saw.
 *
 *   futex-measure <mutex|mutex2|mutex3> uncontended <iterations>
 *   futex-measure <mutex|mutex2|mutex3> contended <threads> <iterations>
 *   futex-measure <mutex|mutex2|mutex3> held <waiters> <rounds> <hold-ms>
 *
 * Build: cc -O2 -pthread -o futex-measure futex-measure.c
 */
#define _GNU_SOURCE
#include <errno.h>
#include <linux/futex.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/syscall.h>
#include <time.h>
#include <unistd.h>

/* Per-thread tallies, summed at the end. No shared counter, so counting does
 * not add contention of its own to the thing it counts. */
struct tally {
  long locks, unlocks;
  long lock_atomics, unlock_atomics;
  long waits, ewouldblock, ewouldblock_held, slept;
  long wakes, wakes_nobody;
  long contended_locks;          /* the first atomic found the mutex taken */
  long max_lock_atomics, max_lock_waits;
};
static __thread struct tally T;

static int version;              /* 1, 2 or 3 */
static int val;                  /* the futex word */

static long sys_futex(int *addr, int op, int v) {
  return syscall(SYS_futex, addr, op, v, NULL, NULL, 0);
}

/* What the first atomic operation of a lock() call found: 0 free, 1 taken,
 * 2 or more taken with waiters said to exist. The page's model classifies
 * every lock call the same way, so the two can be compared call by call. */
static __thread int first_pending, first_found;
static int seen(int r) {
  if (first_pending) { first_pending = 0; first_found = r; }
  return r;
}

/* Appendix A of the paper, with the counting added. */
static int atomic_inc(int *p) { T.lock_atomics++; return seen(__atomic_fetch_add(p, 1, __ATOMIC_SEQ_CST)); }
static int cmpxchg(int *p, int o, int n) { T.lock_atomics++; return seen(__sync_val_compare_and_swap(p, o, n)); }
static int xchg(int *p, int n) { T.lock_atomics++; return seen(__atomic_exchange_n(p, n, __ATOMIC_SEQ_CST)); }
static int atomic_dec(int *p) { T.unlock_atomics++; return __atomic_fetch_sub(p, 1, __ATOMIC_SEQ_CST); }

/* Set by the holder in "held" mode for exactly as long as it holds the lock
 * and nobody else can. A wait that fails while it is set failed with the
 * mutex held throughout, which is the waste Drepper's first bug is about. */
static int holding;

static __thread long lock_waits_now;
static void futex_wait(int *p, int v) {
  int held = __atomic_load_n(&holding, __ATOMIC_SEQ_CST);
  long r = sys_futex(p, FUTEX_WAIT, v);
  T.waits++; lock_waits_now++;
  if (r == -1 && errno == EAGAIN) {
    T.ewouldblock++;
    if (held && __atomic_load_n(&holding, __ATOMIC_SEQ_CST)) T.ewouldblock_held++;
  } else if (r == 0) T.slept++;
}
static void futex_wake(int *p, int n) {
  long r = sys_futex(p, FUTEX_WAKE, n);
  T.wakes++;
  if (r == 0) T.wakes_nobody++;
}

/* Section 4, Mutex, Take 1. */
static void lock1(void) {
  int c;
  while ((c = atomic_inc(&val)) != 0)
    futex_wait(&val, c + 1);
}
static void unlock1(void) {
  __atomic_store_n(&val, 0, __ATOMIC_SEQ_CST);   /* "val = 0", a plain store */
  futex_wake(&val, 1);
}

/* Section 5, Mutex, Take 2. */
static void lock2(void) {
  int c;
  if ((c = cmpxchg(&val, 0, 1)) != 0)
    do {
      if (c == 2 || cmpxchg(&val, 1, 2) != 0)
        futex_wait(&val, 2);
    } while ((c = cmpxchg(&val, 0, 2)) != 0);
}
static void unlock23(void) {
  if (atomic_dec(&val) != 1) {
    __atomic_store_n(&val, 0, __ATOMIC_SEQ_CST);
    futex_wake(&val, 1);
  }
}

/* Section 6, Mutex, Take 3. */
static void lock3(void) {
  int c;
  if ((c = cmpxchg(&val, 0, 1)) != 0) {
    if (c != 2)
      c = xchg(&val, 2);
    while (c != 0) {
      futex_wait(&val, 2);
      c = xchg(&val, 2);
    }
  }
}

/* Every lock() call's cost, as (what the first atomic found, atomics, waits).
 * Calls past the table's edge are counted, not dropped, so the test can say
 * how many it could not compare. */
#define HMAX 16
static __thread long hist[3][HMAX][HMAX];
static __thread long hist_over;
static long hist_total[3][HMAX][HMAX], hist_over_total;

static void lock(void) {
  long a0 = T.lock_atomics;
  lock_waits_now = 0;
  first_pending = 1;
  if (version == 1) lock1(); else if (version == 2) lock2(); else lock3();
  T.locks++;
  long used = T.lock_atomics - a0;
  if (used > 1 || lock_waits_now) T.contended_locks++;
  if (used > T.max_lock_atomics) T.max_lock_atomics = used;
  if (lock_waits_now > T.max_lock_waits) T.max_lock_waits = lock_waits_now;
  int cls = first_found <= 0 ? 0 : first_found == 1 ? 1 : 2;
  if (used < HMAX && lock_waits_now < HMAX) hist[cls][used][lock_waits_now]++;
  else hist_over++;
}
static void unlock(void) {
  if (version == 1) unlock1(); else unlock23();
  T.unlocks++;
}

static pthread_mutex_t sum_lock = PTHREAD_MUTEX_INITIALIZER;
static struct tally total;
static void fold(void) {
  pthread_mutex_lock(&sum_lock);
  total.locks += T.locks; total.unlocks += T.unlocks;
  total.lock_atomics += T.lock_atomics; total.unlock_atomics += T.unlock_atomics;
  total.waits += T.waits; total.ewouldblock += T.ewouldblock; total.slept += T.slept;
  total.ewouldblock_held += T.ewouldblock_held;
  total.wakes += T.wakes; total.wakes_nobody += T.wakes_nobody;
  total.contended_locks += T.contended_locks;
  if (T.max_lock_atomics > total.max_lock_atomics) total.max_lock_atomics = T.max_lock_atomics;
  if (T.max_lock_waits > total.max_lock_waits) total.max_lock_waits = T.max_lock_waits;
  for (int c = 0; c < 3; c++)
    for (int a = 0; a < HMAX; a++)
      for (int w = 0; w < HMAX; w++) hist_total[c][a][w] += hist[c][a][w];
  hist_over_total += hist_over;
  pthread_mutex_unlock(&sum_lock);
}

static long iters;
static volatile long shared_counter;   /* the critical section's work */

/* Contention has to be arranged, not hoped for. The first version started each
 * thread and let it run, and under strace every thread finished its thousand
 * iterations before the next one existed: four threads, nought contended
 * locks. So they all wait at a start line, and the critical section does a
 * little work, long enough that somebody else usually arrives during it. */
static int started, ready;
static long cs_work;

static void *worker(void *arg) {
  (void)arg;
  memset(&T, 0, sizeof T);
  __atomic_add_fetch(&ready, 1, __ATOMIC_SEQ_CST);
  while (!__atomic_load_n(&started, __ATOMIC_SEQ_CST)) { }
  for (long i = 0; i < iters; i++) {
    lock();
    shared_counter++;
    for (volatile long w = 0; w < cs_work; w++) { }
    unlock();
  }
  fold();
  return NULL;
}

/* One thread holds the lock while the others arrive, so the waits that fail
 * while somebody holds it can be counted with nobody releasing it. */
static void *held_waiter(void *arg) {
  (void)arg;
  memset(&T, 0, sizeof T);
  lock();
  unlock();
  fold();
  return NULL;
}

int main(int argc, char **argv) {
  if (argc < 3) { fprintf(stderr, "usage: see the header\n"); return 2; }
  version = !strcmp(argv[1], "mutex") ? 1 : !strcmp(argv[1], "mutex2") ? 2
          : !strcmp(argv[1], "mutex3") ? 3 : 0;
  if (!version) { fprintf(stderr, "unknown mutex %s\n", argv[1]); return 2; }
  const char *mode = argv[2];
  long threads = 1, rounds = 1, hold_ms = 0;

  if (!strcmp(mode, "uncontended") && argc == 4) {
    iters = atol(argv[3]);
    started = 1;                 /* one thread: nobody to wait for */
    worker(NULL);
  } else if (!strcmp(mode, "contended") && argc == 5) {
    threads = atol(argv[3]); iters = atol(argv[4]);
    cs_work = 2000;
    pthread_t *t = calloc(threads, sizeof *t);
    for (long i = 0; i < threads; i++) pthread_create(&t[i], NULL, worker, NULL);
    while (__atomic_load_n(&ready, __ATOMIC_SEQ_CST) < threads) { }
    __atomic_store_n(&started, 1, __ATOMIC_SEQ_CST);
    for (long i = 0; i < threads; i++) pthread_join(t[i], NULL);
    if (shared_counter != threads * iters) {
      fprintf(stderr, "mutual exclusion failed: %ld of %ld\n", shared_counter, threads * iters);
      return 1;
    }
  } else if (!strcmp(mode, "held") && argc == 6) {
    threads = atol(argv[3]); rounds = atol(argv[4]); hold_ms = atol(argv[5]);
    memset(&T, 0, sizeof T);
    for (long r = 0; r < rounds; r++) {
      lock();
      __atomic_store_n(&holding, 1, __ATOMIC_SEQ_CST);
      pthread_t *t = calloc(threads, sizeof *t);
      for (long i = 0; i < threads; i++) pthread_create(&t[i], NULL, held_waiter, NULL);
      struct timespec ts = { hold_ms / 1000, (hold_ms % 1000) * 1000000L };
      nanosleep(&ts, NULL);
      __atomic_store_n(&holding, 0, __ATOMIC_SEQ_CST);
      unlock();
      for (long i = 0; i < threads; i++) pthread_join(t[i], NULL);
      free(t);
    }
    fold();
  } else {
    fprintf(stderr, "bad arguments for %s\n", mode);
    return 2;
  }
  /* The address, so a trace of every futex call in the process -- glibc's
   * own, for pthread_join, included -- can be narrowed to this one word. */
  printf("{\"futex_addr\": \"%p\", ", (void *)&val);
  printf("\"mutex\": \"%s\", \"mode\": \"%s\", \"threads\": %ld, \"iterations\": %ld, "
         "\"rounds\": %ld, \"hold_ms\": %ld, "
         "\"locks\": %ld, \"unlocks\": %ld, \"lock_atomics\": %ld, \"unlock_atomics\": %ld, "
         "\"futex_wait\": %ld, \"ewouldblock\": %ld, \"ewouldblock_while_held\": %ld, "
         "\"slept\": %ld, "
         "\"futex_wake\": %ld, \"wakes_nobody\": %ld, \"contended_locks\": %ld, "
         "\"max_lock_atomics\": %ld, \"max_lock_waits\": %ld, ",
         argv[1], mode, threads, iters, rounds, hold_ms,
         total.locks, total.unlocks, total.lock_atomics, total.unlock_atomics,
         total.waits, total.ewouldblock, total.ewouldblock_held, total.slept,
         total.wakes, total.wakes_nobody, total.contended_locks,
         total.max_lock_atomics, total.max_lock_waits);
  /* "found:atomics:waits": calls */
  printf("\"lock_costs\": {");
  const char *sep = "";
  for (int c = 0; c < 3; c++)
    for (int a = 0; a < HMAX; a++)
      for (int w = 0; w < HMAX; w++)
        if (hist_total[c][a][w]) {
          printf("%s\"%d:%d:%d\": %ld", sep, c, a, w, hist_total[c][a][w]);
          sep = ", ";
        }
  printf("}, \"lock_costs_past_table\": %ld}\n", hist_over_total);
  return 0;
}