3 files changed, 45 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) {