/* A UDP load generator speaking the tracker's own protocol.

   Each thread keeps a socket of its own, does the BEP 15 connect handshake
   once and then announces in a loop, counting replies. Separate sockets
   matter: the kernel spreads datagrams across a SO_REUSEPORT group by a hash
   of the four tuple, so a generator sending everything from one source port
   would pile every packet onto one of the tracker's workers and measure the
   wrong thing.

     cc -O2 -pthread -o udp_load tests/udp_load.c
     ./udp_load 127.0.0.1 6969 4 5 10000 64

   The last argument is how many announces each thread keeps in flight. With
   one, a thread sends and waits, so the offered load is threads divided by
   the round trip and the figure says more about scheduling than about the
   tracker: measured here, a tracker answering 60k/s that way was using
   three quarters of one core. Raising the window offers load without
   needing a thread per outstanding request, which is what lets the tracker
   be pushed to the point where it is actually the limit.

   The last argument is how many distinct torrents to announce. It is what
   keeps the figure honest: a fresh info hash every time measures how fast
   the tracker can insert a million torrents into sorted vectors, which is
   not what a tracker spends its life doing, and which punishes whichever
   configuration is quickest because it reaches the expensive end of the run
   soonest. Cycling a fixed set reaches a steady state where almost every
   announce renews a peer that is already there. */

#include <arpa/inet.h>
#include <netinet/in.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <time.h>
#include <unistd.h>

#define CONNECT_MAGIC 0x41727101980ULL

static struct sockaddr_in g_tracker;
static long               g_torrents = 10000;
static long               g_window   = 64;
static _Atomic long long  g_sent     = 0;
static _Atomic int        g_running  = 1;
static _Atomic long long  g_replies  = 0;
static _Atomic long long  g_timeouts = 0;

static double             seconds(void) {
  struct timespec now;
  clock_gettime(CLOCK_MONOTONIC, &now);
  return now.tv_sec + now.tv_nsec / 1e9;
}

static void put64(unsigned char *at, uint64_t value) {
  int i;
  for (i = 0; i < 8; ++i)
    at[i] = (unsigned char)(value >> (56 - 8 * i));
}
static void put32(unsigned char *at, uint32_t value) {
  at[0] = (unsigned char)(value >> 24);
  at[1] = (unsigned char)(value >> 16);
  at[2] = (unsigned char)(value >> 8);
  at[3] = (unsigned char)value;
}
static uint64_t get64(const unsigned char *at) {
  uint64_t value = 0;
  int      i;
  for (i = 0; i < 8; ++i)
    value = (value << 8) | at[i];
  return value;
}

