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 "localerr.h"
  37 #include "ustcomm.h"
  38 #include "share.h"
  39
  40 /* return value: 0 = subbuffer is finished, it won't produce data anymore
  41  *               1 = got subbuffer successfully
  42  *               <0 = error
  43  */
  44
  45 #define GET_SUBBUF_OK 1
  46 #define GET_SUBBUF_DONE 0
  47 #define GET_SUBBUF_DIED 2
  48
  49 #define PUT_SUBBUF_OK 1
  50 #define PUT_SUBBUF_DIED 0
  51 #define PUT_SUBBUF_PUSHED 2
  52
  53 char *sock_path=NULL;
  54 char *trace_path=NULL;
  55 int daemon_mode = 0;
  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 'h':
 520                         usage();
 521                         exit(0);
 522                 case 'V':
 523                         printf("Version 0.0\n");
 524                         break;
 525
 526                 default:
 527                         /* unknown option or other error; error is
 528                         printed by getopt, just return */
 529                         return -1;
 530                 }
 531         }
 532
 533         return 0;
 534 }
 535
 536 void sigterm_handler(int sig)
 537 {
 538         terminate_req = 1;
 539 }
 540
 541 int start_ustd(int fd)
 542 {
 543         struct ustcomm_ustd ustd;
 544         int result;
 545         sigset_t sigset;
 546         struct sigaction sa;
 547
 548         result = sigemptyset(&sigset);
 549         if(result == -1) {
 550                 PERROR("sigemptyset");
 551                 return 1;
 552         }
 553         sa.sa_handler = sigterm_handler;
 554         sa.sa_mask = sigset;
 555         sa.sa_flags = SA_RESTART;
 556         result = sigaction(SIGTERM, &sa, NULL);
 557         if(result == -1) {
 558                 PERROR("sigaction");
 559                 return 1;
 560         }
 561
 562         result = ustcomm_init_ustd(&ustd, sock_path);
 563         if(result == -1) {
 564                 ERR("failed to initialize socket");
 565                 return 1;
 566         }
 567
 568         /* setup handler for SIGPIPE */
 569         result = sigemptyset(&sigset);
 570         if(result == -1) {
 571                 PERROR("sigemptyset");
 572                 return 1;
 573         }
 574         result = sigaddset(&sigset, SIGPIPE);
 575         if(result == -1) {
 576                 PERROR("sigaddset");
 577                 return 1;
 578         }
 579         result = sigprocmask(SIG_BLOCK, &sigset, NULL);
 580         if(result == -1) {
 581                 PERROR("sigprocmask");
 582                 return 1;
 583         }
 584
 585         /* Notify parent that we are successfully started. */
 586         if(fd != -1) {
 587                 /* write any one character */
 588                 result = write(fd, "!", 1);
 589                 if(result == -1) {
 590                         PERROR("write");
 591                         return -1;
 592                 }
 593                 if(result != 1) {
 594                         ERR("Problem sending confirmation of daemon start to parent");
 595                         return -1;
 596                 }
 597                 result = close(fd);
 598                 if(result == -1) {
 599                         PERROR("close");
 600                 }
 601         }
 602
 603         /* app loop */
 604         for(;;) {
 605                 char *recvbuf;
 606
 607                 /* check for requests on our public socket */
 608                 result = ustcomm_ustd_recv_message(&ustd, &recvbuf, NULL, 100);
 609                 if(result == -1) {
 610                         ERR("error in ustcomm_ustd_recv_message");
 611                         goto loop_end;
 612                 }
 613                 if(result > 0) {
 614                         if(!strncmp(recvbuf, "collect", 7)) {
 615                                 pid_t pid;
 616                                 char *bufname;
 617                                 int result;
 618
 619                                 result = sscanf(recvbuf, "%*s %d %50as", &pid, &bufname);
 620                                 if(result != 2) {
 621                                         ERR("parsing error: %s", recvbuf);
 622                                         goto free_bufname;
 623                                 }
 624
 625                                 result = start_consuming_buffer(pid, bufname);
 626                                 if(result < 0) {
 627                                         ERR("error in add_buffer");
 628                                         goto free_bufname;
 629                                 }
 630
 631                                 free_bufname:
 632                                 free(bufname);
 633                         }
 634
 635                         free(recvbuf);
 636                 }
 637
 638                 loop_end:
 639
 640                 if(terminate_req) {
 641                         pthread_mutex_lock(&active_buffers_mutex);
 642                         if(active_buffers == 0) {
 643                                 pthread_mutex_unlock(&active_buffers_mutex);
 644                                 break;
 645                         }
 646                         pthread_mutex_unlock(&active_buffers_mutex);
 647                 }
 648         }
 649
 650         return 0;
 651 }
 652
 653 int start_ustd_daemon()
 654 {
 655         int result;
 656         int fd[2];
 657
 658         result = pipe(fd);
 659
 660         result = fork();
 661         if(result == -1) {
 662                 PERROR("fork");
 663                 return -1;
 664         }
 665         else if(result == 0) {
 666                 return start_ustd(fd[1]);
 667         }
 668         else {
 669                 char buf;
 670
 671                 result = read(fd[0], &buf, 1);
 672                 if(result == -1) {
 673                         PERROR("read");
 674                         return -1;
 675                 }
 676                 if(result != 1) {
 677                         ERR("did not receive valid confirmation that the daemon is started");
 678                         return -1;
 679                 }
 680
 681                 result = close(fd[0]);
 682                 if(result == -1) {
 683                         PERROR("close");
 684                 }
 685
 686                 DBG("The daemon is now successfully started");
 687         }
 688
 689         /* Wait for confirmation that the server is ready. */
 690
 691
 692         return 0;
 693 }
 694
 695 int main(int argc, char **argv)
 696 {
 697         int result;
 698
 699         result = parse_args(argc, argv);
 700         if(result == -1) {
 701                 exit(1);
 702         }
 703
 704         if(daemon_mode) {
 705                 result = start_ustd_daemon();
 706         }
 707         else {
 708                 result = start_ustd(-1);
 709         }
 710
 711         return result;
 712 }