Implement task timeout

* Add argument: --task-timeout=1[s,m,h]
* Timed out tasks are SIGKILL'd
* If killing a task fails, the entire program ends

7 files changed, 51 insertions, 0 deletions

diff --git a/src/cli/stasis/args.c b/src/cli/stasis/args.c
index fbda494..9410958 100644
--- a/src/cli/stasis/args.c
+++ b/src/cli/stasis/args.c
@@ -13,6 +13,7 @@ struct option long_options[] = {
     {"unbuffered", no_argument, 0, 'U'},
     {"update-base", no_argument, 0, OPT_ALWAYS_UPDATE_BASE},
     {"fail-fast", no_argument, 0, OPT_FAIL_FAST},
+    {"task-timeout", required_argument, 0, OPT_TASK_TIMEOUT},
     {"overwrite", no_argument, 0, OPT_OVERWRITE},
     {"no-docker", no_argument, 0, OPT_NO_DOCKER},
     {"no-artifactory", no_argument, 0, OPT_NO_ARTIFACTORY},
@@ -37,6 +38,7 @@ const char *long_options_help[] = {
     "Disable line buffering",
     "Update conda installation prior to STASIS environment creation",
     "On error, immediately terminate all tasks",
+    "Terminate task after timeout is reached (#s, #m, #h)",
     "Overwrite an existing release",
     "Do not build docker images",
     "Do not upload artifacts to Artifactory",
@@ -104,3 +106,29 @@ void usage(char *progname) {
         puts(output);
     }
 }
+
+int str_to_timeout(char *s) {
+    if (!s) {
+        return 0; // no timeout
+    }
+
+    char *scale = NULL;
+    int value = (int) strtol(s, &scale, 10);
+    if (scale) {
+        if (*scale == 's') {
+            value *= 1; // seconds, no-op
+        } else if (*scale == 'm') {
+            value *= 60; // minutes
+        } else if (*scale == 'h') {
+            value *= 3200; // hours
+        } else {
+            return STR_TO_TIMEOUT_INVALID_TIME_SCALE; // invalid time scale
+        }
+    }
+
+    if (value < 0) {
+        return STR_TO_TIMEOUT_NEGATIVE; // cannot be negative
+    }
+    return value;
+}
+
diff --git a/src/cli/stasis/include/args.h b/src/cli/stasis/include/args.h
index 850be4a..d75fe29 100644
--- a/src/cli/stasis/include/args.h
+++ b/src/cli/stasis/include/args.h
@@ -18,8 +18,13 @@
 #define OPT_NO_PARALLEL 1010
 #define OPT_POOL_STATUS_INTERVAL 1011
 #define OPT_NO_TASK_LOGGING 1012
+#define OPT_TASK_TIMEOUT 1013
 
 extern struct option long_options[];
 void usage(char *progname);
 
+#define STR_TO_TIMEOUT_NEGATIVE (-1)
+#define STR_TO_TIMEOUT_INVALID_TIME_SCALE (-2)
+int str_to_timeout(char *s);
+
 #endif //STASIS_ARGS_H
diff --git a/src/cli/stasis/stasis_main.c b/src/cli/stasis/stasis_main.c
index 967ecaf..44ee6d7 100644
--- a/src/cli/stasis/stasis_main.c
+++ b/src/cli/stasis/stasis_main.c
@@ -540,6 +540,18 @@ int main(int argc, char *argv[]) {
             case OPT_FAIL_FAST:
                 globals.parallel_fail_fast = true;
                 break;
+            case OPT_TASK_TIMEOUT:
+                globals.task_timeout = str_to_timeout(optarg);
+                if (globals.task_timeout < 0) {
+                    fprintf(stderr, "Invalid timeout: %s\n", optarg);
+                    if (globals.task_timeout == STR_TO_TIMEOUT_INVALID_TIME_SCALE) {
+                        fprintf(stderr, "Use format '#s' (seconds), '#m' (minutes), '#h' (hours)\n");
+                    } else if (globals.task_timeout == STR_TO_TIMEOUT_NEGATIVE) {
+                        fprintf(stderr, "Timeout cannot be negative\n");
+                    }
+                    exit(1);
+                }
+                break;
             case OPT_POOL_STATUS_INTERVAL:
                 globals.pool_status_interval = (int) strtol(optarg, NULL, 10);
                 if (globals.pool_status_interval < 1) {
diff --git a/src/lib/core/globals.c b/src/lib/core/globals.c
index a262d6c..834213b 100644
--- a/src/lib/core/globals.c
+++ b/src/lib/core/globals.c
@@ -44,6 +44,7 @@ struct STASIS_GLOBAL globals = {
         .enable_task_logging = true, ///< Toggle logging for multiprocess tasks
         .parallel_fail_fast = false, ///< Kill ALL multiprocessing tasks immediately on error
         .pool_status_interval = 30, ///< Report "Task is running"
+        .task_timeout = 0, ///< Time in seconds before task is terminated
 };
 
 void globals_free() {
diff --git a/src/lib/core/include/core.h b/src/lib/core/include/core.h
index e96e010..5a3fa85 100644
--- a/src/lib/core/include/core.h
+++ b/src/lib/core/include/core.h
@@ -51,6 +51,7 @@ struct STASIS_GLOBAL {
     char *tmpdir; //!< Path to temporary storage directory
     char *conda_install_prefix; //!< Path to install conda
     char *sysconfdir; //!< Path where STASIS reads its configuration files (mission directory, etc)
+    int task_timeout; ///< Time in seconds before task is terminated
     struct {
         char *tox_posargs;
         char *conda_reactivate;
diff --git a/src/lib/core/include/multiprocessing.h b/src/lib/core/include/multiprocessing.h
index 6477818..ab7b416 100644
--- a/src/lib/core/include/multiprocessing.h
+++ b/src/lib/core/include/multiprocessing.h
@@ -15,6 +15,7 @@ struct MultiProcessingTask {
     pid_t parent_pid; ///< Program PID (parent process)
     int status; ///< Child process exit status
     int signaled_by; ///< Last signal received, if any
+    int timeout; ///< Seconds to elapse before killing the process
     time_t _now; ///< Current time
     time_t _seconds; ///< Time elapsed since status interval (used by MultiprocessingPool.status_interval)
     time_t _startup; ///< Time elapsed since task started
diff --git a/src/lib/core/multiprocessing.c b/src/lib/core/multiprocessing.c
index 69719e8..ff4453c 100644
--- a/src/lib/core/multiprocessing.c
+++ b/src/lib/core/multiprocessing.c
@@ -167,6 +167,9 @@ struct MultiProcessingTask *mp_pool_task(struct MultiProcessingPool *pool, const
     memset(slot->cmd, 0, slot->cmd_len);
     strncpy(slot->cmd, cmd, slot->cmd_len);
 
+    // Set task timeout
+    slot->timeout = globals.task_timeout;
+
     return slot;
 }