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
  56 /* Number of active buffers and the mutex to protect it. */
  57 int active_buffers = 0;
  58 pthread_mutex_t active_buffers_mutex = PTHREAD_MUTEX_INITIALIZER;
  59 /* Whether a request to end the program was received. */
  60 sig_atomic_t terminate_req = 0;
  61
  62 int get_subbuffer(struct buffer_info *buf)
  63 {
  64         char *send_msg=NULL;
  65         char *received_msg=NULL;
  66         char *rep_code=NULL;
  67         int retval;
  68         int result;
  69
  70         asprintf(&send_msg, "get_subbuffer %s", buf->name);
  71         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
  72         if((result == -1 && errno == EPIPE) || result == 0) {
  73                 DBG("app died while being traced");
  74                 retval = GET_SUBBUF_DIED;
  75                 goto end;
  76         }
  77         else if(result < 0) {
  78                 ERR("get_subbuffer: ustcomm_send_request failed");
  79                 retval = -1;
  80                 goto end;
  81         }
  82
  83         result = sscanf(received_msg, "%as %ld", &rep_code, &buf->consumed_old);
  84         if(result != 2 && result != 1) {
  85                 ERR("unable to parse response to get_subbuffer");
  86                 retval = -1;
  87                 goto end_rep;
  88         }
  89
  90         DBG("received msg is %s", received_msg);
  91
  92         if(!strcmp(rep_code, "OK")) {
  93                 DBG("got subbuffer %s", buf->name);
  94                 retval = GET_SUBBUF_OK;
  95         }
  96         else if(nth_token_is(received_msg, "END", 0) == 1) {
  97                 retval = GET_SUBBUF_DONE;
  98                 goto end_rep;
  99         }
 100         else {
 101                 DBG("error getting subbuffer %s", buf->name);
 102                 retval = -1;
 103         }
 104
 105         /* FIMXE: free correctly the stuff */
 106 end_rep:
 107         if(rep_code)
 108                 free(rep_code);
 109 end:
 110         if(send_msg)
 111                 free(send_msg);
 112         if(received_msg)
 113                 free(received_msg);
 114
 115         return retval;
 116 }
 117
 118 int put_subbuffer(struct buffer_info *buf)
 119 {
 120         char *send_msg=NULL;
 121         char *received_msg=NULL;
 122         char *rep_code=NULL;
 123         int retval;
 124         int result;
 125
 126         asprintf(&send_msg, "put_subbuffer %s %ld", buf->name, buf->consumed_old);
 127         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 128         if(result < 0 && errno == ECONNRESET) {
 129                 retval = PUT_SUBBUF_DIED;
 130                 goto end;
 131         }
 132         if(result < 0) {
 133                 ERR("put_subbuffer: send_message failed");
 134                 retval = -1;
 135                 goto end;
 136         }
 137
 138         result = sscanf(received_msg, "%as", &rep_code);
 139         if(result != 1) {
 140                 ERR("unable to parse response to put_subbuffer");
 141                 retval = -1;
 142                 goto end_rep;
 143         }
 144
 145         if(!strcmp(rep_code, "OK")) {
 146                 DBG("subbuffer put %s", buf->name);
 147                 retval = PUT_SUBBUF_OK;
 148         }
 149         else {
 150                 DBG("put_subbuffer: received error, we were pushed");
 151                 retval = PUT_SUBBUF_PUSHED;
 152                 goto end_rep;
 153         }
 154
 155 end_rep:
 156         if(rep_code)
 157                 free(rep_code);
 158
 159 end:
 160         if(send_msg)
 161                 free(send_msg);
 162         if(received_msg)
 163                 free(received_msg);
 164
 165         return retval;
 166 }
 167
 168 void decrement_active_buffers(void *arg)
 169 {
 170         pthread_mutex_lock(&active_buffers_mutex);
 171         active_buffers--;
 172         pthread_mutex_unlock(&active_buffers_mutex);
 173 }
 174
 175 int create_dir_if_needed(char *dir)
 176 {
 177         int result;
 178         result = mkdir(dir, 0777);
 179         if(result == -1) {
 180                 if(errno != EEXIST) {
 181                         PERROR("mkdir");
 182                         return -1;
 183                 }
 184         }
 185
 186         return 0;
 187 }
 188
 189 int is_directory(const char *dir)
 190 {
 191         int result;
 192         struct stat st;
 193
 194         result = stat(dir, &st);
 195         if(result == -1) {
 196                 PERROR("stat");
 197                 return 0;
 198         }
 199
 200         if(!S_ISDIR(st.st_mode)) {
 201                 return 0;
 202         }
 203
 204         return 1;
 205 }
 206
 207 struct buffer_info *connect_buffer(pid_t pid, const char *bufname)
 208 {
 209         struct buffer_info *buf;
 210         char *send_msg;
 211         char *received_msg;
 212         int result;
 213         char *tmp;
 214         int fd;
 215         struct shmid_ds shmds;
 216
 217         buf = (struct buffer_info *) malloc(sizeof(struct buffer_info));
 218         if(buf == NULL) {
 219                 ERR("add_buffer: insufficient memory");
 220                 return NULL;
 221         }
 222
 223         buf->name = bufname;
 224         buf->pid = pid;
 225
 226         /* connect to app */
 227         result = ustcomm_connect_app(buf->pid, &buf->conn);
 228         if(result) {
 229                 WARN("unable to connect to process, it probably died before we were able to connect");
 230                 return NULL;
 231         }
 232
 233         /* get pidunique */
 234         asprintf(&send_msg, "get_pidunique");
 235         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 236         free(send_msg);
 237         if(result == -1) {
 238                 ERR("problem in ustcomm_send_request(get_pidunique)");
 239                 return NULL;
 240         }
 241
 242         result = sscanf(received_msg, "%lld", &buf->pidunique);
 243         if(result != 1) {
 244                 ERR("unable to parse response to get_pidunique");
 245                 return NULL;
 246         }
 247         free(received_msg);
 248         DBG("got pidunique %lld", buf->pidunique);
 249
 250         /* get shmid */
 251         asprintf(&send_msg, "get_shmid %s", buf->name);
 252         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 253         free(send_msg);
 254         if(result == -1) {
 255                 ERR("problem in ustcomm_send_request(get_shmid)");
 256                 return NULL;
 257         }
 258
 259         result = sscanf(received_msg, "%d %d", &buf->shmid, &buf->bufstruct_shmid);
 260         if(result != 2) {
 261                 ERR("unable to parse response to get_shmid");
 262                 return NULL;
 263         }
 264         free(received_msg);
 265         DBG("got shmids %d %d", buf->shmid, buf->bufstruct_shmid);
 266
 267         /* get n_subbufs */
 268         asprintf(&send_msg, "get_n_subbufs %s", buf->name);
 269         result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 270         free(send_msg);
 271         if(result == -1) {
 272                 ERR("problem in ustcomm_send_request(g_n_subbufs)");
 273                 return NULL;
 274         }
 275
 276         result = sscanf(received_msg, "%d", &buf->n_subbufs);
 277         if(result != 1) {
 278                 ERR("unable to parse response to get_n_subbufs");
 279                 return NULL;
 280         }
 281         free(received_msg);
 282         DBG("got n_subbufs %d", buf->n_subbufs);
 283
 284         /* get subbuf size */
 285         asprintf(&send_msg, "get_subbuf_size %s", buf->name);
 286         ustcomm_send_request(&buf->conn, send_msg, &received_msg);
 287         free(send_msg);
 288
 289         result = sscanf(received_msg, "%d", &buf->subbuf_size);
 290         if(result != 1) {
 291                 ERR("unable to parse response to get_subbuf_size");
 292                 return NULL;
 293         }
 294         free(received_msg);
 295         DBG("got subbuf_size %d", buf->subbuf_size);
 296
 297         /* attach memory */
 298         buf->mem = shmat(buf->shmid, NULL, 0);
 299         if(buf->mem == (void *) 0) {
 300                 PERROR("shmat");
 301                 return NULL;
 302         }
 303         DBG("successfully attached buffer memory");
 304
 305         buf->bufstruct_mem = shmat(buf->bufstruct_shmid, NULL, 0);
 306         if(buf->bufstruct_mem == (void *) 0) {
 307                 PERROR("shmat");
 308                 return NULL;
 309         }
 310         DBG("successfully attached buffer bufstruct memory");
 311
 312         /* obtain info on the memory segment */
 313         result = shmctl(buf->shmid, IPC_STAT, &shmds);
 314         if(result == -1) {
 315                 PERROR("shmctl");
 316                 return NULL;
 317         }
 318         buf->memlen = shmds.shm_segsz;
 319
 320         /* open file for output */
 321         if(!trace_path) {
 322                 /* Only create the directory if using the default path, because
 323                  * of the risk of typo when using trace path override. We don't
 324                  * want to risk creating plenty of useless directories in that case.
 325                  */
 326                 result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
 327                 if(result == -1) {
 328                         ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
 329                         return NULL;
 330                 }
 331
 332                 trace_path = USTD_DEFAULT_TRACE_PATH;
 333         }
 334
 335         asprintf(&tmp, "%s/%u_%lld", trace_path, buf->pid, buf->pidunique);
 336         result = create_dir_if_needed(tmp);
 337         if(result == -1) {
 338                 ERR("could not create directory %s", tmp);
 339                 free(tmp);
 340                 return NULL;
 341         }
 342         free(tmp);
 343
 344         asprintf(&tmp, "%s/%u_%lld/%s_0", trace_path, buf->pid, buf->pidunique, buf->name);
 345         result = fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC | O_EXCL, 00600);
 346         if(result == -1) {
 347                 PERROR("open");
 348                 ERR("failed opening trace file %s", tmp);
 349                 return NULL;
 350         }
 351         buf->file_fd = fd;
 352         free(tmp);
 353
 354         pthread_mutex_lock(&active_buffers_mutex);
 355         active_buffers++;
 356         pthread_mutex_unlock(&active_buffers_mutex);
 357
 358         return buf;
 359 }
 360
 361 int consumer_loop(struct buffer_info *buf)
 362 {
 363         int result;
 364
 365         pthread_cleanup_push(decrement_active_buffers, NULL);
 366
 367         for(;;) {
 368                 /* get the subbuffer */
 369                 result = get_subbuffer(buf);
 370                 if(result == -1) {
 371                         ERR("error getting subbuffer");
 372                         continue;
 373                 }
 374                 else if(result == GET_SUBBUF_DONE) {
 375                         /* this is done */
 376                         break;
 377                 }
 378                 else if(result == GET_SUBBUF_DIED) {
 379                         finish_consuming_dead_subbuffer(buf);
 380                         break;
 381                 }
 382
 383                 /* write data to file */
 384                 result = patient_write(buf->file_fd, buf->mem + (buf->consumed_old & (buf->n_subbufs * buf->subbuf_size-1)), buf->subbuf_size);
 385                 if(result == -1) {
 386                         PERROR("write");
 387                         /* FIXME: maybe drop this trace */
 388                 }
 389
 390                 /* put the subbuffer */
 391                 /* FIXME: we actually should unput the buffer before consuming... */
 392                 result = put_subbuffer(buf);
 393                 if(result == -1) {
 394                         ERR("unknown error putting subbuffer (channel=%s)", buf->name);
 395                         break;
 396                 }
 397                 else if(result == PUT_SUBBUF_PUSHED) {
 398                         ERR("Buffer overflow (channel=%s), reader pushed. This channel will not be usable passed this point.", buf->name);
 399                         break;
 400                 }
 401                 else if(result == PUT_SUBBUF_DIED) {
 402                         WARN("application died while putting subbuffer");
 403                         /* FIXME: probably need to skip the first subbuffer in finish_consuming_dead_subbuffer */
 404                         finish_consuming_dead_subbuffer(buf);
 405                         break;
 406                 }
 407                 else if(result == PUT_SUBBUF_OK) {
 408                 }
 409         }
 410
 411         DBG("thread for buffer %s is stopping", buf->name);
 412
 413         /* FIXME: destroy, unalloc... */
 414
 415         pthread_cleanup_pop(1);
 416
 417         return 0;
 418 }
 419
 420 void free_buffer(struct buffer_info *buf)
 421 {
 422 }
 423
 424 struct consumer_thread_args {
 425         pid_t pid;
 426         const char *bufname;
 427 };
 428
 429 void *consumer_thread(void *arg)
 430 {
 431         struct buffer_info *buf = (struct buffer_info *) arg;
 432         struct consumer_thread_args *args = (struct consumer_thread_args *) arg;
 433
 434         DBG("GOT ARGS: pid %d bufname %s", args->pid, args->bufname);
 435
 436         buf = connect_buffer(args->pid, args->bufname);
 437         if(buf == NULL) {
 438                 ERR("failed to connect to buffer");
 439                 goto end;
 440         }
 441
 442         consumer_loop(buf);
 443
 444         free_buffer(buf);
 445
 446         end:
 447         /* bufname is free'd in free_buffer() */
 448         free(args);
 449         return NULL;
 450 }
 451
 452 int start_consuming_buffer(pid_t pid, const char *bufname)
 453 {
 454         pthread_t thr;
 455         struct consumer_thread_args *args;
 456
 457         DBG("beginning of start_consuming_buffer: args: pid %d bufname %s", pid, bufname);
 458
 459         args = (struct consumer_thread_args *) malloc(sizeof(struct consumer_thread_args));
 460
 461         args->pid = pid;
 462         args->bufname = strdup(bufname);
 463         DBG("beginning2 of start_consuming_buffer: args: pid %d bufname %s", args->pid, args->bufname);
 464
 465         pthread_create(&thr, NULL, consumer_thread, args);
 466         DBG("end of start_consuming_buffer: args: pid %d bufname %s", args->pid, args->bufname);
 467
 468         return 0;
 469 }
 470
 471 void usage(void)
 472 {
 473         fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
 474                         "\t-h\t\tDisplay this usage.\n"
 475                         "\t-o DIR\t\tSpecify the directory where to output the traces.\n"
 476                         "\t-s PATH\t\tSpecify the path to use for the daemon socket.\n");
 477 }
 478
 479 int parse_args(int argc, char **argv)
 480 {
 481         int c;
 482
 483         while (1) {
 484                 int option_index = 0;
 485                 static struct option long_options[] = {
 486                         {"help", 0, 0, 'h'},
 487                         {"version", 0, 0, 'V'},
 488                         {0, 0, 0, 0}
 489                 };
 490
 491                 c = getopt_long(argc, argv, "hs:o:", long_options, &option_index);
 492                 if (c == -1)
 493                         break;
 494
 495                 switch (c) {
 496                 case 0:
 497                         printf("option %s", long_options[option_index].name);
 498                         if (optarg)
 499                                 printf(" with arg %s", optarg);
 500                         printf("\n");
 501                         break;
 502                 case 's':
 503                         sock_path = optarg;
 504                         break;
 505                 case 'o':
 506                         trace_path = optarg;
 507                         if(!is_directory(trace_path)) {
 508                                 ERR("Not a valid directory. (%s)", trace_path);
 509                                 return -1;
 510                         }
 511                         break;
 512                 case 'h':
 513                         usage();
 514                         exit(0);
 515                 case 'V':
 516                         printf("Version 0.0\n");
 517                         break;
 518
 519                 default:
 520                         /* unknown option or other error; error is
 521                         printed by getopt, just return */
 522                         return -1;
 523                 }
 524         }
 525
 526         return 0;
 527 }
 528
 529 void sigterm_handler(int sig)
 530 {
 531         terminate_req = 1;
 532 }
 533
 534 int main(int argc, char **argv)
 535 {
 536         struct ustcomm_ustd ustd;
 537         int result;
 538         sigset_t sigset;
 539         struct sigaction sa;
 540
 541         result = sigemptyset(&sigset);
 542         if(result == -1) {
 543                 PERROR("sigemptyset");
 544                 return 1;
 545         }
 546         sa.sa_handler = sigterm_handler;
 547         sa.sa_mask = sigset;
 548         sa.sa_flags = SA_RESTART;
 549         result = sigaction(SIGTERM, &sa, NULL);
 550         if(result == -1) {
 551                 PERROR("sigaction");
 552                 return 1;
 553         }
 554
 555         result = parse_args(argc, argv);
 556         if(result == -1) {
 557                 exit(1);
 558         }
 559
 560         result = ustcomm_init_ustd(&ustd, sock_path);
 561         if(result == -1) {
 562                 ERR("failed to initialize socket");
 563                 return 1;
 564         }
 565
 566         /* setup handler for SIGPIPE */
 567         result = sigemptyset(&sigset);
 568         if(result == -1) {
 569                 PERROR("sigemptyset");
 570                 return 1;
 571         }
 572         result = sigaddset(&sigset, SIGPIPE);
 573         if(result == -1) {
 574                 PERROR("sigaddset");
 575                 return 1;
 576         }
 577         result = sigprocmask(SIG_BLOCK, &sigset, NULL);
 578         if(result == -1) {
 579                 PERROR("sigprocmask");
 580                 return 1;
 581         }
 582
 583         /* app loop */
 584         for(;;) {
 585                 char *recvbuf;
 586
 587                 /* check for requests on our public socket */
 588                 result = ustcomm_ustd_recv_message(&ustd, &recvbuf, NULL, 100);
 589                 if(result == -1) {
 590                         ERR("error in ustcomm_ustd_recv_message");
 591                         goto loop_end;
 592                 }
 593                 if(result > 0) {
 594                         if(!strncmp(recvbuf, "collect", 7)) {
 595                                 pid_t pid;
 596                                 char *bufname;
 597                                 int result;
 598
 599                                 result = sscanf(recvbuf, "%*s %d %50as", &pid, &bufname);
 600                                 if(result != 2) {
 601                                         ERR("parsing error: %s", recvbuf);
 602                                         goto free_bufname;
 603                                 }
 604
 605                                 result = start_consuming_buffer(pid, bufname);
 606                                 if(result < 0) {
 607                                         ERR("error in add_buffer");
 608                                         goto free_bufname;
 609                                 }
 610
 611                                 free_bufname:
 612                                 free(bufname);
 613                         }
 614
 615                         free(recvbuf);
 616                 }
 617
 618                 loop_end:
 619
 620                 if(terminate_req) {
 621                         pthread_mutex_lock(&active_buffers_mutex);
 622                         if(active_buffers == 0) {
 623                                 pthread_mutex_unlock(&active_buffers_mutex);
 624                                 break;
 625                         }
 626                         pthread_mutex_unlock(&active_buffers_mutex);
 627                 }
 628         }
 629
 630         return 0;
 631 }