Add options

* -o sets the output directory
* cleaned up various stdout/stderr messages
* add ability to spawn a process and monitor it

1 files changed, 133 insertions, 33 deletions

diff --git a/mstat.c b/mstat.c
index 18e19b6..867a335 100644
--- a/mstat.c
+++ b/mstat.c
@@ -2,18 +2,30 @@
 #include <dirent.h>
 #include <signal.h>
 #include <time.h>
+#include <sys/wait.h>
 #include "common.h"
 
+#ifndef PATH_MAX
+#define PATH_MAX 4096
+#endif
 
 struct Option {
     /** Increased verbosity */
     unsigned char verbose;
     /** Overwrite existing file(s) */
     unsigned char clobber;
+    /** Program to execute */
+    char *program;
     /** PID to track */
     pid_t pid;
+    /** PID subprocess status */
+    int status;
     /** Output file handle to track */
     FILE *file;
+    /** Output root*/
+    char root[PATH_MAX];
+    /** Output filename */
+    char filename[PATH_MAX];
     /** Number of times per second mstat samples a pid */
     double sample_rate;
 } option;
@@ -25,12 +37,21 @@ struct Option {
  */
 static void handle_interrupt(int sig) {
     switch (sig) {
+        case SIGCHLD:
+            waitpid(option.pid, &option.status, WNOHANG|WUNTRACED);
+            if (WIFEXITED(option.status)) {
+                printf("pid %d returned %d\n", option.pid, WEXITSTATUS(option.status));
+            } else {
+                fprintf(stderr, "warning: pid %d is likely defunct\n", option.pid);
+            }
         case SIGUSR1:
             if (option.file) {
-                fprintf(stderr, "flushed handle: %p\n", option.file);
+                if (option.verbose)
+                    fprintf(stderr, "flushing %s\n", option.filename);
                 fflush(option.file);
             } else {
-                fprintf(stderr, "flush request ignored. no handle");
+                if (option.verbose)
+                    fprintf(stderr, "flush request ignored. no handle");
             }
             return;
         case 0:
@@ -39,6 +60,7 @@ static void handle_interrupt(int sig) {
             if (option.file) {
                 fflush(option.file);
                 fclose(option.file);
+                printf("data written: %s\n", option.filename);
             }
             exit(0);
         default:
@@ -55,11 +77,14 @@ static void usage(char *prog) {
     if (sep) {
         name = sep + 1;
     }
-    printf("usage: %s [OPTIONS] <PID>\n"
+    printf("usage: %s [OPTIONS] [-p PID] | {PROGRAM... ARGS}\n"
+           "-c        clobber 'PID#.mstat' if it exists\n"
            "-h        this help message\n"
-           "-c        clobber output file if it exists\n"
-           "-s        set sample rate (default: %0.2lf)\n"
-           "\n", name, option.sample_rate);
+           "-o DIR    path to output directory (must exist)\n"
+           "-p PID    process id to monitor\n"
+           "-s RATE   samples per second (default: %0.2lf)\n"
+           "-v        increased verbosity\n"
+           "", name, option.sample_rate);
 }
 
 /**
@@ -67,12 +92,12 @@ static void usage(char *prog) {
  * @param argc
  * @param argv
  */
-void parse_options(int argc, char *argv[]) {
+int parse_options(int argc, char *argv[]) {
     if (argc < 2) {
         usage(argv[0]);
         exit(1);
     }
-    for (int x = 0, i = 1; i < argc; i++) {
+    for (int i = 1; i < argc; i++) {
         if (strlen(argv[i]) > 1 && *argv[i] == '-') {
             char *arg = argv[i] + 1;
             if (!strcmp(arg, "h")) {
@@ -81,17 +106,24 @@ void parse_options(int argc, char *argv[]) {
             }
             if (!strcmp(arg, "v")) {
                 option.verbose = 1;
-            } else if (!strcmp(arg, "s")) {
-               option.sample_rate = strtod(argv[i+1], NULL);
-               i++;
             } else if (!strcmp(arg, "c")) {
                 option.clobber = 1;
+            } else if (!strcmp(arg, "o")) {
+                strncpy(option.root, argv[i+1], PATH_MAX - 1);
+                i++;
+            } else if (!strcmp(arg, "s")) {
+                option.sample_rate = strtod(argv[i+1], NULL);
+                i++;
+            } else if (!strcmp(arg, "p")) {
+                option.pid = (pid_t) strtol(argv[i+1], NULL, 10);
+                i++;
             }
         } else {
-            option.pid = (pid_t) strtol(argv[i], NULL, 10);
-            x++;
+            // Positional arguments begin here, return index
+            return i;
         }
     }
+    return -1;
 }
 
 /**
@@ -118,56 +150,116 @@ int smaps_rollup_usable(pid_t pid) {
 
 
 int main(int argc, char *argv[]) {
-    struct mstat_record_t record;
-    char path[PATH_MAX];
+    struct mstat_record_t record = {0};
+    int positional = 0;
 
     // Set default options
     option.sample_rate = 1;
     option.verbose = 0;
     option.clobber = 0;
     // Set options based on arguments
-    parse_options(argc, argv);
+    positional = parse_options(argc, argv);
+    if (!option.pid && positional < 0) {
+        fprintf(stderr, "missing: -p PID, or PROGRAM with arguments\n\n");
+        usage(argv[0]);
+        exit(1);
+    }
 
+    // Wait for our children
+    signal(SIGCHLD, handle_interrupt);
     // Allow user to flush the data stream with USR1
     signal(SIGUSR1, handle_interrupt);
     // Always attempt to exit cleanly
     signal(SIGINT, handle_interrupt);
     signal(SIGTERM, handle_interrupt);
 
-    // For each PID passed to the program, begin gathering stats
-    sprintf(path, "%d.mstat", option.pid);
+    // Set up output directory root and file path
+    snprintf(option.filename, PATH_MAX - 1, "%d.mstat", option.pid);
+    if (strlen(option.root)) {
+        // Strip trailing slash from path
+        if (strlen(option.root) > 1 && strrchr(option.root, '/')) {
+            option.root[strlen(option.root) - 1] = '\0';
+        }
 
-    if (pid_exists(option.pid) < 0) {
-        fprintf(stderr, "no pid %d\n", option.pid);
-        exit(1);
+        // Die if the output directory doesn't exist
+        if (access(option.root, X_OK) < 0) {
+            perror(option.root);
+            exit(1);
+        }
+
+        // Construct new filename
+        char tmppath[PATH_MAX];
+        snprintf(tmppath, PATH_MAX - 1, "%s/%s", option.root, option.filename);
+        strncpy(option.filename, tmppath, PATH_MAX - 1);
+    }
+
+    // Figure out what we are going to monitor.
+    // Will it be a user-defined PID or a new process?
+    if (option.pid) {
+        // User-defined PID
+        if (pid_exists(option.pid) < 0) {
+            fprintf(stderr, "no pid %d\n", option.pid);
+            exit(1);
+        }
+    } else {
+        // New process
+        pid_t p = fork();
+        if (p == -1) {
+            perror("fork");
+            exit(1);
+        }
+
+        if (p == 0) {
+            // "where" is the path to the program to execute
+            char where[PATH_MAX] = {0};
+            if (mstat_find_program(argv[positional], where)) {
+                perror(argv[positional]);
+                exit(1);
+            }
+
+            // Execute the requested program (with arguments)
+            if (execv(where, &argv[positional]) < 0) {
+                perror("execv");
+                exit(1);
+            }
+
+            if (waitpid(option.pid, &option.status, WNOHANG | WUNTRACED) < 0) {
+                perror("waitpid failed");
+                exit(1);
+            }
+        }
+        option.pid = p;
     }
 
+    // Verify /proc/PID/smaps_rollup is present
     if (smaps_rollup_usable(option.pid) < 0) {
         fprintf(stderr, "pid %d: %s\n", option.pid, strerror(errno));
         exit(1);
     }
 
-    if (access(path, F_OK) == 0) {
+    // Remove previous mstat data file if clobber is enabled
+    if (access(option.filename, F_OK) == 0) {
         if (option.clobber) {
-            remove(path);
-            fprintf(stderr, "%s clobbered\n", path);
+            remove(option.filename);
+            fprintf(stderr, "%s clobbered\n", option.filename);
         } else {
-            fprintf(stderr, "%s file already exists\n", path);
+            fprintf(stderr, "%s file already exists\n", option.filename);
             exit(1);
         }
     }
 
-    option.file = mstat_open(path);
+    // Initialize mstat data file
+    option.file = mstat_open(option.filename);
     if (!option.file) {
-        perror(path);
+        perror(option.filename);
         exit(1);
     }
 
+    // Write mstat data file header, if necessary
     if (mstat_check_header(option.file)) {
         mstat_write_header(option.file);
     }
 
-
     size_t i;
     struct timespec ts_start, ts_end;
 
@@ -175,7 +267,8 @@ int main(int argc, char *argv[]) {
     clock_gettime(CLOCK_MONOTONIC, &ts_start);
 
     // Begin sample loop
-    printf("PID: %d\nSamples per second: %.2lf\n(interrupt with ctrl-c...)\n", option.pid, option.sample_rate);
+    printf("PID: %d\nSamples per second: %.2lf\n(interrupt with ctrl-c...)\n",
+           option.pid, option.sample_rate);
     i = 0;
     while (1) {
         memset(&record, 0, sizeof(record));
@@ -187,22 +280,29 @@ int main(int argc, char *argv[]) {
 
         // Sample memory values
         if (mstat_attach(&record, record.pid) < 0) {
-            fprintf(stderr, "lost pid %d\n", option.pid);
+            if (positional < 0) {
+                // '-p' monitoring: let the user know when the PID disappears
+                fprintf(stderr, "pid: %d disappeared\n", option.pid);
+            }
             break;
         }
 
         if (option.verbose) {
-            printf("pid: %d, sample: %zu, elapsed: %lf, rss: %li\n", record.pid, i, record.timestamp, record.rss);
+            printf("pid: %d, sample: %zu, elapsed: %lf, rss: %li\n",
+                   record.pid, i, record.timestamp, record.rss);
         }
 
         if (mstat_write(option.file, &record) < 0) {
-            fprintf(stderr, "Unable to write record to mstat file for pid %d\n", option.pid);
+            fprintf(stderr, "Unable to write record to mstat file for pid %d: %s\n",
+                    option.pid, strerror(errno));
             break;
         }
+
+        // Perform n samples per second
         usleep((int) (1e6 / option.sample_rate));
         i++;
     }
 
     handle_interrupt(0);
-    return 0;
+    return option.status;
 }