/* A receiver with the udp arrangements we are choosing between, so they can
   be compared without any of them being built into the tracker first.

   It echoes what it is sent, so tests/udp_load.c drives it unchanged. There
   is no tracker work here on purpose: with nothing else happening, whatever
   the kernel's receive path costs is the whole figure, which is the clearest
   view of it. Pass a work time to dilute that back towards realism.

     cc -O2 -pthread -o udp_modes tests/udp_modes.c
     ./udp_modes 6969 4 shared 1
     ./udp_modes 6969 4 shared 32
     ./udp_modes 6969 4 per-thread 1
     ./udp_modes 6969 4 per-thread 32

   Arguments: port, threads, shared|per-thread, datagrams per syscall,
   and optionally nanoseconds of pretend work per datagram. */

#define _GNU_SOURCE /* struct mmsghdr on glibc */

#include <errno.h>
#include <netinet/in.h>
#include <pthread.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#if defined(__linux__) || defined(__FreeBSD__)
#define HAVE_MMSG 1
#endif

#define MAX_BATCH   256
#define BUFFER_SIZE 2048

static int         g_port;
static long        g_threads;
static long        g_batch;
static long        g_work_ns;
static int         g_shared_socket = -1; /* -1 means one socket per thread */
static _Atomic int g_running       = 1;

static void        stop(int sig) {
  (void)sig;
  atomic_store(&g_running, 0);
}

/* Which option the platform offers for spreading datagrams over a group of
   sockets bound to one address. They are not the same thing: FreeBSD's plain
   SO_REUSEPORT permits the bind without spreading anything, and Darwin's
   hands every datagram to whichever socket bound last. */
static const char *share_the_port(int sock) {
  int one = 1;

  setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
#ifdef SO_REUSEPORT_LB
  if (!setsockopt(sock, SOL_SOCKET, SO_REUSEPORT_LB, &one, sizeof(one)))
    return "SO_REUSEPORT_LB";
#endif
#ifdef SO_REUSEPORT
  if (!setsockopt(sock, SOL_SOCKET, SO_REUSEPORT, &one, sizeof(one)))
    return "SO_REUSEPORT";
#endif
  return "neither";
}

static int bind_one(int share, const char **option) {
  struct sockaddr_in address;
  int                sock = socket(AF_INET, SOCK_DGRAM, 0);

  if (sock < 0)
    return -1;
  *option = share ? share_the_port(sock) : "exclusive";

  memset(&address, 0, sizeof(address));
  address.sin_family      = AF_INET;
  address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
  address.sin_port        = htons((uint16_t)g_port);
  if (bind(sock, (struct sockaddr *)&address, sizeof(address)) < 0) {
    close(sock);
    return -1;
  }
  return sock;
}

static void pretend_to_work(void) {
  struct timespec started, now;
  long            waited = 0;

  if (!g_work_ns)
    return;
  clock_gettime(CLOCK_MONOTONIC, &started);
  while (waited < g_work_ns) {
    clock_gettime(CLOCK_MONOTONIC, &now);
    waited = (now.tv_sec - started.tv_sec) * 1000000000L + (now.tv_nsec - started.tv_nsec);
  }
}

static void one_at_a_time(int sock) {
  char               buffer[BUFFER_SIZE];
  struct sockaddr_in from;
  socklen_t          length;
  ssize_t            got;

  while (atomic_load_explicit(&g_running, memory_order_relaxed)) {
    length = sizeof(from);
    got    = recvfrom(sock, buffer, sizeof(buffer), 0, (struct sockaddr *)&from, &length);
    if (got < 0) {
      if (errno == EINTR || errno == EAGAIN)
        continue;
      return;
    }
    pretend_to_work();
    sendto(sock, buffer, (size_t)got, 0, (struct sockaddr *)&from, length);
  }
}

