ustd: add system to elegantly end daemon on SIGTERM, after waiting that all traces...
[ust.git] / ustd / ustd.c
index 5b45b6c812b94263d92ce81b8f4a8a06a81b70f3..fdf8d13f8d2f1fe418b1604a1b6baad3a1114fd1 100644 (file)
@@ -18,6 +18,7 @@
 #define _GNU_SOURCE
 
 #include <sys/types.h>
+#include <sys/stat.h>
 #include <sys/shm.h>
 #include <fcntl.h>
 #include <unistd.h>
@@ -29,6 +30,7 @@
 #include <string.h>
 #include <errno.h>
 #include <assert.h>
+#include <getopt.h>
 
 #include "ustd.h"
 #include "localerr.h"
 #define PUT_SUBBUF_DIED 0
 #define PUT_SUBBUF_PUSHED 2
 
+char *sock_path=NULL;
+char *trace_path=NULL;
+
+/* Number of active buffers and the mutex to protect it. */
+int active_buffers = 0;
+pthread_mutex_t active_buffers_mutex = PTHREAD_MUTEX_INITIALIZER;
+/* Whether a request to end the program was received. */
+sig_atomic_t terminate_req = 0;
+
 int test_sigpipe(void)
 {
        sigset_t sigset;
@@ -191,11 +202,20 @@ ssize_t patient_write(int fd, const void *buf, size_t count)
        return bufc-(const char *)buf;
 }
 
+void decrement_active_buffers(void *arg)
+{
+       pthread_mutex_lock(&active_buffers_mutex);
+       active_buffers--;
+       pthread_mutex_unlock(&active_buffers_mutex);
+}
+
 void *consumer_thread(void *arg)
 {
        struct buffer_info *buf = (struct buffer_info *) arg;
        int result;
 
+       pthread_cleanup_push(decrement_active_buffers, NULL);
+
        for(;;) {
                /* get the subbuffer */
                result = get_subbuffer(buf);
@@ -242,6 +262,8 @@ void *consumer_thread(void *arg)
 
        /* FIXME: destroy, unalloc... */
 
+       pthread_cleanup_pop(1);
+
        return NULL;
 }
 
@@ -259,6 +281,24 @@ int create_dir_if_needed(char *dir)
        return 0;
 }
 
+int is_directory(const char *dir)
+{
+       int result;
+       struct stat st;
+
+       result = stat(dir, &st);
+       if(result == -1) {
+               PERROR("stat");
+               return 0;
+       }
+
+       if(!S_ISDIR(st.st_mode)) {
+               return 0;
+       }
+
+       return 1;
+}
+
 int add_buffer(pid_t pid, char *bufname)
 {
        struct buffer_info *buf;
@@ -357,13 +397,21 @@ int add_buffer(pid_t pid, char *bufname)
        buf->memlen = shmds.shm_segsz;
 
        /* open file for output */
-       result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
-       if(result == -1) {
-               ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
-               return -1;
+       if(!trace_path) {
+               /* Only create the directory if using the default path, because
+                * of the risk of typo when using trace path override. We don't
+                * want to risk creating plenty of useless directories in that case.
+                */
+               result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
+               if(result == -1) {
+                       ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
+                       return -1;
+               }
+
+               trace_path = USTD_DEFAULT_TRACE_PATH;
        }
 
-       asprintf(&tmp, "%s/%u", USTD_DEFAULT_TRACE_PATH, buf->pid);
+       asprintf(&tmp, "%s/%u", trace_path, buf->pid);
        result = create_dir_if_needed(tmp);
        if(result == -1) {
                ERR("could not create directory %s", tmp);
@@ -372,7 +420,7 @@ int add_buffer(pid_t pid, char *bufname)
        }
        free(tmp);
 
-       asprintf(&tmp, "%s/%u/%s_0", USTD_DEFAULT_TRACE_PATH, buf->pid, buf->name);
+       asprintf(&tmp, "%s/%u/%s_0", trace_path, buf->pid, buf->name);
        result = fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC, 00600);
        if(result == -1) {
                PERROR("open");
@@ -382,23 +430,111 @@ int add_buffer(pid_t pid, char *bufname)
        buf->file_fd = fd;
        free(tmp);
 
+       pthread_mutex_lock(&active_buffers_mutex);
+       active_buffers++;
+       pthread_mutex_unlock(&active_buffers_mutex);
+
        pthread_create(&thr, NULL, consumer_thread, buf);
 
        return 0;
 }
 
+void usage(void)
+{
+       fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
+                       "\t-h\t\tDisplay this usage.\n"
+                       "\t-o DIR\t\tSpecify the directory where to output the traces.\n"
+                       "\t-s PATH\t\tSpecify the path to use for the daemon socket.\n");
+}
+
+int parse_args(int argc, char **argv)
+{
+       int c;
+
+       while (1) {
+               int option_index = 0;
+               static struct option long_options[] = {
+                       {"help", 0, 0, 'h'},
+                       {"version", 0, 0, 'V'},
+                       {0, 0, 0, 0}
+               };
+
+               c = getopt_long(argc, argv, "hs:o:", long_options, &option_index);
+               if (c == -1)
+                       break;
+
+               switch (c) {
+               case 0:
+                       printf("option %s", long_options[option_index].name);
+                       if (optarg)
+                               printf(" with arg %s", optarg);
+                       printf("\n");
+                       break;
+               case 's':
+                       sock_path = optarg;
+                       break;
+               case 'o':
+                       trace_path = optarg;
+                       if(!is_directory(trace_path)) {
+                               ERR("Not a valid directory. (%s)", trace_path);
+                               return -1;
+                       }
+                       break;
+               case 'h':
+                       usage();
+                       exit(0);
+               case 'V':
+                       printf("Version 0.0\n");
+                       break;
+
+               default:
+                       /* unknown option or other error; error is
+                       printed by getopt, just return */
+                       return -1;
+               }
+       }
+
+       return 0;
+}
+
+void sigterm_handler(int sig)
+{
+       terminate_req = 1;
+}
+
 int main(int argc, char **argv)
 {
        struct ustcomm_ustd ustd;
        int result;
        sigset_t sigset;
+       struct sigaction sa;
 
-       result = ustcomm_init_ustd(&ustd);
+       result = sigemptyset(&sigset);
+       if(result == -1) {
+               perror("sigemptyset");
+               return 1;
+       }
+       sa.sa_handler = sigterm_handler;
+       sa.sa_mask = sigset;
+       sa.sa_flags = SA_RESTART;
+       result = sigaction(SIGTERM, &sa, NULL);
+       if(result == -1) {
+               PERROR("sigaction");
+               return 1;
+       }
+
+       result = parse_args(argc, argv);
+       if(result == -1) {
+               exit(1);
+       }
+
+       result = ustcomm_init_ustd(&ustd, sock_path);
        if(result == -1) {
                ERR("failed to initialize socket");
                return 1;
        }
 
+       /* setup handler for SIGPIPE */
        result = sigemptyset(&sigset);
        if(result == -1) {
                perror("sigemptyset");
@@ -445,6 +581,15 @@ int main(int argc, char **argv)
 
                        free(recvbuf);
                }
+
+               if(terminate_req) {
+                       pthread_mutex_lock(&active_buffers_mutex);
+                       if(active_buffers == 0) {
+                               pthread_mutex_unlock(&active_buffers_mutex);
+                               break;
+                       }
+                       pthread_mutex_unlock(&active_buffers_mutex);
+               }
        }
 
        return 0;
This page took 0.03182 seconds and 4 git commands to generate.