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 (\"%s\")", received_msg);
 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", 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 write_current_subbuffer(struct buffer_info *buf)
 363 {
 364         int result;
 365
 366         void *subbuf_mem = buf->mem + (buf->consumed_old & (buf->n_subbufs * buf->subbuf_size-1));
 367
 368         size_t cur_sb_size = subbuffer_data_size(subbuf_mem);
 369
 370         result = patient_write(buf->file_fd, subbuf_mem, cur_sb_size);
 371         if(result == -1) {
 372                 PERROR("write");
 373                 /* FIXME: maybe drop this trace */
 374                 return 0;
 375         }
 376
 377         return 0;
 378 }
 379
 380 int consumer_loop(struct buffer_info *buf)
 381 {
 382         int result;
 383
 384         pthread_cleanup_push(decrement_active_buffers, NULL);
 385
 386         for(;;) {
 387                 /* get the subbuffer */
 388                 result = get_subbuffer(buf);
 389                 if(result == -1) {
 390                         ERR("error getting subbuffer");
 391                         continue;
 392                 }
 393                 else if(result == GET_SUBBUF_DONE) {
 394                         /* this is done */
 395                         break;
 396                 }
 397                 else if(result == GET_SUBBUF_DIED) {
 398                         finish_consuming_dead_subbuffer(buf);
 399                         break;
 400                 }
 401
 402                 /* write data to file */
 403                 write_current_subbuffer(buf);
 404                 /* FIXME: handle return value? */
 405
 406                 /* put the subbuffer */
 407                 /* FIXME: we actually should unput the buffer before consuming... */
 408                 result = put_subbuffer(buf);
 409                 if(result == -1) {
 410                         ERR("unknown error putting subbuffer (channel=%s)", buf->name);
 411                         break;
 412                 }
 413                 else if(result == PUT_SUBBUF_PUSHED) {
 414                         ERR("Buffer overflow (channel=%s), reader pushed. This channel will not be usable passed this point.", buf->name);
 415                         break;
 416                 }
 417                 else if(result == PUT_SUBBUF_DIED) {
 418                         WARN("application died while putting subbuffer");
 419                         /* FIXME: probably need to skip the first subbuffer in finish_consuming_dead_subbuffer */
 420                         finish_consuming_dead_subbuffer(buf);
 421                         break;
 422                 }
 423                 else if(result == PUT_SUBBUF_OK) {
 424                 }
 425         }
 426
 427         DBG("thread for buffer %s is stopping", buf->name);
 428
 429         /* FIXME: destroy, unalloc... */
 430
 431         pthread_cleanup_pop(1);
 432
 433         return 0;
 434 }
 435
 436 void free_buffer(struct buffer_info *buf)
 437 {
 438 }
 439
 440 struct consumer_thread_args {
 441         pid_t pid;
 442         const char *bufname;
 443 };
 444
 445 void *consumer_thread(void *arg)
 446 {
 447         struct buffer_info *buf = (struct buffer_info *) arg;
 448         struct consumer_thread_args *args = (struct consumer_thread_args *) arg;
 449
 450         DBG("GOT ARGS: pid %d bufname %s", args->pid, args->bufname);
 451
 452         buf = connect_buffer(args->pid, args->bufname);
 453         if(buf == NULL) {
 454                 ERR("failed to connect to buffer");
 455                 goto end;
 456         }
 457
 458         consumer_loop(buf);
 459
 460         free_buffer(buf);
 461
 462         end:
 463         /* bufname is free'd in free_buffer() */
 464         free(args);
 465         return NULL;
 466 }
 467
 468 int start_consuming_buffer(pid_t pid, const char *bufname)
 469 {
 470         pthread_t thr;
 471         struct consumer_thread_args *args;
 472
 473         DBG("beginning of start_consuming_buffer: args: pid %d bufname %s", pid, bufname);
 474
 475         args = (struct consumer_thread_args *) malloc(sizeof(struct consumer_thread_args));
 476
 477         args->pid = pid;
 478         args->bufname = strdup(bufname);
 479         DBG("beginning2 of start_consuming_buffer: args: pid %d bufname %s", args->pid, args->bufname);
 480
 481         pthread_create(&thr, NULL, consumer_thread, args);
 482         DBG("end of start_consuming_buffer: args: pid %d bufname %s", args->pid, args->bufname);
 483
 484         return 0;
 485 }
 486
 487 void usage(void)
 488 {
 489         fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
 490                         "\t-h\t\tDisplay this usage.\n"
 491                         "\t-o DIR\t\tSpecify the directory where to output the traces.\n"
 492                         "\t-s PATH\t\tSpecify the path to use for the daemon socket.\n"
 493                         "\t-d\t\tStart as a daemon.\n"
 494                         "\t--pidfile FILE\tWrite the PID in this file (when using -d).\n");
 495 }
 496
 497 int parse_args(int argc, char **argv)
 498 {
 499         int c;
 500
 501         while (1) {
 502                 int option_index = 0;
 503                 static struct option long_options[] = {
 504                         {"pidfile", 1, 0, 'p'},
 505                         {"help", 0, 0, 'h'},
 506                         {"version", 0, 0, 'V'},
 507                         {0, 0, 0, 0}
 508                 };
 509
 510                 c = getopt_long(argc, argv, "hs:o:d", long_options, &option_index);
 511                 if (c == -1)
 512                         break;
 513
 514                 switch (c) {
 515                 case 0:
 516                         printf("option %s", long_options[option_index].name);
 517                         if (optarg)
 518                                 printf(" with arg %s", optarg);
 519                         printf("\n");
 520                         break;
 521                 case 's':
 522                         sock_path = optarg;
 523                         break;
 524                 case 'o':
 525                         trace_path = optarg;
 526                         if(!is_directory(trace_path)) {
 527                                 ERR("Not a valid directory. (%s)", trace_path);
 528                                 return -1;
 529                         }
 530                         break;
 531                 case 'd':
 532                         daemon_mode = 1;
 533                         break;
 534                 case 'p':
 535                         pidfile = strdup(optarg);
 536                         break;
 537                 case 'h':
 538                         usage();
 539                         exit(0);
 540                 case 'V':
 541                         printf("Version 0.0\n");
 542                         break;
 543
 544                 default:
 545                         /* unknown option or other error; error is
 546                         printed by getopt, just return */
 547                         return -1;
 548                 }
 549         }
 550
 551         return 0;
 552 }
 553
 554 void sigterm_handler(int sig)
 555 {
 556         terminate_req = 1;
 557 }
 558
 559 static int write_pidfile(const char *file_name, pid_t pid)
 560 {
 561         FILE *pidfp;
 562
 563         pidfp = fopen(file_name, "w");
 564         if(!pidfp) {
 565                 PERROR("fopen (%s)", pidfile);
 566                 WARN("killing child process");
 567                 return -1;
 568         }
 569
 570         fprintf(pidfp, "%d\n", pid);
 571
 572         fclose(pidfp);
 573
 574         return 0;
 575 }
 576
 577 int start_ustd(int fd)
 578 {
 579         struct ustcomm_ustd ustd;
 580         int result;
 581         sigset_t sigset;
 582         struct sigaction sa;
 583
 584         result = sigemptyset(&sigset);
 585         if(result == -1) {
 586                 PERROR("sigemptyset");
 587                 return 1;
 588         }
 589         sa.sa_handler = sigterm_handler;
 590         sa.sa_mask = sigset;
 591         sa.sa_flags = SA_RESTART;
 592         result = sigaction(SIGTERM, &sa, NULL);
 593         if(result == -1) {
 594                 PERROR("sigaction");
 595                 return 1;
 596         }
 597
 598         result = ustcomm_init_ustd(&ustd, sock_path);
 599         if(result == -1) {
 600                 ERR("failed to initialize socket");
 601                 return 1;
 602         }
 603
 604         /* setup handler for SIGPIPE */
 605         result = sigemptyset(&sigset);
 606         if(result == -1) {
 607                 PERROR("sigemptyset");
 608                 return 1;
 609         }
 610         result = sigaddset(&sigset, SIGPIPE);
 611         if(result == -1) {
 612                 PERROR("sigaddset");
 613                 return 1;
 614         }
 615         result = sigprocmask(SIG_BLOCK, &sigset, NULL);
 616         if(result == -1) {
 617                 PERROR("sigprocmask");
 618                 return 1;
 619         }
 620
 621         /* Write pidfile */
 622         if(pidfile) {
 623                 result = write_pidfile(pidfile, getpid());
 624                 if(result == -1) {
 625                         ERR("failed to write pidfile");
 626                         return 1;
 627                 }
 628         }
 629
 630         /* Notify parent that we are successfully started. */
 631         if(fd != -1) {
 632                 /* write any one character */
 633                 result = write(fd, "!", 1);
 634                 if(result == -1) {
 635                         PERROR("write");
 636                         return -1;
 637                 }
 638                 if(result != 1) {
 639                         ERR("Problem sending confirmation of daemon start to parent");
 640                         return -1;
 641                 }
 642                 result = close(fd);
 643                 if(result == -1) {
 644                         PERROR("close");
 645                 }
 646         }
 647
 648         /* app loop */
 649         for(;;) {
 650                 char *recvbuf;
 651
 652                 /* check for requests on our public socket */
 653                 result = ustcomm_ustd_recv_message(&ustd, &recvbuf, NULL, 100);
 654                 if(result == -1) {
 655                         ERR("error in ustcomm_ustd_recv_message");
 656                         goto loop_end;
 657                 }
 658                 if(result > 0) {
 659                         if(!strncmp(recvbuf, "collect", 7)) {
 660                                 pid_t pid;
 661                                 char *bufname;
 662                                 int result;
 663
 664                                 result = sscanf(recvbuf, "%*s %d %50as", &pid, &bufname);
 665                                 if(result != 2) {
 666                                         ERR("parsing error: %s", recvbuf);
 667                                         goto free_bufname;
 668                                 }
 669
 670                                 result = start_consuming_buffer(pid, bufname);
 671                                 if(result < 0) {
 672                                         ERR("error in add_buffer");
 673                                         goto free_bufname;
 674                                 }
 675
 676                                 free_bufname:
 677                                 free(bufname);
 678                         }
 679
 680                         free(recvbuf);
 681                 }
 682
 683                 loop_end:
 684
 685                 if(terminate_req) {
 686                         pthread_mutex_lock(&active_buffers_mutex);
 687                         if(active_buffers == 0) {
 688                                 pthread_mutex_unlock(&active_buffers_mutex);
 689                                 break;
 690                         }
 691                         pthread_mutex_unlock(&active_buffers_mutex);
 692                 }
 693         }
 694
 695         return 0;
 696 }
 697
 698 int start_ustd_daemon()
 699 {
 700         int result;
 701         int fd[2];
 702         pid_t child_pid;
 703
 704         result = pipe(fd);
 705
 706         result = child_pid = fork();
 707         if(result == -1) {
 708                 PERROR("fork");
 709                 return -1;
 710         }
 711         else if(result == 0) {
 712                 return start_ustd(fd[1]);
 713         }
 714         else {
 715                 char buf;
 716
 717                 result = read(fd[0], &buf, 1);
 718                 if(result == -1) {
 719                         PERROR("read");
 720                         return -1;
 721                 }
 722                 if(result != 1) {
 723                         ERR("did not receive valid confirmation that the daemon is started");
 724                         return -1;
 725                 }
 726
 727                 result = close(fd[0]);
 728                 if(result == -1) {
 729                         PERROR("close");
 730                 }
 731
 732                 DBG("The daemon is now successfully started");
 733         }
 734
 735         /* Wait for confirmation that the server is ready. */
 736
 737
 738         return 0;
 739 }
 740
 741 int main(int argc, char **argv)
 742 {
 743         int result;
 744
 745         result = parse_args(argc, argv);
 746         if(result == -1) {
 747                 exit(1);
 748         }
 749
 750         if(daemon_mode) {
 751                 result = start_ustd_daemon();
 752         }
 753         else {
 754                 result = start_ustd(-1);
 755         }
 756
 757         return result;
 758 }