git.sesse.net Git - jam/blob - jam.c

   1 #include <stdio.h>
   2 #include <math.h>
   3 #include <string.h>
   4 #include <stdlib.h>
   5 #include <getopt.h>
   6 #include <netdb.h>
   7 #include <unistd.h>
   8 #include <pthread.h>
   9 #include <sys/socket.h>
  10 #include <netinet/in.h>
  11 #include <sys/epoll.h>
  12 #include <sys/ioctl.h>
  13 #include <errno.h>
  14
  15 unsigned short port = 2007;
  16 unsigned update_frequency = 1048576;
  17
  18 struct in_addr *destinations = NULL;
  19 unsigned num_destinations = 0;
  20 unsigned room_destinations = 0;
  21
  22 struct in_addr *sources = NULL;
  23 unsigned num_sources = 0;
  24 unsigned room_sources = 0;
  25
  26 unsigned num_senders = 128;
  27 unsigned num_sockets_per_sender = 16;
  28 unsigned num_receivers = 128;
  29 unsigned do_listen = 1;
  30
  31 // quite arbitrary, should probably be a bit higher than num_sockets_per_sender
  32 unsigned epoll_room_in_receiver = 128;
  33
  34 unsigned long long total_bytes_received = 0;
  35 pthread_mutex_t receive_mutex = PTHREAD_MUTEX_INITIALIZER;
  36
  37 unsigned long long total_bytes_sent = 0;
  38 pthread_mutex_t send_mutex = PTHREAD_MUTEX_INITIALIZER;
  39
  40 const static struct option longopts[] = {
  41         { "source-list", required_argument, NULL, 's' },
  42         { "destination-list", required_argument, NULL, 'd' },
  43         { "num-senders", required_argument, NULL, 'n' },
  44         { "num-sockets-per-sender", required_argument, NULL, 'N' },
  45         { "num-receivers", required_argument, NULL, 'r' },
  46         { "port", required_argument, NULL, 'p' },
  47         { "sender-only", no_argument, NULL, 'o' },
  48         { NULL, 0, NULL, 0 }
  49 };
  50
  51 // generates from [0,1>
  52 double gen_uniform_random()
  53 {
  54         return rand() / (RAND_MAX+1.0);
  55 }
  56
  57 double gen_pareto_random(double min, double k)
  58 {
  59         double u = gen_uniform_random();
  60         return min * pow(u, -1.0 / k);
  61 }
  62
  63 void read_ip_list(char *filename, struct in_addr **addr_list, unsigned *num, unsigned *room)
  64 {
  65         char buf[256];
  66         FILE *in = fopen(filename, "r");
  67         if (in == NULL) {
  68                 perror(filename);
  69                 exit(1);
  70         }
  71
  72         for ( ;; ) {
  73                 char *ptr;
  74                 struct in_addr addr;
  75                 struct hostent *he;
  76
  77                 if (fgets(buf, 256, in) == NULL)
  78                         break;
  79
  80                 ptr = strchr(buf, '\n');
  81                 if (ptr != NULL)
  82                         *ptr = 0;
  83
  84                 ptr = strchr(buf, '\r');
  85                 if (ptr != NULL)
  86                         *ptr = 0;
  87
  88                 ptr = buf + strspn(buf, " \t");
  89
  90                 if (ptr[0] == '#' || ptr[0] == 0)
  91                         continue;
  92
  93                 he = gethostbyname(ptr);
  94                 if (he == NULL) {
  95                         perror(ptr);
  96                         exit(1);
  97                 }
  98
  99                 // just pick the first for now
 100                 memcpy(&addr.s_addr, he->h_addr_list[0], sizeof(addr.s_addr));
 101
 102                 if (*num >= *room) {
 103                         if (*room == 0) {
 104                                 *room = 16;
 105                         } else {
 106                                 *room <<= 1;
 107                         }
 108                         *addr_list = (struct in_addr *)realloc(*addr_list, *room * sizeof(struct in_addr));
 109                 }
 110
 111                 (*addr_list)[*num] = addr;
 112                 ++*num;
 113         }
 114
 115         fclose(in);
 116 }
 117
 118 void parse_options(int argc, char **argv)
 119 {
 120         int option_index = 0;
 121
 122         for ( ;; ) {
 123                 int c = getopt_long(argc, argv, "s:d:n:N:r:p:o", longopts, &option_index);
 124                 switch (c) {
 125                 case 's':
 126                         read_ip_list(optarg, &sources, &num_sources, &room_sources);
 127                         break;
 128                 case 'd':
 129                         read_ip_list(optarg, &destinations, &num_destinations, &room_destinations);
 130                         break;
 131                 case 'n':
 132                         num_senders = atoi(optarg);
 133                         break;
 134                 case 'r':
 135                         num_receivers = atoi(optarg);
 136                         break;
 137                 case 'N':
 138                         num_sockets_per_sender = atoi(optarg);
 139                         break;
 140                 case 'p':
 141                         port = atoi(optarg);
 142                         break;
 143                 case 'o':
 144                         do_listen = 0;
 145                         break;
 146                 case -1:
 147                         return;       // end of argument list
 148                 default:
 149                         fprintf(stderr, "Invalid option\n");
 150                         exit(1);
 151                 }
 152         }
 153 }
 154
 155 struct sender {
 156         int fd;
 157         unsigned long long bytes_left;
 158 };
 159
 160 void generate_new_sender(int ep_fd, struct sender *s)
 161 {
 162         int sock, one = 1;
 163         unsigned src_num, dst_num;
 164         unsigned long long num_bytes;
 165         struct sockaddr_in sin;
 166         struct epoll_event ev;
 167
 168         sock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 169         if (sock == -1) {
 170                 perror("socket()");
 171                 exit(1);
 172         }
 173
 174         // find the right parameters
 175         src_num = (unsigned)(num_sources * gen_uniform_random());
 176         dst_num = (unsigned)(num_destinations * gen_uniform_random());
 177         num_bytes = (unsigned)gen_pareto_random(1048576.0, 1.0);
 178
 179         // FIXME: bind to the right source
 180
 181         if (ioctl(sock, FIONBIO, &one) == -1) {
 182                 perror("FIONBIO");
 183                 exit(1);
 184         }
 185
 186         sin.sin_family = AF_INET;
 187         sin.sin_port = htons(port);
 188         sin.sin_addr = destinations[dst_num];
 189
 190         // non-blocking connect (will be detected by epoll later)
 191         if (connect(sock, (struct sockaddr *)&sin, sizeof(sin)) == -1 && errno != EINPROGRESS) {
 192                 perror("connect()");
 193                 exit(1);
 194         }
 195
 196         // stick it in the epoll set (FIXME: make edge-triggered?)
 197         ev.events = EPOLLOUT | EPOLLHUP | EPOLLERR;
 198         ev.data.ptr = s;
 199
 200         s->fd = sock;
 201         s->bytes_left = num_bytes;
 202
 203         if (epoll_ctl(ep_fd, EPOLL_CTL_ADD, sock, &ev) == -1) {
 204                 perror("EPOLL_CTL_ADD");
 205                 exit(1);
 206         }
 207 }
 208
 209 void *sender_worker(void *arg)
 210 {
 211         unsigned i;
 212         char buf[65536];
 213         unsigned long long bytes_sent = 0;
 214         int ep_fd;
 215         struct epoll_event *events;
 216
 217         ep_fd = epoll_create(num_sockets_per_sender);
 218         if (ep_fd == -1) {
 219                 perror("epoll_create");
 220                 exit(1);
 221         }
 222
 223         // malloc, since there might not be enough room on the stack
 224         events = (struct epoll_event *)malloc(sizeof(struct epoll_event) * num_sockets_per_sender);
 225         if (events == NULL) {
 226                 perror("malloc");
 227                 exit(1);
 228         }
 229
 230         // fill the buffer with random junk
 231         for (i = 0; i < 65536; ++i)
 232                 buf[i] = rand() & 0xff;
 233
 234         // allocate all the senders
 235         for (i = 0; i < num_sockets_per_sender; ++i) {
 236                 struct sender *s = (struct sender *)malloc(sizeof(struct sender));
 237                 if (s == NULL) {
 238                         perror("malloc()");
 239                         exit(1);
 240                 }
 241
 242                 generate_new_sender(ep_fd, s);
 243         }
 244
 245         for ( ;; ) {
 246                 int num_active = epoll_wait(ep_fd, events, num_sockets_per_sender, -1);
 247                 if (num_active == -1) {
 248                         if (errno == EINTR)
 249                                 continue;
 250                         perror("epoll_wait");
 251                         exit(1);
 252                 }
 253
 254                 for (i = 0; i < num_active; ++i) {
 255                         struct sender *s = (struct sender *)events[i].data.ptr;
 256                         unsigned long long bytes_to_send = s->bytes_left;
 257                         unsigned ret;
 258
 259                         if (bytes_to_send > 65536) {
 260                                 bytes_to_send = 65536;
 261                         }
 262
 263                         ret = send(s->fd, buf, bytes_to_send, MSG_NOSIGNAL);
 264                         if (ret == -1) {
 265                                 if (errno == EAGAIN || errno == EINTR)
 266                                         continue;
 267
 268                                 perror("send()");
 269                                 exit(1);
 270                         }
 271
 272                         s->bytes_left -= ret;
 273                         bytes_sent += ret;
 274
 275                         // update the central counter after every 1MB (8ms
 276                         // at gigabit speeds, should be enough) of sent data
 277                         if (bytes_sent > update_frequency) {
 278                                 pthread_mutex_lock(&send_mutex);
 279                                 total_bytes_sent += bytes_sent;
 280                                 pthread_mutex_unlock(&send_mutex);
 281
 282                                 bytes_sent = 0;
 283                         }
 284
 285                         if (s->bytes_left == 0) {
 286                                 if (epoll_ctl(ep_fd, EPOLL_CTL_DEL, s->fd, NULL) == -1) {
 287                                         perror("EPOLL_CTL_DEL");
 288                                         exit(1);
 289                                 }
 290                                 close(s->fd);
 291                                 generate_new_sender(ep_fd, s);
 292                         }
 293                 }
 294         }
 295
 296         pthread_mutex_lock(&send_mutex);
 297         total_bytes_sent += bytes_sent;
 298         pthread_mutex_unlock(&send_mutex);
 299
 300         free(events);
 301         close(ep_fd);
 302
 303         pthread_exit(0);
 304 }
 305
 306 void *receiver_worker(void *arg)
 307 {
 308         int server_sock = (int)arg;
 309         char buf[65536];
 310         unsigned long long bytes_received = 0;
 311         int ep_fd, i;
 312         struct epoll_event ev, *events;
 313
 314         ep_fd = epoll_create(num_sockets_per_sender);
 315         if (ep_fd == -1) {
 316                 perror("epoll_create");
 317                 exit(1);
 318         }
 319
 320         // malloc, since there might not be enough room on the stack
 321         events = (struct epoll_event *)malloc(sizeof(struct epoll_event) * epoll_room_in_receiver);
 322         if (events == NULL) {
 323                 perror("malloc");
 324                 exit(1);
 325         }
 326
 327         // stick the receiver socket in the epoll set
 328         ev.events = EPOLLIN;
 329         ev.data.fd = server_sock;
 330
 331         if (epoll_ctl(ep_fd, EPOLL_CTL_ADD, server_sock, &ev) == -1) {
 332                 perror("EPOLL_CTL_ADD");
 333                 exit(1);
 334         }
 335
 336         for ( ;; ) {
 337                 int num_active = epoll_wait(ep_fd, events, epoll_room_in_receiver, -1);
 338                 if (num_active == -1) {
 339                         if (errno == EINTR)
 340                                 continue;
 341                         perror("epoll_wait");
 342                         exit(1);
 343                 }
 344
 345                 for (i = 0; i < num_active; ++i) {
 346                         int sock = events[i].data.fd;
 347                         int ret;
 348                         if (sock == server_sock) {
 349                                 struct sockaddr_in addr;
 350                                 int sock;
 351                                 socklen_t addr_len = sizeof(addr);
 352
 353                                 sock = accept(server_sock, (struct sockaddr *)&addr, &addr_len);
 354                                 if (sock == -1) {
 355                                         if (errno == EAGAIN || errno == EINTR) {
 356                                                 // another thread snatched it, ignore
 357                                                 continue;
 358                                         } else {
 359                                                 perror("accept()");
 360                                                 exit(1);
 361                                         }
 362                                 }
 363
 364                                 // add it to the epoll set
 365                                 ev.events = EPOLLIN | EPOLLHUP | EPOLLERR;
 366                                 ev.data.fd = sock;
 367
 368                                 if (epoll_ctl(ep_fd, EPOLL_CTL_ADD, sock, &ev) == -1) {
 369                                         perror("EPOLL_CTL_ADD");
 370                                         exit(1);
 371                                 }
 372                                 continue;
 373                         }
 374
 375                         ret = read(sock, buf, 65536);
 376                         if (ret == 0) {
 377                                 if (epoll_ctl(ep_fd, EPOLL_CTL_DEL, sock, NULL) == -1) {
 378                                         perror("EPOLL_CTL_DEL");
 379                                         exit(1);
 380                                 }
 381                                 close(sock);
 382                                 continue;
 383                         }
 384
 385                         bytes_received += ret;
 386
 387                         // update the central counter after every 1MB (8ms
 388                         // at gigabit speeds, should be enough) of received data
 389                         if (bytes_received > update_frequency) {
 390                                 pthread_mutex_lock(&receive_mutex);
 391                                 total_bytes_received += bytes_received;
 392                                 pthread_mutex_unlock(&receive_mutex);
 393
 394                                 bytes_received = 0;
 395                         }
 396                 }
 397         }
 398
 399         pthread_mutex_lock(&receive_mutex);
 400         total_bytes_received += bytes_received;
 401         pthread_mutex_unlock(&receive_mutex);
 402
 403         pthread_exit(0);
 404 }
 405
 406 // We're keeping this separate receiver dispatcher around because we might eventually
 407 // want to move away from the model where each receiver worker does the accept(). Thus,
 408 // it makes sense to keep this dispatcher around, even though all it does at the moment
 409 // is spawn off a few new threads and then die itself.
 410 void *receiver_dispatcher(void *arg)
 411 {
 412         int server_sock = (int)arg;
 413         int i;
 414         pthread_attr_t attr;
 415         pthread_t thread;
 416
 417         // FIXME: these do not really set errno
 418         if (pthread_attr_init(&attr) != 0) {
 419                 perror("pthread_attr_init()");
 420                 exit(1);
 421         }
 422
 423         if (pthread_attr_setstacksize(&attr, PTHREAD_STACK_MIN + 65536 + 0x4000) != 0) {
 424                 perror("pthread_attr_setstacksize");
 425                 exit(1);
 426         }
 427
 428         if (pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED) != 0) {
 429                 perror("pthread_attr_setdetachstate");
 430                 exit(1);
 431         }
 432
 433         for (i = 0; i < num_receivers; ++i) {
 434                 if (pthread_create(&thread, &attr, receiver_worker, (void *)server_sock) != 0) {
 435                         perror("pthread_create()");
 436                         exit(1);
 437                 }
 438         }
 439
 440         return NULL;
 441 }
 442
 443 int get_server_socket(unsigned short port)
 444 {
 445         int server_sock;
 446         struct sockaddr_in sin;
 447         unsigned one = 1;
 448
 449         server_sock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 450         if (server_sock == -1) {
 451                 perror("socket()");
 452                 exit(1);
 453         }
 454
 455         if (setsockopt(server_sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) == -1) {
 456                 perror("setsocket(SO_REUSEADDR)");
 457                 exit(1);
 458         }
 459
 460         sin.sin_family = AF_INET;
 461         sin.sin_port = htons(port);
 462         sin.sin_addr.s_addr = INADDR_ANY;
 463
 464         if (bind(server_sock, (struct sockaddr *)&sin, sizeof(struct sockaddr)) == -1) {
 465                 perror("bind()");
 466                 exit(1);
 467         }
 468
 469         if (listen(server_sock, 255) == -1) {
 470                 perror("listen()");
 471                 exit(1);
 472         }
 473
 474         return server_sock;
 475 }
 476
 477 int main(int argc, char **argv)
 478 {
 479         int server_sock;
 480         unsigned i;
 481         pthread_attr_t attr;
 482
 483         parse_options(argc, argv);
 484
 485         if (destinations == NULL || sources == NULL) {
 486                 fprintf(stderr, "Missing or empty source or destination host list. Aborting.\n");
 487                 exit(1);
 488         }
 489
 490         if (do_listen) {
 491                 server_sock = get_server_socket(port);
 492         }
 493
 494         printf("Sending data on port %u from %u sources to %u destinations.\n\n",
 495                 port, num_sources, num_destinations);
 496
 497         // FIXME: these do not really set errno
 498         if (pthread_attr_init(&attr) != 0) {
 499                 perror("pthread_attr_init()");
 500                 exit(1);
 501         }
 502
 503         if (pthread_attr_setstacksize(&attr, PTHREAD_STACK_MIN + 65536 + 0x4000) != 0) {
 504                 perror("pthread_attr_setstacksize");
 505                 exit(1);
 506         }
 507
 508         if (pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED) != 0) {
 509                 perror("pthread_attr_setdetachstate");
 510                 exit(1);
 511         }
 512
 513         // Fire off the master receiver.
 514         if (do_listen) {
 515                 pthread_t thread;
 516                 if (pthread_create(&thread, &attr, receiver_dispatcher, (void *)server_sock) != 0) {
 517                         perror("pthread_create()");
 518                         exit(1);
 519                 }
 520         }
 521
 522         printf("Waiting five seconds before starting senders...\n");
 523         sleep(5);
 524
 525         // Fire off sender workers.
 526         for (i = 0; i < num_senders; ++i) {
 527                 pthread_t thread;
 528
 529                 if (pthread_create(&thread, &attr, sender_worker, NULL) != 0) {
 530                         perror("pthread_create()");
 531                         exit(1);
 532                 }
 533         }
 534
 535         // Just stay around collecting statistics until we're done.
 536         for ( ;; ) {
 537                 unsigned long long sent, received;
 538                 static unsigned long long last_sent = 0, last_received = 0;
 539                 double recv_rate, send_rate;
 540
 541                 pthread_mutex_lock(&send_mutex);
 542                 sent = total_bytes_sent;
 543                 pthread_mutex_unlock(&send_mutex);
 544
 545                 pthread_mutex_lock(&receive_mutex);
 546                 received = total_bytes_received;
 547                 pthread_mutex_unlock(&receive_mutex);
 548
 549                 send_rate = (sent - last_sent) * 8.0 / 1048576.0;
 550                 recv_rate = (received - last_received) * 8.0 / 1048576.0;
 551
 552                 printf("%12llu %12llu  %5.0f Mbit/sec  %5.0f Mbit/sec\n", sent, received,
 553                         send_rate, recv_rate);
 554
 555                 last_sent = sent;
 556                 last_received = received;
 557
 558                 sleep(1);
 559         }
 560
 561         exit(0);
 562 }