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                         perror("epoll_wait");
 249                         exit(1);
 250                 }
 251
 252                 for (i = 0; i < num_active; ++i) {
 253                         struct sender *s = (struct sender *)events[i].data.ptr;
 254                         unsigned long long bytes_to_send = s->bytes_left;
 255                         unsigned ret;
 256
 257                         if (bytes_to_send > 65536) {
 258                                 bytes_to_send = 65536;
 259                         }
 260
 261                         ret = send(s->fd, buf, bytes_to_send, MSG_NOSIGNAL);
 262                         if (ret == -1) {
 263                                 if (errno == EAGAIN)
 264                                         continue;
 265
 266                                 perror("send()");
 267                                 exit(1);
 268                         }
 269
 270                         s->bytes_left -= ret;
 271                         bytes_sent += ret;
 272
 273                         // update the central counter after every 1MB (8ms
 274                         // at gigabit speeds, should be enough) of sent data
 275                         if (bytes_sent > update_frequency) {
 276                                 pthread_mutex_lock(&send_mutex);
 277                                 total_bytes_sent += bytes_sent;
 278                                 pthread_mutex_unlock(&send_mutex);
 279
 280                                 bytes_sent = 0;
 281                         }
 282
 283                         if (s->bytes_left == 0) {
 284                                 if (epoll_ctl(ep_fd, EPOLL_CTL_DEL, s->fd, NULL) == -1) {
 285                                         perror("EPOLL_CTL_DEL");
 286                                         exit(1);
 287                                 }
 288                                 close(s->fd);
 289                                 generate_new_sender(ep_fd, s);
 290                         }
 291                 }
 292         }
 293
 294         pthread_mutex_lock(&send_mutex);
 295         total_bytes_sent += bytes_sent;
 296         pthread_mutex_unlock(&send_mutex);
 297
 298         free(events);
 299         close(ep_fd);
 300
 301         pthread_exit(0);
 302 }
 303
 304 void *receiver_worker(void *arg)
 305 {
 306         int server_sock = (int)arg;
 307         char buf[65536];
 308         unsigned long long bytes_received = 0;
 309         int ep_fd, i;
 310         struct epoll_event ev, *events;
 311
 312         ep_fd = epoll_create(num_sockets_per_sender);
 313         if (ep_fd == -1) {
 314                 perror("epoll_create");
 315                 exit(1);
 316         }
 317
 318         // malloc, since there might not be enough room on the stack
 319         events = (struct epoll_event *)malloc(sizeof(struct epoll_event) * epoll_room_in_receiver);
 320         if (events == NULL) {
 321                 perror("malloc");
 322                 exit(1);
 323         }
 324
 325         // stick the receiver socket in the epoll set
 326         ev.events = EPOLLIN;
 327         ev.data.fd = server_sock;
 328
 329         if (epoll_ctl(ep_fd, EPOLL_CTL_ADD, server_sock, &ev) == -1) {
 330                 perror("EPOLL_CTL_ADD");
 331                 exit(1);
 332         }
 333
 334         for ( ;; ) {
 335                 int num_active = epoll_wait(ep_fd, events, epoll_room_in_receiver, -1);
 336                 if (num_active == -1) {
 337                         perror("epoll_wait");
 338                         exit(1);
 339                 }
 340
 341                 for (i = 0; i < num_active; ++i) {
 342                         int sock = events[i].data.fd;
 343                         int ret;
 344                         if (sock == server_sock) {
 345                                 struct sockaddr_in addr;
 346                                 int sock;
 347                                 socklen_t addr_len = sizeof(addr);
 348
 349                                 sock = accept(server_sock, (struct sockaddr *)&addr, &addr_len);
 350                                 if (sock == -1) {
 351                                         if (errno == EAGAIN) {
 352                                                 // another thread snatched it, ignore
 353                                                 continue;
 354                                         } else {
 355                                                 perror("accept()");
 356                                                 exit(1);
 357                                         }
 358                                 }
 359
 360                                 // add it to the epoll set
 361                                 ev.events = EPOLLIN | EPOLLHUP | EPOLLERR;
 362                                 ev.data.fd = sock;
 363
 364                                 if (epoll_ctl(ep_fd, EPOLL_CTL_ADD, sock, &ev) == -1) {
 365                                         perror("EPOLL_CTL_ADD");
 366                                         exit(1);
 367                                 }
 368                                 continue;
 369                         }
 370
 371                         ret = read(sock, buf, 65536);
 372                         if (ret == 0) {
 373                                 if (epoll_ctl(ep_fd, EPOLL_CTL_DEL, sock, NULL) == -1) {
 374                                         perror("EPOLL_CTL_DEL");
 375                                         exit(1);
 376                                 }
 377                                 close(sock);
 378                                 continue;
 379                         }
 380
 381                         bytes_received += ret;
 382
 383                         // update the central counter after every 1MB (8ms
 384                         // at gigabit speeds, should be enough) of received data
 385                         if (bytes_received > update_frequency) {
 386                                 pthread_mutex_lock(&receive_mutex);
 387                                 total_bytes_received += bytes_received;
 388                                 pthread_mutex_unlock(&receive_mutex);
 389
 390                                 bytes_received = 0;
 391                         }
 392                 }
 393         }
 394
 395         pthread_mutex_lock(&receive_mutex);
 396         total_bytes_received += bytes_received;
 397         pthread_mutex_unlock(&receive_mutex);
 398
 399         pthread_exit(0);
 400 }
 401
 402 // We're keeping this separate receiver dispatcher around because we might eventually
 403 // want to move away from the model where each receiver worker does the accept(). Thus,
 404 // it makes sense to keep this dispatcher around, even though all it does at the moment
 405 // is spawn off a few new threads and then die itself.
 406 void *receiver_dispatcher(void *arg)
 407 {
 408         int server_sock = (int)arg;
 409         int i;
 410         pthread_attr_t attr;
 411         pthread_t thread;
 412
 413         // FIXME: these do not really set errno
 414         if (pthread_attr_init(&attr) != 0) {
 415                 perror("pthread_attr_init()");
 416                 exit(1);
 417         }
 418
 419         if (pthread_attr_setstacksize(&attr, PTHREAD_STACK_MIN + 65536 + 0x4000) != 0) {
 420                 perror("pthread_attr_setstacksize");
 421                 exit(1);
 422         }
 423
 424         if (pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED) != 0) {
 425                 perror("pthread_attr_setdetachstate");
 426                 exit(1);
 427         }
 428
 429         for (i = 0; i < num_receivers; ++i) {
 430                 if (pthread_create(&thread, &attr, receiver_worker, (void *)server_sock) != 0) {
 431                         perror("pthread_create()");
 432                         exit(1);
 433                 }
 434         }
 435
 436         return NULL;
 437 }
 438
 439 int get_server_socket(unsigned short port)
 440 {
 441         int server_sock;
 442         struct sockaddr_in sin;
 443         unsigned one = 1;
 444
 445         server_sock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 446         if (server_sock == -1) {
 447                 perror("socket()");
 448                 exit(1);
 449         }
 450
 451         if (setsockopt(server_sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) == -1) {
 452                 perror("setsocket(SO_REUSEADDR)");
 453                 exit(1);
 454         }
 455
 456         sin.sin_family = AF_INET;
 457         sin.sin_port = htons(port);
 458         sin.sin_addr.s_addr = INADDR_ANY;
 459
 460         if (bind(server_sock, (struct sockaddr *)&sin, sizeof(struct sockaddr)) == -1) {
 461                 perror("bind()");
 462                 exit(1);
 463         }
 464
 465         if (listen(server_sock, 255) == -1) {
 466                 perror("listen()");
 467                 exit(1);
 468         }
 469
 470         return server_sock;
 471 }
 472
 473 int main(int argc, char **argv)
 474 {
 475         int server_sock;
 476         unsigned i;
 477         pthread_attr_t attr;
 478
 479         parse_options(argc, argv);
 480
 481         if (destinations == NULL || sources == NULL) {
 482                 fprintf(stderr, "Missing or empty source or destination host list. Aborting.\n");
 483                 exit(1);
 484         }
 485
 486         if (do_listen) {
 487                 server_sock = get_server_socket(port);
 488         }
 489
 490         printf("Sending data on port %u from %u sources to %u destinations.\n\n",
 491                 port, num_sources, num_destinations);
 492
 493         // FIXME: these do not really set errno
 494         if (pthread_attr_init(&attr) != 0) {
 495                 perror("pthread_attr_init()");
 496                 exit(1);
 497         }
 498
 499         if (pthread_attr_setstacksize(&attr, PTHREAD_STACK_MIN + 65536 + 0x4000) != 0) {
 500                 perror("pthread_attr_setstacksize");
 501                 exit(1);
 502         }
 503
 504         if (pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED) != 0) {
 505                 perror("pthread_attr_setdetachstate");
 506                 exit(1);
 507         }
 508
 509         // Fire off the master receiver.
 510         if (do_listen) {
 511                 pthread_t thread;
 512                 if (pthread_create(&thread, &attr, receiver_dispatcher, (void *)server_sock) != 0) {
 513                         perror("pthread_create()");
 514                         exit(1);
 515                 }
 516         }
 517
 518         printf("Waiting five seconds before starting senders...\n");
 519         sleep(5);
 520
 521         // Fire off sender workers.
 522         for (i = 0; i < num_senders; ++i) {
 523                 pthread_t thread;
 524
 525                 if (pthread_create(&thread, &attr, sender_worker, NULL) != 0) {
 526                         perror("pthread_create()");
 527                         exit(1);
 528                 }
 529         }
 530
 531         // Just stay around collecting statistics until we're done.
 532         for ( ;; ) {
 533                 unsigned long long sent, received;
 534                 static unsigned long long last_sent = 0, last_received = 0;
 535                 double recv_rate, send_rate;
 536
 537                 pthread_mutex_lock(&send_mutex);
 538                 sent = total_bytes_sent;
 539                 pthread_mutex_unlock(&send_mutex);
 540
 541                 pthread_mutex_lock(&receive_mutex);
 542                 received = total_bytes_received;
 543                 pthread_mutex_unlock(&receive_mutex);
 544
 545                 send_rate = (sent - last_sent) * 8.0 / 1048576.0;
 546                 recv_rate = (received - last_received) * 8.0 / 1048576.0;
 547
 548                 printf("%12llu %12llu  %5.0f Mbit/sec  %5.0f Mbit/sec\n", sent, received,
 549                         send_rate, recv_rate);
 550
 551                 last_sent = sent;
 552                 last_received = received;
 553
 554                 sleep(1);
 555         }
 556
 557         exit(0);
 558 }