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