#ifdef HAVE_MMSG
static void in_batches(int sock) {
  static _Thread_local char       buffers[MAX_BATCH][BUFFER_SIZE];
  static _Thread_local struct     sockaddr_in froms[MAX_BATCH];
  static _Thread_local struct     iovec iovecs[MAX_BATCH];
  static _Thread_local struct     mmsghdr messages[MAX_BATCH];
  int                             got, i;

  while (atomic_load_explicit(&g_running, memory_order_relaxed)) {
    for (i = 0; i < g_batch; ++i) {
      iovecs[i].iov_base             = buffers[i];
      iovecs[i].iov_len              = BUFFER_SIZE;
      memset(messages + i, 0, sizeof(messages[i]));
      messages[i].msg_hdr.msg_name    = froms + i;
      messages[i].msg_hdr.msg_namelen = sizeof(froms[i]);
      messages[i].msg_hdr.msg_iov     = iovecs + i;
      messages[i].msg_hdr.msg_iovlen  = 1;
    }

    /* WAITFORONE so a quiet receiver behaves like the unbatched one rather
       than holding the first datagram until the batch fills */
    got = recvmmsg(sock, messages, (unsigned)g_batch, MSG_WAITFORONE, NULL);
    if (got < 0) {
      if (errno == EINTR || errno == EAGAIN)
        continue;
      return;
    }

    for (i = 0; i < got; ++i) {
      pretend_to_work();
      iovecs[i].iov_len              = messages[i].msg_len;
      messages[i].msg_hdr.msg_namelen = sizeof(froms[i]);
    }
    sendmmsg(sock, messages, (unsigned)got, 0);
  }
}
#endif

static void *worker(void *args) {
  const char *option = "";
  int         sock   = g_shared_socket;

  (void)args;
  if (sock < 0) {
    sock = bind_one(1, &option);
    if (sock < 0) {
      fprintf(stderr, "  a worker could not bind: %s\n", strerror(errno));
      return NULL;
    }
  }

#ifdef HAVE_MMSG
  if (g_batch > 1)
    in_batches(sock);
  else
#endif
    one_at_a_time(sock);

  if (g_shared_socket < 0)
    close(sock);
  return NULL;
}

int main(int argc, char **argv) {
  const char *how     = argc > 3 ? argv[3] : "shared";
  const char *option  = "";
  pthread_t  *threads;
  long        i;

  g_port    = argc > 1 ? atoi(argv[1]) : 6969;
  g_threads = argc > 2 ? atol(argv[2]) : 1;
  g_batch   = argc > 4 ? atol(argv[4]) : 1;
  g_work_ns = argc > 5 ? atol(argv[5]) : 0;
  if (g_batch < 1)
    g_batch = 1;
  if (g_batch > MAX_BATCH)
    g_batch = MAX_BATCH;

#ifndef HAVE_MMSG
  if (g_batch > 1) {
    fprintf(stderr, "  no recvmmsg on this platform; batching unavailable\n");
    return 2;
  }
#endif

  signal(SIGINT, stop);
  signal(SIGTERM, stop);

  if (!strcmp(how, "shared")) {
    g_shared_socket = bind_one(0, &option);
    if (g_shared_socket < 0) {
      fprintf(stderr, "  bind failed: %s\n", strerror(errno));
      return 1;
    }
  } else {
    /* Probe once so the mode line can say what the platform gave us */
    int probe = bind_one(1, &option);
    if (probe < 0) {
      fprintf(stderr, "  bind failed: %s\n", strerror(errno));
      return 1;
    }
    close(probe);
  }

  printf("  port %d, %ld threads, %s socket%s (%s), %ld per syscall, %ld ns work\n",
         g_port, g_threads, strcmp(how, "shared") ? "one per thread" : "one shared",
         strcmp(how, "shared") ? "s" : "", option, g_batch, g_work_ns);
  fflush(stdout);

  threads = calloc((size_t)g_threads, sizeof(pthread_t));
  for (i = 0; i < g_threads; ++i)
    pthread_create(threads + i, NULL, worker, NULL);
  for (i = 0; i < g_threads; ++i)
    pthread_join(threads[i], NULL);

  if (g_shared_socket >= 0)
    close(g_shared_socket);
  free(threads);
  return 0;
}