fwd.c: Check return code from `chdir'.

[fwd] / blast.c

#include <assert.h>
#include <ctype.h>
#include <errno.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/types.h>
#include <sys/time.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>

#include <mLib/alloc.h>
#include <mLib/conn.h>
#include <mLib/mdwopt.h>
#include <mLib/quis.h>
#include <mLib/report.h>
#include <mLib/sel.h>
#include <mLib/tv.h>

typedef struct blast {
  conn c;
  sel_timer t;
} blast;

#define sin saddrin
static struct sockaddr_in sin;
static sel_state sel;
static struct timeval ctv = { 0, 500000 };
static sel_timer sec;
static unsigned count = 0;

static void timers(void);

static void stats(struct timeval *tv, void *p)
{
  static char baton[5] = "/-\\|";
  static char *bt = baton;
  struct timeval ttv;
  printf("\r%c %u connections/second", *bt++, count);
  if (!*bt)
    bt = baton;
  fflush(stdout);
  count = 0;
  TV_ADDL(&ttv, tv, 1, 0);
  sel_addtimer(&sel, &sec, &ttv, stats, p);
}

static void newconn(blast *b);

static void retry(struct timeval *tv, void *p)
{
  blast *b = p;
  newconn(b);
}

static void backoff(blast *b)
{
  struct timeval tv;
  unsigned r = rand();
  double q = RAND_MAX / 5;
  gettimeofday(&tv, 0);
  TV_ADDL(&tv, &tv, r / q, (r % (unsigned)q) * (MILLION / q));
  sel_addtimer(&sel, &b->t, &tv, retry, b);
}

static void connected(int fd, void *p)
{
  blast *b = p;
  sel_rmtimer(&b->t);
  if (fd == -1)
    backoff(b);
  else {
    count++;
    close(fd);
  }
  newconn(b);
}

static void timeout(struct timeval *tv, void *p)
{
  blast *b = p;
  conn_kill(&b->c);
  newconn(b);
}

static void timers(void)
{
  struct tab { void (*func)(struct timeval *tv, void *p); const char *name; }
  tab[] = { { retry, "retry" }, {timeout, "timeout"}, { stats, "stats" }, { 0,
  0 }};
  sel_timer *t = sel.timers;
  while (t) {
    struct tab *q; for (q = tab; q->func != t->func; q++) ;
/*     assert(t->prev->next == t); */
    printf("%lu.%06lu  %p  %s\n", t->tv.tv_sec, t->tv.tv_usec, t->p, q->name);
    assert(t != t->next);
    t = t->next;
  }
  puts("");
}

static void newconn(blast *b)
{
  int fd = socket(PF_INET, SOCK_STREAM, 0);
  struct timeval tv;
  if (fd < 0)
    goto fail;
  gettimeofday(&tv, 0);
  TV_ADD(&tv, &tv, &ctv);
  sel_addtimer(&sel, &b->t, &tv, timeout, b);
  if (conn_init(&b->c, &sel, fd, (struct sockaddr *)&sin, sizeof(sin),
		connected, b))
    goto fail;
  return;

fail:
  sel_rmtimer(&b->t);
  backoff(b);
}

int main(int argc, char *argv[])
{
  blast *b;
  size_t n = 256;

  ego(argv[0]);

  for (;;) {
    int i = getopt(argc, argv, "t:n:");
    if (i < 0)
      break;
    switch (i) {
      case 't': {
	double t = strtod(optarg, 0);
	double s = modf(t, &t);
	ctv.tv_sec = t;
	ctv.tv_usec = s * MILLION;
      } break;
      case 'n':
	n = atoi(optarg);
	break;
      default:
	exit(1);
    }
  }

  argv += optind;
  argc -= optind;
  if (argc != 2) {
    pquis(stderr, "Usage: $ [-t TIME] [-n COUNT] HOST PORT\n");
    exit(1);
  }

  sel_init(&sel);
  sin.sin_family = AF_INET;

  {
    struct hostent *h = gethostbyname(argv[0]);
    if (!h)
      die(1, "bad hostname `%s'", argv[0]);
    memcpy(&sin.sin_addr, h->h_addr, sizeof(struct in_addr));
  }

  if (isdigit((unsigned char)argv[1][0]))
    sin.sin_port = htons(atoi(argv[1]));
  else {
    struct servent *s = getservbyname(argv[1], "tcp");
    if (!s)
      die(1, "bad service name `%s'", argv[1]);
    sin.sin_port = s->s_port;
  }

  b = xmalloc(n * sizeof(blast));

  {
    int i;
    for (i = 0; i < n; i++)
      newconn(&b[i]);
  }

  {
    struct timeval tv;
    gettimeofday(&tv, 0);
    tv.tv_sec++;
    sel_addtimer(&sel, &sec, &tv, stats, 0);
  }

  for (;;) {
    sel_select(&sel);
  }
}