ustd/ustd.c

   1 /* Copyright (C) 2009  Pierre-Marc Fournier
   2  *
   3  * This library is free software; you can redistribute it and/or
   4  * modify it under the terms of the GNU Lesser General Public
   5  * License as published by the Free Software Foundation; either
   6  * version 2.1 of the License, or (at your option) any later version.
   7  *
   8  * This library is distributed in the hope that it will be useful,
   9  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  10  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  11  * Lesser General Public License for more details.
  12  *
  13  * You should have received a copy of the GNU Lesser General Public
  14  * License along with this library; if not, write to the Free Software
  15  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
  16  */
  17
  18 #define _GNU_SOURCE
  19
  20 #include <sys/types.h>
  21 #include <sys/stat.h>
  22 #include <sys/shm.h>
  23 #include <fcntl.h>
  24 #include <unistd.h>
  25 #include <pthread.h>
  26 #include <signal.h>
  27
  28 #include <stdlib.h>
  29 #include <stdio.h>
  30 #include <string.h>
  31 #include <errno.h>
  32 #include <assert.h>
  33 #include <getopt.h>
  34
  35 #include "ustd.h"
  36 #include "usterr.h"
  37 #include "ustcomm.h"
  38
  39 /* return value: 0 = subbuffer is finished, it won't produce data anymore
  40  *               1 = got subbuffer successfully
  41  *               <0 = error
  42  */
  43
  44 #define GET_SUBBUF_OK 1
  45 #define GET_SUBBUF_DONE 0
  46 #define GET_SUBBUF_DIED 2
  47
  48 #define PUT_SUBBUF_OK 1
  49 #define PUT_SUBBUF_DIED 0
  50 #define PUT_SUBBUF_PUSHED 2
  51
  52 char *sock_path=NULL;
  53 char *trace_path=NULL;
  54 int daemon_mode = 0;
  55 char *pidfile = NULL;
  56
  57 /* Number of active buffers and the mutex to protect it. */
  58 int active_buffers = 0;
  59 pthread_mutex_t active_buffers_mutex = PTHREAD_MUTEX_INITIALIZER;
  60 /* Whether a request to end the program was received. */
  61 sig_atomic_t terminate_req = 0;
  62
  63 int get_subbuffer(struct buffer_info *buf)
  64 {
  65         char *send_msg=NULL;
  66         char *received_msg=NULL;
  67         char *rep_code=NULL;
  68         int retval;
  69         int result;
  70
  71         asprintf(&send_msg, "get_subbuffer %s", buf->name);
  72         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
  73         if((result == -1 && errno == EPIPE) || result == 0) {
  74                 DBG("app died while being traced");
  75                 retval = GET_SUBBUF_DIED;
  76                 goto end;
  77         }
  78         else if(result < 0) {
  79                 ERR("get_subbuffer: ustcomm_send_request failed");
  80                 retval = -1;
  81                 goto end;
  82         }
  83
  84         result = sscanf(received_msg, "%as %ld", &rep_code, &buf->consumed_old);
  85         if(result != 2 && result != 1) {
  86                 ERR("unable to parse response to get_subbuffer");
  87                 retval = -1;
  88                 goto end_rep;
  89         }
  90
  91         DBG("received msg is %s", received_msg);
  92
  93         if(!strcmp(rep_code, "OK")) {
  94                 DBG("got subbuffer %s", buf->name);
  95                 retval = GET_SUBBUF_OK;
  96         }
  97         else if(nth_token_is(received_msg, "END", 0) == 1) {
  98                 retval = GET_SUBBUF_DONE;
  99                 goto end_rep;
 100         }
 101         else {
 102                 DBG("error getting subbuffer %s", buf->name);
 103                 retval = -1;
 104         }
 105
 106         /* FIMXE: free correctly the stuff */
 107 end_rep:
 108         if(rep_code)
 109                 free(rep_code);
 110 end:
 111         if(send_msg)
 112                 free(send_msg);
 113         if(received_msg)
 114                 free(received_msg);
 115
 116         return retval;
 117 }
 118
 119 int put_subbuffer(struct buffer_info *buf)
 120 {
 121         char *send_msg=NULL;
 122         char *received_msg=NULL;
 123         char *rep_code=NULL;
 124         int retval;
 125         int result;
 126
 127         asprintf(&send_msg, "put_subbuffer %s %ld", buf->name, buf->consumed_old);
 128         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 129         if(result < 0 && errno == ECONNRESET) {
 130                 retval = PUT_SUBBUF_DIED;
 131                 goto end;
 132         }
 133         if(result < 0) {
 134                 ERR("put_subbuffer: send_message failed");
 135                 retval = -1;
 136                 goto end;
 137         }
 138
 139         result = sscanf(received_msg, "%as", &rep_code);
 140         if(result != 1) {
 141                 ERR("unable to parse response to put_subbuffer");
 142                 retval = -1;
 143                 goto end_rep;
 144         }
 145
 146         if(!strcmp(rep_code, "OK")) {
 147                 DBG("subbuffer put %s", buf->name);
 148                 retval = PUT_SUBBUF_OK;
 149         }
 150         else {
 151                 DBG("put_subbuffer: received error, we were pushed");
 152                 retval = PUT_SUBBUF_PUSHED;
 153                 goto end_rep;
 154         }
 155
 156 end_rep:
 157         if(rep_code)
 158                 free(rep_code);
 159
 160 end:
 161         if(send_msg)
 162                 free(send_msg);
 163         if(received_msg)
 164                 free(received_msg);
 165
 166         return retval;
 167 }
 168
 169 void decrement_active_buffers(void *arg)
 170 {
 171         pthread_mutex_lock(&active_buffers_mutex);
 172         active_buffers--;
 173         pthread_mutex_unlock(&active_buffers_mutex);
 174 }
 175
 176 int create_dir_if_needed(char *dir)
 177 {
 178         int result;
 179         result = mkdir(dir, 0777);
 180         if(result == -1) {
 181                 if(errno != EEXIST) {
 182                         PERROR("mkdir");
 183                         return -1;
 184                 }
 185         }
 186
 187         return 0;
 188 }
 189
 190 int is_directory(const char *dir)
 191 {
 192         int result;
 193         struct stat st;
 194
 195         result = stat(dir, &st);
 196         if(result == -1) {
 197                 PERROR("stat");
 198                 return 0;
 199         }
 200
 201         if(!S_ISDIR(st.st_mode)) {
 202                 return 0;
 203         }
 204
 205         return 1;
 206 }
 207
 208 struct buffer_info *connect_buffer(pid_t pid, const char *bufname)
 209 {
 210         struct buffer_info *buf;
 211         char *send_msg;
 212         char *received_msg;
 213         int result;
 214         char *tmp;
 215         int fd;
 216         struct shmid_ds shmds;
 217
 218         buf = (struct buffer_info *) malloc(sizeof(struct buffer_info));
 219         if(buf == NULL) {
 220                 ERR("add_buffer: insufficient memory");
 221                 return NULL;
 222         }
 223
 224         buf->name = bufname;
 225         buf->pid = pid;
 226
 227         /* connect to app */
 228         result = ustcomm_connect_app(buf->pid, &buf->conn);
 229         if(result) {
 230                 WARN("unable to connect to process, it probably died before we were able to connect");
 231                 return NULL;
 232         }
 233
 234         /* get pidunique */
 235         asprintf(&send_msg, "get_pidunique");
 236         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 237         free(send_msg);
 238         if(result == -1) {
 239                 ERR("problem in ustcomm_send_request(get_pidunique)");
 240                 return NULL;
 241         }
 242
 243         result = sscanf(received_msg, "%lld", &buf->pidunique);
 244         if(result != 1) {
 245                 ERR("unable to parse response to get_pidunique");
 246                 return NULL;
 247         }
 248         free(received_msg);
 249         DBG("got pidunique %lld", buf->pidunique);
 250
 251         /* get shmid */
 252         asprintf(&send_msg, "get_shmid %s", buf->name);
 253         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 254         free(send_msg);
 255         if(result == -1) {
 256                 ERR("problem in ustcomm_send_request(get_shmid)");
 257                 return NULL;
 258         }
 259
 260         result = sscanf(received_msg, "%d %d", &buf->shmid, &buf->bufstruct_shmid);
 261         if(result != 2) {
 262                 ERR("unable to parse response to get_shmid");
 263                 return NULL;
 264         }
 265         free(received_msg);
 266         DBG("got shmids %d %d", buf->shmid, buf->bufstruct_shmid);
 267
 268         /* get n_subbufs */
 269         asprintf(&send_msg, "get_n_subbufs %s", buf->name);
 270         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 271         free(send_msg);
 272         if(result == -1) {
 273                 ERR("problem in ustcomm_send_request(g_n_subbufs)");
 274                 return NULL;
 275         }
 276
 277         result = sscanf(received_msg, "%d", &buf->n_subbufs);
 278         if(result != 1) {
 279                 ERR("unable to parse response to get_n_subbufs");
 280                 return NULL;
 281         }
 282         free(received_msg);
 283         DBG("got n_subbufs %d", buf->n_subbufs);
 284
 285         /* get subbuf size */
 286         asprintf(&send_msg, "get_subbuf_size %s", buf->name);
 287         ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 288         free(send_msg);
 289
 290         result = sscanf(received_msg, "%d", &buf->subbuf_size);
 291         if(result != 1) {
 292                 ERR("unable to parse response to get_subbuf_size");
 293                 return NULL;
 294         }
 295         free(received_msg);
 296         DBG("got subbuf_size %d", buf->subbuf_size);
 297
 298         /* attach memory */
 299         buf->mem = shmat(buf->shmid, NULL, 0);
 300         if(buf->mem == (void *) 0) {
 301                 PERROR("shmat");
 302                 return NULL;
 303         }
 304         DBG("successfully attached buffer memory");
 305
 306         buf->bufstruct_mem = shmat(buf->bufstruct_shmid, NULL, 0);
 307         if(buf->bufstruct_mem == (void *) 0) {
 308                 PERROR("shmat");
 309                 return NULL;
 310         }
 311         DBG("successfully attached buffer bufstruct memory");
 312
 313         /* obtain info on the memory segment */
 314         result = shmctl(buf->shmid, IPC_STAT, &shmds);
 315         if(result == -1) {
 316                 PERROR("shmctl");
 317                 return NULL;
 318         }
 319         buf->memlen = shmds.shm_segsz;
 320
 321         /* open file for output */
 322         if(!trace_path) {
 323                 /* Only create the directory if using the default path, because
 324                  * of the risk of typo when using trace path override. We don't
 325                  * want to risk creating plenty of useless directories in that case.
 326                  */
 327                 result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
 328                 if(result == -1) {
 329                         ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
 330                         return NULL;
 331                 }
 332
 333                 trace_path = USTD_DEFAULT_TRACE_PATH;
 334         }
 335
 336         asprintf(&tmp, "%s/%u_%lld", trace_path, buf->pid, buf->pidunique);
 337         result = create_dir_if_needed(tmp);
 338         if(result == -1) {
 339                 ERR("could not create directory %s", tmp);
 340                 free(tmp);
 341                 return NULL;
 342         }
 343         free(tmp);
 344
 345         asprintf(&tmp, "%s/%u_%lld/%s_0", trace_path, buf->pid, buf->pidunique, buf->name);
 346         result = fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC | O_EXCL, 00600);
 347         if(result == -1) {
 348                 PERROR("open");
 349                 ERR("failed opening trace file %s", tmp);
 350                 return NULL;
 351         }
 352         buf->file_fd = fd;
 353         free(tmp);
 354
 355         pthread_mutex_lock(&active_buffers_mutex);
 356         active_buffers++;
 357         pthread_mutex_unlock(&active_buffers_mutex);
 358
 359         return buf;
 360 }
 361
 362 int consumer_loop(struct buffer_info *buf)
 363 {
 364         int result;
 365
 366         pthread_cleanup_push(decrement_active_buffers, NULL);
 367
 368         for(;;) {
 369                 /* get the subbuffer */
 370                 result = get_subbuffer(buf);
 371                 if(result == -1) {
 372                         ERR("error getting subbuffer");
 373                         continue;
 374                 }
 375                 else if(result == GET_SUBBUF_DONE) {
 376                         /* this is done */
 377                         break;
 378                 }
 379                 else if(result == GET_SUBBUF_DIED) {
 380                         finish_consuming_dead_subbuffer(buf);
 381                         break;
 382                 }
 383
 384                 /* write data to file */
 385                 result = patient_write(buf->file_fd, buf->mem + (buf->consumed_old & (buf->n_subbufs * buf->subbuf_size-1)), buf->subbuf_size);
 386                 if(result == -1) {
 387                         PERROR("write");
 388                         /* FIXME: maybe drop this trace */
 389                 }
 390
 391                 /* put the subbuffer */
 392                 /* FIXME: we actually should unput the buffer before consuming... */
 393                 result = put_subbuffer(buf);
 394                 if(result == -1) {
 395                         ERR("unknown error putting subbuffer (channel=%s)", buf->name);
 396                         break;
 397                 }
 398                 else if(result == PUT_SUBBUF_PUSHED) {
 399                         ERR("Buffer overflow (channel=%s), reader pushed. This channel will not be usable passed this point.", buf->name);
 400                         break;
 401                 }
 402                 else if(result == PUT_SUBBUF_DIED) {
 403                         WARN("application died while putting subbuffer");
 404                         /* FIXME: probably need to skip the first subbuffer in finish_consuming_dead_subbuffer */
 405                         finish_consuming_dead_subbuffer(buf);
 406                         break;
 407                 }
 408                 else if(result == PUT_SUBBUF_OK) {
 409                 }
 410         }
 411
 412         DBG("thread for buffer %s is stopping", buf->name);
 413
 414         /* FIXME: destroy, unalloc... */
 415
 416         pthread_cleanup_pop(1);
 417
 418         return 0;
 419 }
 420
 421 void free_buffer(struct buffer_info *buf)
 422 {
 423 }
 424
 425 struct consumer_thread_args {
 426         pid_t pid;
 427         const char *bufname;
 428 };
 429
 430 void *consumer_thread(void *arg)
 431 {
 432         struct buffer_info *buf = (struct buffer_info *) arg;
 433         struct consumer_thread_args *args = (struct consumer_thread_args *) arg;
 434
 435         DBG("GOT ARGS: pid %d bufname %s", args->pid, args->bufname);
 436
 437         buf = connect_buffer(args->pid, args->bufname);
 438         if(buf == NULL) {
 439                 ERR("failed to connect to buffer");
 440                 goto end;
 441         }
 442
 443         consumer_loop(buf);
 444
 445         free_buffer(buf);
 446
 447         end:
 448         /* bufname is free'd in free_buffer() */
 449         free(args);
 450         return NULL;
 451 }
 452
 453 int start_consuming_buffer(pid_t pid, const char *bufname)
 454 {
 455         pthread_t thr;
 456         struct consumer_thread_args *args;
 457
 458         DBG("beginning of start_consuming_buffer: args: pid %d bufname %s", pid, bufname);
 459
 460         args = (struct consumer_thread_args *) malloc(sizeof(struct consumer_thread_args));
 461
 462         args->pid = pid;
 463         args->bufname = strdup(bufname);
 464         DBG("beginning2 of start_consuming_buffer: args: pid %d bufname %s", args->pid, args->bufname);
 465
 466         pthread_create(&thr, NULL, consumer_thread, args);
 467         DBG("end of start_consuming_buffer: args: pid %d bufname %s", args->pid, args->bufname);
 468
 469         return 0;
 470 }
 471
 472 void usage(void)
 473 {
 474         fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
 475                         "\t-h\t\tDisplay this usage.\n"
 476                         "\t-o DIR\t\tSpecify the directory where to output the traces.\n"
 477                         "\t-s PATH\t\tSpecify the path to use for the daemon socket.\n"
 478                         "\t-d\t\tStart as a daemon.\n"
 479                         "\t--pidfile FILE\tWrite the PID in this file (when using -d).\n");
 480 }
 481
 482 int parse_args(int argc, char **argv)
 483 {
 484         int c;
 485
 486         while (1) {
 487                 int option_index = 0;
 488                 static struct option long_options[] = {
 489                         {"pidfile", 1, 0, 'p'},
 490                         {"help", 0, 0, 'h'},
 491                         {"version", 0, 0, 'V'},
 492                         {0, 0, 0, 0}
 493                 };
 494
 495                 c = getopt_long(argc, argv, "hs:o:d", long_options, &option_index);
 496                 if (c == -1)
 497                         break;
 498
 499                 switch (c) {
 500                 case 0:
 501                         printf("option %s", long_options[option_index].name);
 502                         if (optarg)
 503                                 printf(" with arg %s", optarg);
 504                         printf("\n");
 505                         break;
 506                 case 's':
 507                         sock_path = optarg;
 508                         break;
 509                 case 'o':
 510                         trace_path = optarg;
 511                         if(!is_directory(trace_path)) {
 512                                 ERR("Not a valid directory. (%s)", trace_path);
 513                                 return -1;
 514                         }
 515                         break;
 516                 case 'd':
 517                         daemon_mode = 1;
 518                         break;
 519                 case 'p':
 520                         pidfile = strdup(optarg);
 521                         break;
 522                 case 'h':
 523                         usage();
 524                         exit(0);
 525                 case 'V':
 526                         printf("Version 0.0\n");
 527                         break;
 528
 529                 default:
 530                         /* unknown option or other error; error is
 531                         printed by getopt, just return */
 532                         return -1;
 533                 }
 534         }
 535
 536         return 0;
 537 }
 538
 539 void sigterm_handler(int sig)
 540 {
 541         terminate_req = 1;
 542 }
 543
 544 static int write_pidfile(const char *file_name, pid_t pid)
 545 {
 546         FILE *pidfp;
 547
 548         pidfp = fopen(file_name, "w");
 549         if(!pidfp) {
 550                 PERROR("fopen (%s)", pidfile);
 551                 WARN("killing child process");
 552                 return -1;
 553         }
 554
 555         fprintf(pidfp, "%d\n", pid);
 556
 557         fclose(pidfp);
 558
 559         return 0;
 560 }
 561
 562 int start_ustd(int fd)
 563 {
 564         struct ustcomm_ustd ustd;
 565         int result;
 566         sigset_t sigset;
 567         struct sigaction sa;
 568
 569         result = sigemptyset(&sigset);
 570         if(result == -1) {
 571                 PERROR("sigemptyset");
 572                 return 1;
 573         }
 574         sa.sa_handler = sigterm_handler;
 575         sa.sa_mask = sigset;
 576         sa.sa_flags = SA_RESTART;
 577         result = sigaction(SIGTERM, &sa, NULL);
 578         if(result == -1) {
 579                 PERROR("sigaction");
 580                 return 1;
 581         }
 582
 583         result = ustcomm_init_ustd(&ustd, sock_path);
 584         if(result == -1) {
 585                 ERR("failed to initialize socket");
 586                 return 1;
 587         }
 588
 589         /* setup handler for SIGPIPE */
 590         result = sigemptyset(&sigset);
 591         if(result == -1) {
 592                 PERROR("sigemptyset");
 593                 return 1;
 594         }
 595         result = sigaddset(&sigset, SIGPIPE);
 596         if(result == -1) {
 597                 PERROR("sigaddset");
 598                 return 1;
 599         }
 600         result = sigprocmask(SIG_BLOCK, &sigset, NULL);
 601         if(result == -1) {
 602                 PERROR("sigprocmask");
 603                 return 1;
 604         }
 605
 606         /* Write pidfile */
 607         if(pidfile) {
 608                 result = write_pidfile(pidfile, getpid());
 609                 if(result == -1) {
 610                         ERR("failed to write pidfile");
 611                         return 1;
 612                 }
 613         }
 614
 615         /* Notify parent that we are successfully started. */
 616         if(fd != -1) {
 617                 /* write any one character */
 618                 result = write(fd, "!", 1);
 619                 if(result == -1) {
 620                         PERROR("write");
 621                         return -1;
 622                 }
 623                 if(result != 1) {
 624                         ERR("Problem sending confirmation of daemon start to parent");
 625                         return -1;
 626                 }
 627                 result = close(fd);
 628                 if(result == -1) {
 629                         PERROR("close");
 630                 }
 631         }
 632
 633         /* app loop */
 634         for(;;) {
 635                 char *recvbuf;
 636
 637                 /* check for requests on our public socket */
 638                 result = ustcomm_ustd_recv_message(&ustd, &recvbuf, NULL, 100);
 639                 if(result == -1) {
 640                         ERR("error in ustcomm_ustd_recv_message");
 641                         goto loop_end;
 642                 }
 643                 if(result > 0) {
 644                         if(!strncmp(recvbuf, "collect", 7)) {
 645                                 pid_t pid;
 646                                 char *bufname;
 647                                 int result;
 648
 649                                 result = sscanf(recvbuf, "%*s %d %50as", &pid, &bufname);
 650                                 if(result != 2) {
 651                                         ERR("parsing error: %s", recvbuf);
 652                                         goto free_bufname;
 653                                 }
 654
 655                                 result = start_consuming_buffer(pid, bufname);
 656                                 if(result < 0) {
 657                                         ERR("error in add_buffer");
 658                                         goto free_bufname;
 659                                 }
 660
 661                                 free_bufname:
 662                                 free(bufname);
 663                         }
 664
 665                         free(recvbuf);
 666                 }
 667
 668                 loop_end:
 669
 670                 if(terminate_req) {
 671                         pthread_mutex_lock(&active_buffers_mutex);
 672                         if(active_buffers == 0) {
 673                                 pthread_mutex_unlock(&active_buffers_mutex);
 674                                 break;
 675                         }
 676                         pthread_mutex_unlock(&active_buffers_mutex);
 677                 }
 678         }
 679
 680         return 0;
 681 }
 682
 683 int start_ustd_daemon()
 684 {
 685         int result;
 686         int fd[2];
 687         pid_t child_pid;
 688
 689         result = pipe(fd);
 690
 691         result = child_pid = fork();
 692         if(result == -1) {
 693                 PERROR("fork");
 694                 return -1;
 695         }
 696         else if(result == 0) {
 697                 return start_ustd(fd[1]);
 698         }
 699         else {
 700                 char buf;
 701
 702                 result = read(fd[0], &buf, 1);
 703                 if(result == -1) {
 704                         PERROR("read");
 705                         return -1;
 706                 }
 707                 if(result != 1) {
 708                         ERR("did not receive valid confirmation that the daemon is started");
 709                         return -1;
 710                 }
 711
 712                 result = close(fd[0]);
 713                 if(result == -1) {
 714                         PERROR("close");
 715                 }
 716
 717                 DBG("The daemon is now successfully started");
 718         }
 719
 720         /* Wait for confirmation that the server is ready. */
 721
 722
 723         return 0;
 724 }
 725
 726 int main(int argc, char **argv)
 727 {
 728         int result;
 729
 730         result = parse_args(argc, argv);
 731         if(result == -1) {
 732                 exit(1);
 733         }
 734
 735         if(daemon_mode) {
 736                 result = start_ustd_daemon();
 737         }
 738         else {
 739                 result = start_ustd(-1);
 740         }
 741
 742         return result;
 743 }