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