static void *loader(void *args) {
  long          index = (long)args;
  unsigned char request[98], reply[1500];
  uint64_t      connection_id = 0;
  int           sock          = socket(AF_INET, SOCK_DGRAM, 0);
  struct timeval patience     = {0, 50000}; /* 50 ms */
  unsigned      seed          = (unsigned)index * 2654435761u + 1;

  if (sock < 0)
    return NULL;
  setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &patience, sizeof(patience));

  /* connect handshake */
  memset(request, 0, sizeof(request));
  put64(request, CONNECT_MAGIC);
  put32(request + 8, 0);
  put32(request + 12, (uint32_t)index);
  if (sendto(sock, request, 16, 0, (struct sockaddr *)&g_tracker, sizeof(g_tracker)) != 16) {
    close(sock);
    return NULL;
  }
  if (recvfrom(sock, reply, sizeof(reply), 0, NULL, NULL) < 16) {
    fprintf(stderr, "  thread %ld: no answer to connect\n", index);
    close(sock);
    return NULL;
  }
  connection_id = get64(reply + 8);

  /* Keep g_window announces outstanding: send that many, then drain what has
     come back, rather than waiting for each reply in turn */
  while (atomic_load_explicit(&g_running, memory_order_relaxed)) {
    ssize_t got;
    int     i;
    long    outstanding;

    for (outstanding = 0; outstanding < g_window; ++outstanding) {
    memset(request, 0, sizeof(request));
    put64(request, connection_id);
    put32(request + 8, 1); /* announce */
    put32(request + 12, (uint32_t)index);
    /* One of a fixed set, so that after the first pass almost every announce
       renews a peer rather than creating a torrent. The index is hashed into
       all twenty bytes so the buckets, which key off the first four, are
       spread rather than walked in order. */
    {
      unsigned which = (unsigned)((seed = seed * 1103515245u + 12345u) >> 8) % (unsigned)g_torrents;
      unsigned spread;
      for (i = 0; i < 20; ++i) {
        spread             = which * 2654435761u + (unsigned)i * 40503u;
        spread            ^= spread >> 13;
        request[16 + i]    = (unsigned char)(spread >> 11);
      }
    }
    for (i = 0; i < 20; ++i)
      request[36 + i] = (unsigned char)('a' + (i % 26));
    put64(request + 56, 0);                             /* downloaded */
    put64(request + 64, 0);                             /* left: a seeder */
    put64(request + 72, 0);                             /* uploaded */
    put32(request + 80, 0);                             /* event */
    put32(request + 84, 0);                             /* ip */
    put32(request + 88, (uint32_t)index);               /* key */
    put32(request + 92, (uint32_t)-1);                  /* num_want */
    /* A port of this thread's own, so each thread is a distinct peer and the
       swarms hold more than one */
    request[96] = (unsigned char)((6881 + index) >> 8);
    request[97] = (unsigned char)((6881 + index) & 0xff);

    if (sendto(sock, request, 98, 0, (struct sockaddr *)&g_tracker, sizeof(g_tracker)) != 98)
      break;
    atomic_fetch_add_explicit(&g_sent, 1, memory_order_relaxed);
    }

    for (outstanding = 0; outstanding < g_window; ++outstanding) {
      got = recvfrom(sock, reply, sizeof(reply), 0, NULL, NULL);
      if (got >= 20)
        atomic_fetch_add_explicit(&g_replies, 1, memory_order_relaxed);
      else {
        atomic_fetch_add_explicit(&g_timeouts, 1, memory_order_relaxed);
        break; /* nothing more is coming; go round and send again */
      }
    }
  }

  close(sock);
  return NULL;
}

int main(int argc, char **argv) {
  const char *host    = argc > 1 ? argv[1] : "127.0.0.1";
  int         port    = argc > 2 ? atoi(argv[2]) : 6969;
  long        threads = argc > 3 ? atol(argv[3]) : 16;
  double      run_for = argc > 4 ? atof(argv[4]) : 5.0;
  long        torrents = argc > 5 ? atol(argv[5]) : 10000;
  long        window   = argc > 6 ? atol(argv[6]) : 64;
  pthread_t  *ids     = calloc((size_t)threads, sizeof(pthread_t));
  long        i;
  double      started, elapsed;

  g_torrents = torrents > 0 ? torrents : 1;
  g_window   = window > 0 ? window : 1;

  memset(&g_tracker, 0, sizeof(g_tracker));
  g_tracker.sin_family = AF_INET;
  g_tracker.sin_port   = htons((uint16_t)port);
  if (inet_pton(AF_INET, host, &g_tracker.sin_addr) != 1) {
    fprintf(stderr, "cannot parse address %s\n", host);
    return 1;
  }

  for (i = 0; i < threads; ++i)
    pthread_create(ids + i, NULL, loader, (void *)i);

  started = seconds();
  usleep((useconds_t)(run_for * 1e6));
  atomic_store(&g_running, 0);
  elapsed = seconds() - started;

  for (i = 0; i < threads; ++i)
    pthread_join(ids[i], NULL);

  printf("%.0f replies/s  (%.0f sent/s, %lld replies, %lld lost, %ld threads x %ld in flight, %.1f s)\n",
         (double)atomic_load(&g_replies) / elapsed, (double)atomic_load(&g_sent) / elapsed,
         (long long)atomic_load(&g_replies), (long long)atomic_load(&g_timeouts), threads, g_window, elapsed);
  free(ids);
  return 0;
}