ustd/ustd.c

   1 /* Copyright (C) 2009  Pierre-Marc Fournier
   2  *               2010  Alexis Halle
   3  *
   4  * This library is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * This library is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with this library; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
  17  */
  18
  19 #define _GNU_SOURCE
  20
  21 #include <sys/types.h>
  22 #include <sys/stat.h>
  23 #include <sys/shm.h>
  24 #include <fcntl.h>
  25 #include <unistd.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 "ust/ustd.h"
  36 #include "usterr.h"
  37
  38 char *sock_path=NULL;
  39 char *trace_path=NULL;
  40 int daemon_mode = 0;
  41 char *pidfile = NULL;
  42
  43 struct libustd_instance *instance;
  44
  45 struct buffer_info_local {
  46         /* output file */
  47         int file_fd;
  48         /* the offset we must truncate to, to unput the last subbuffer */
  49         off_t previous_offset;
  50 };
  51
  52 static int write_pidfile(const char *file_name, pid_t pid)
  53 {
  54         FILE *pidfp;
  55
  56         pidfp = fopen(file_name, "w");
  57         if(!pidfp) {
  58                 PERROR("fopen (%s)", file_name);
  59                 WARN("killing child process");
  60                 return -1;
  61         }
  62
  63         fprintf(pidfp, "%d\n", pid);
  64
  65         fclose(pidfp);
  66
  67         return 0;
  68 }
  69
  70 int create_dir_if_needed(char *dir)
  71 {
  72         int result;
  73         result = mkdir(dir, 0777);
  74         if(result == -1) {
  75                 if(errno != EEXIST) {
  76                         PERROR("mkdir");
  77                         return -1;
  78                 }
  79         }
  80
  81         return 0;
  82 }
  83
  84 int unwrite_last_subbuffer(struct buffer_info *buf)
  85 {
  86         int result;
  87         struct buffer_info_local *buf_local = buf->user_data;
  88
  89         result = ftruncate(buf_local->file_fd, buf_local->previous_offset);
  90         if(result == -1) {
  91                 PERROR("ftruncate");
  92                 return -1;
  93         }
  94
  95         result = lseek(buf_local->file_fd, buf_local->previous_offset, SEEK_SET);
  96         if(result == (int)(off_t)-1) {
  97                 PERROR("lseek");
  98                 return -1;
  99         }
 100
 101         return 0;
 102 }
 103
 104 int write_current_subbuffer(struct buffer_info *buf)
 105 {
 106         int result;
 107         struct buffer_info_local *buf_local = buf->user_data;
 108
 109         void *subbuf_mem = buf->mem + (buf->consumed_old & (buf->n_subbufs * buf->subbuf_size-1));
 110
 111         size_t cur_sb_size = subbuffer_data_size(subbuf_mem);
 112
 113         off_t cur_offset = lseek(buf_local->file_fd, 0, SEEK_CUR);
 114         if(cur_offset == (off_t)-1) {
 115                 PERROR("lseek");
 116                 return -1;
 117         }
 118
 119         buf_local->previous_offset = cur_offset;
 120         DBG("previous_offset: %ld", cur_offset);
 121
 122         result = patient_write(buf_local->file_fd, subbuf_mem, cur_sb_size);
 123         if(result == -1) {
 124                 PERROR("write");
 125                 return -1;
 126         }
 127
 128         return 0;
 129 }
 130
 131 int on_read_subbuffer(struct libustd_callbacks *data, struct buffer_info *buf)
 132 {
 133         return write_current_subbuffer(buf);
 134 }
 135
 136 int on_read_partial_subbuffer(struct libustd_callbacks *data, struct buffer_info *buf,
 137                                 long subbuf_index, unsigned long valid_length)
 138 {
 139         struct buffer_info_local *buf_local = buf->user_data;
 140         char *tmp;
 141         int result;
 142         unsigned long pad_size;
 143
 144         result = patient_write(buf_local->file_fd, buf->mem + subbuf_index * buf->subbuf_size, valid_length);
 145         if(result == -1) {
 146                 ERR("Error writing to buffer file");
 147                 return;
 148         }
 149
 150         /* pad with empty bytes */
 151         pad_size = PAGE_ALIGN(valid_length)-valid_length;
 152         if(pad_size) {
 153                 tmp = zmalloc(pad_size);
 154                 result = patient_write(buf_local->file_fd, tmp, pad_size);
 155                 if(result == -1) {
 156                         ERR("Error writing to buffer file");
 157                         return;
 158                 }
 159                 free(tmp);
 160         }
 161
 162 }
 163
 164 int on_open_buffer(struct libustd_callbacks *data, struct buffer_info *buf)
 165 {
 166         char *tmp;
 167         int result;
 168         int fd;
 169         struct buffer_info_local *buf_local =
 170                 zmalloc(sizeof(struct buffer_info_local));
 171
 172         if(!buf_local) {
 173                 ERR("could not allocate buffer_info_local struct");
 174                 return 1;
 175         }
 176
 177         buf->user_data = buf_local;
 178
 179         /* open file for output */
 180         if(!trace_path) {
 181                 /* Only create the directory if using the default path, because
 182                  * of the risk of typo when using trace path override. We don't
 183                  * want to risk creating plenty of useless directories in that case.
 184                  */
 185                 result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
 186                 if(result == -1) {
 187                         ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
 188                         return 1;
 189                 }
 190
 191                 trace_path = USTD_DEFAULT_TRACE_PATH;
 192         }
 193
 194         asprintf(&tmp, "%s/%u_%lld", trace_path, buf->pid, buf->pidunique);
 195         result = create_dir_if_needed(tmp);
 196         if(result == -1) {
 197                 ERR("could not create directory %s", tmp);
 198                 free(tmp);
 199                 return 1;
 200         }
 201         free(tmp);
 202
 203         asprintf(&tmp, "%s/%u_%lld/%s", trace_path, buf->pid, buf->pidunique, buf->name);
 204         result = fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC | O_EXCL, 00600);
 205         if(result == -1) {
 206                 PERROR("open");
 207                 ERR("failed opening trace file %s", tmp);
 208                 return 1;
 209         }
 210         buf_local->file_fd = fd;
 211         free(tmp);
 212
 213         return 0;
 214 }
 215
 216 int on_close_buffer(struct libustd_callbacks *data, struct buffer_info *buf)
 217 {
 218         struct buffer_info_local *buf_local = buf->user_data;
 219         int result = close(buf_local->file_fd);
 220         free(buf_local);
 221         if(result == -1) {
 222                 PERROR("close");
 223         }
 224         return 0;
 225 }
 226
 227 int on_put_error(struct libustd_callbacks *data, struct buffer_info *buf)
 228 {
 229         unwrite_last_subbuffer(buf);
 230 }
 231
 232 struct libustd_callbacks *new_callbacks()
 233 {
 234         struct libustd_callbacks *callbacks =
 235                 zmalloc(sizeof(struct libustd_callbacks));
 236
 237         if(!callbacks)
 238                 return NULL;
 239
 240         callbacks->on_open_buffer = on_open_buffer;
 241         callbacks->on_close_buffer = on_close_buffer;
 242         callbacks->on_read_subbuffer = on_read_subbuffer;
 243         callbacks->on_read_partial_subbuffer = on_read_partial_subbuffer;
 244         callbacks->on_put_error = on_put_error;
 245         callbacks->on_new_thread = NULL;
 246         callbacks->on_close_thread = NULL;
 247         callbacks->on_trace_end = NULL;
 248
 249         return callbacks;
 250
 251 }
 252
 253 int is_directory(const char *dir)
 254 {
 255         int result;
 256         struct stat st;
 257
 258         result = stat(dir, &st);
 259         if(result == -1) {
 260                 PERROR("stat");
 261                 return 0;
 262         }
 263
 264         if(!S_ISDIR(st.st_mode)) {
 265                 return 0;
 266         }
 267
 268         return 1;
 269 }
 270
 271 void usage(void)
 272 {
 273         fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
 274                         "\t-h\t\tDisplay this usage.\n"
 275                         "\t-o DIR\t\tSpecify the directory where to output the traces.\n"
 276                         "\t-s PATH\t\tSpecify the path to use for the daemon socket.\n"
 277                         "\t-d\t\tStart as a daemon.\n"
 278                         "\t--pidfile FILE\tWrite the PID in this file (when using -d).\n");
 279 }
 280
 281 int parse_args(int argc, char **argv)
 282 {
 283         int c;
 284
 285         while (1) {
 286                 int option_index = 0;
 287                 static struct option long_options[] = {
 288                         {"pidfile", 1, 0, 'p'},
 289                         {"help", 0, 0, 'h'},
 290                         {"version", 0, 0, 'V'},
 291                         {0, 0, 0, 0}
 292                 };
 293
 294                 c = getopt_long(argc, argv, "hs:o:d", long_options, &option_index);
 295                 if (c == -1)
 296                         break;
 297
 298                 switch (c) {
 299                 case 0:
 300                         printf("option %s", long_options[option_index].name);
 301                         if (optarg)
 302                                 printf(" with arg %s", optarg);
 303                         printf("\n");
 304                         break;
 305                 case 's':
 306                         sock_path = optarg;
 307                         break;
 308                 case 'o':
 309                         trace_path = optarg;
 310                         if(!is_directory(trace_path)) {
 311                                 ERR("Not a valid directory. (%s)", trace_path);
 312                                 return -1;
 313                         }
 314                         break;
 315                 case 'd':
 316                         daemon_mode = 1;
 317                         break;
 318                 case 'p':
 319                         pidfile = strdup(optarg);
 320                         break;
 321                 case 'h':
 322                         usage();
 323                         exit(0);
 324                 case 'V':
 325                         printf("Version 0.0\n");
 326                         break;
 327
 328                 default:
 329                         /* unknown option or other error; error is
 330                         printed by getopt, just return */
 331                         return -1;
 332                 }
 333         }
 334
 335         return 0;
 336 }
 337
 338 void sigterm_handler(int sig)
 339 {
 340         libustd_stop_instance(instance, 0);
 341 }
 342
 343 int start_ustd(int fd)
 344 {
 345         int result;
 346         sigset_t sigset;
 347         struct sigaction sa;
 348
 349         struct libustd_callbacks *callbacks = new_callbacks();
 350         if(!callbacks) {
 351                 PERROR("new_callbacks");
 352                 return 1;
 353         }
 354
 355         result = sigemptyset(&sigset);
 356         if(result == -1) {
 357                 PERROR("sigemptyset");
 358                 return 1;
 359         }
 360         sa.sa_handler = sigterm_handler;
 361         sa.sa_mask = sigset;
 362         sa.sa_flags = 0;
 363         result = sigaction(SIGTERM, &sa, NULL);
 364         if(result == -1) {
 365                 PERROR("sigaction");
 366                 return 1;
 367         }
 368         result = sigaction(SIGINT, &sa, NULL);
 369         if(result == -1) {
 370                 PERROR("sigaction");
 371                 return 1;
 372         }
 373
 374         instance = libustd_new_instance(callbacks, sock_path);
 375         if(!instance) {
 376                 ERR("failed to create libustd instance");
 377                 return 1;
 378         }
 379
 380         result = libustd_init_instance(instance);
 381         if(result) {
 382                 ERR("failed to initialize libustd instance");
 383                 return 1;
 384         }
 385
 386         /* setup handler for SIGPIPE */
 387         result = sigemptyset(&sigset);
 388         if(result == -1) {
 389                 PERROR("sigemptyset");
 390                 return 1;
 391         }
 392         result = sigaddset(&sigset, SIGPIPE);
 393         if(result == -1) {
 394                 PERROR("sigaddset");
 395                 return 1;
 396         }
 397         result = sigprocmask(SIG_BLOCK, &sigset, NULL);
 398         if(result == -1) {
 399                 PERROR("sigprocmask");
 400                 return 1;
 401         }
 402
 403         /* Write pidfile */
 404         if(pidfile) {
 405                 result = write_pidfile(pidfile, getpid());
 406                 if(result == -1) {
 407                         ERR("failed to write pidfile");
 408                         return 1;
 409                 }
 410         }
 411
 412         /* Notify parent that we are successfully started. */
 413         if(fd != -1) {
 414                 /* write any one character */
 415                 result = write(fd, "!", 1);
 416                 if(result == -1) {
 417                         PERROR("write");
 418                         return -1;
 419                 }
 420                 if(result != 1) {
 421                         ERR("Problem sending confirmation of daemon start to parent");
 422                         return -1;
 423                 }
 424                 result = close(fd);
 425                 if(result == -1) {
 426                         PERROR("close");
 427                 }
 428         }
 429
 430         libustd_start_instance(instance);
 431
 432         free(callbacks);
 433
 434         return 0;
 435 }
 436
 437 int start_ustd_daemon()
 438 {
 439         int result;
 440         int fd[2];
 441         pid_t child_pid;
 442
 443         result = pipe(fd);
 444
 445         result = child_pid = fork();
 446         if(result == -1) {
 447                 PERROR("fork");
 448                 return -1;
 449         }
 450         else if(result == 0) {
 451                 return start_ustd(fd[1]);
 452         }
 453         else {
 454                 char buf;
 455
 456                 result = read(fd[0], &buf, 1);
 457                 if(result == -1) {
 458                         PERROR("read");
 459                         return -1;
 460                 }
 461                 if(result != 1) {
 462                         ERR("did not receive valid confirmation that the daemon is started");
 463                         return -1;
 464                 }
 465
 466                 result = close(fd[0]);
 467                 if(result == -1) {
 468                         PERROR("close");
 469                 }
 470
 471                 DBG("The daemon is now successfully started");
 472         }
 473
 474         /* Wait for confirmation that the server is ready. */
 475
 476
 477         return 0;
 478 }
 479
 480 int main(int argc, char **argv)
 481 {
 482         int result;
 483
 484         result = parse_args(argc, argv);
 485         if(result == -1) {
 486                 exit(1);
 487         }
 488
 489         if(daemon_mode) {
 490                 result = start_ustd_daemon();
 491         }
 492         else {
 493                 result = start_ustd(-1);
 494         }
 495
 496         return result;
 497 }