mdw@git.distorted.org.uk Git - fwd/blob - blast.c

   1 #include <assert.h>
   2 #include <ctype.h>
   3 #include <errno.h>
   4 #include <math.h>
   5 #include <stdio.h>
   6 #include <stdlib.h>
   7 #include <string.h>
   8
   9 #include <sys/types.h>
  10 #include <sys/time.h>
  11 #include <unistd.h>
  12 #include <sys/socket.h>
  13 #include <netinet/in.h>
  14 #include <arpa/inet.h>
  15 #include <netdb.h>
  16
  17 #include <mLib/alloc.h>
  18 #include <mLib/conn.h>
  19 #include <mLib/mdwopt.h>
  20 #include <mLib/quis.h>
  21 #include <mLib/report.h>
  22 #include <mLib/sel.h>
  23 #include <mLib/tv.h>
  24
  25 typedef struct blast {
  26   conn c;
  27   sel_timer t;
  28 } blast;
  29
  30 #define sin saddrin
  31 static struct sockaddr_in sin;
  32 static sel_state sel;
  33 static struct timeval ctv = { 0, 500000 };
  34 static sel_timer sec;
  35 static unsigned count = 0;
  36
  37 static void timers(void);
  38
  39 static void stats(struct timeval *tv, void *p)
  40 {
  41   static char baton[5] = "/-\\|";
  42   static char *bt = baton;
  43   struct timeval ttv;
  44   printf("\r%c %u connections/second", *bt++, count);
  45   if (!*bt)
  46     bt = baton;
  47   fflush(stdout);
  48   count = 0;
  49   TV_ADDL(&ttv, tv, 1, 0);
  50   sel_addtimer(&sel, &sec, &ttv, stats, p);
  51 }
  52
  53 static void newconn(blast *b);
  54
  55 static void retry(struct timeval *tv, void *p)
  56 {
  57   blast *b = p;
  58   newconn(b);
  59 }
  60
  61 static void backoff(blast *b)
  62 {
  63   struct timeval tv;
  64   unsigned r = rand();
  65   double q = RAND_MAX / 5;
  66   gettimeofday(&tv, 0);
  67   TV_ADDL(&tv, &tv, r / q, (r % (unsigned)q) * (MILLION / q));
  68   sel_addtimer(&sel, &b->t, &tv, retry, b);
  69 }
  70
  71 static void connected(int fd, void *p)
  72 {
  73   blast *b = p;
  74   sel_rmtimer(&b->t);
  75   if (fd == -1)
  76     backoff(b);
  77   else {
  78     count++;
  79     close(fd);
  80   }
  81   newconn(b);
  82 }
  83
  84 static void timeout(struct timeval *tv, void *p)
  85 {
  86   blast *b = p;
  87   conn_kill(&b->c);
  88   newconn(b);
  89 }
  90
  91 static void timers(void)
  92 {
  93   struct tab { void (*func)(struct timeval *tv, void *p); const char *name; }
  94   tab[] = { { retry, "retry" }, {timeout, "timeout"}, { stats, "stats" }, { 0,
  95   0 }};
  96   sel_timer *t = sel.timers;
  97   while (t) {
  98     struct tab *q; for (q = tab; q->func != t->func; q++) ;
  99 /*     assert(t->prev->next == t); */
 100     printf("%lu.%06lu  %p  %s\n", t->tv.tv_sec, t->tv.tv_usec, t->p, q->name);
 101     assert(t != t->next);
 102     t = t->next;
 103   }
 104   puts("");
 105 }
 106
 107 static void newconn(blast *b)
 108 {
 109   int fd = socket(PF_INET, SOCK_STREAM, 0);
 110   struct timeval tv;
 111   if (fd < 0)
 112     goto fail;
 113   gettimeofday(&tv, 0);
 114   TV_ADD(&tv, &tv, &ctv);
 115   if (conn_init(&b->c, &sel, fd, (struct sockaddr *)&sin, sizeof(sin),
 116                 connected, b))
 117     goto fail;
 118   sel_addtimer(&sel, &b->t, &tv, timeout, b);
 119   return;
 120
 121 fail:
 122   backoff(b);
 123 }
 124
 125 int main(int argc, char *argv[])
 126 {
 127   blast *b;
 128   size_t n = 256;
 129
 130   ego(argv[0]);
 131
 132   for (;;) {
 133     int i = getopt(argc, argv, "t:n:");
 134     if (i < 0)
 135       break;
 136     switch (i) {
 137       case 't': {
 138         double t = strtod(optarg, 0);
 139         double s = modf(t, &t);
 140         ctv.tv_sec = t;
 141         ctv.tv_usec = s * MILLION;
 142       } break;
 143       case 'n':
 144         n = atoi(optarg);
 145         break;
 146       default:
 147         exit(1);
 148     }
 149   }
 150
 151   argv += optind;
 152   argc -= optind;
 153   if (argc != 2) {
 154     pquis(stderr, "Usage: $ [-t TIME] [-n COUNT] HOST PORT\n");
 155     exit(1);
 156   }
 157
 158   sel_init(&sel);
 159   sin.sin_family = AF_INET;
 160
 161   {
 162     struct hostent *h = gethostbyname(argv[0]);
 163     if (!h)
 164       die(1, "bad hostname `%s'", argv[0]);
 165     memcpy(&sin.sin_addr, h->h_addr, sizeof(struct in_addr));
 166   }
 167
 168   if (isdigit((unsigned char)argv[1][0]))
 169     sin.sin_port = htons(atoi(argv[1]));
 170   else {
 171     struct servent *s = getservbyname(argv[1], "tcp");
 172     if (!s)
 173       die(1, "bad service name `%s'", argv[1]);
 174     sin.sin_port = s->s_port;
 175   }
 176
 177   b = xmalloc(n * sizeof(blast));
 178
 179   {
 180     int i;
 181     for (i = 0; i < n; i++)
 182       newconn(&b[i]);
 183   }
 184
 185   {
 186     struct timeval tv;
 187     gettimeofday(&tv, 0);
 188     tv.tv_sec++;
 189     sel_addtimer(&sel, &sec, &tv, stats, 0);
 190   }
 191
 192   for (;;) {
 193     sel_select(&sel);
 194   }
 195 }