rt: Add RUST_MAX_STACK env var with 8MB default

Closes #1489

9 files changed, 107 insertions, 58 deletions

diff --git a/src/rt/rust_env.cpp b/src/rt/rust_env.cpp
index b220c459c0b..bdec4c28435 100644
--- a/src/rt/rust_env.cpp
+++ b/src/rt/rust_env.cpp
@@ -8,6 +8,7 @@
 // The environment variables that the runtime knows about
 #define RUST_THREADS "RUST_THREADS"
 #define RUST_MIN_STACK "RUST_MIN_STACK"
+#define RUST_MAX_STACK "RUST_MAX_STACK"
 #define RUST_LOG "RUST_LOG"
 #define CHECK_CLAIMS "CHECK_CLAIMS"
 #define DETAILED_LEAKS "DETAILED_LEAKS"
@@ -69,15 +70,26 @@ get_num_threads()
 
 static size_t
 get_min_stk_size() {
-    char *stack_size = getenv(RUST_MIN_STACK);
-    if(stack_size) {
-        return strtol(stack_size, NULL, 0);
+    char *minsz = getenv(RUST_MIN_STACK);
+    if(minsz) {
+        return strtol(minsz, NULL, 0);
     }
     else {
         return 0x300;
     }
 }
 
+static size_t
+get_max_stk_size() {
+    char *maxsz = getenv(RUST_MAX_STACK);
+    if (maxsz) {
+	return strtol(maxsz, NULL, 0);
+    }
+    else {
+	return 1024*1024*8;
+    }
+}
+
 static char*
 copyenv(const char* name) {
     char *envvar = getenv(name);
@@ -99,6 +111,7 @@ load_env() {
 
     env->num_sched_threads = (size_t)get_num_threads();
     env->min_stack_size = get_min_stk_size();
+    env->max_stack_size = get_max_stk_size();
     env->logspec = copyenv(RUST_LOG);
     env->check_claims = getenv(CHECK_CLAIMS) != NULL;
     env->detailed_leaks = getenv(DETAILED_LEAKS) != NULL;
diff --git a/src/rt/rust_env.h b/src/rt/rust_env.h
index 6dc9cc120cd..eecdb7931b4 100644
--- a/src/rt/rust_env.h
+++ b/src/rt/rust_env.h
@@ -1,6 +1,7 @@
 struct rust_env {
     size_t num_sched_threads;
     size_t min_stack_size;
+    size_t max_stack_size;
     char* logspec;
     bool check_claims;
     bool detailed_leaks;
diff --git a/src/rt/rust_task.cpp b/src/rt/rust_task.cpp
index eb3040ebfae..4d136a59ae6 100644
--- a/src/rt/rust_task.cpp
+++ b/src/rt/rust_task.cpp
@@ -122,12 +122,6 @@ VALGRIND_STACK_DEREGISTER(stk->valgrind_id);
 }
 
 static void
-free_stk(rust_task *task, stk_seg *stk) {
-    LOGPTR(task->sched, "freeing stk segment", (uintptr_t)stk);
-    task->free(stk);
-}
-
-static void
 add_stack_canary(stk_seg *stk) {
     memcpy(stk->data, stack_canary, sizeof(stack_canary));
     assert(sizeof(stack_canary) == 16 && "Stack canary was not the expected size");
@@ -139,6 +133,21 @@ check_stack_canary(stk_seg *stk) {
       && "Somebody killed the canary");
 }
 
+// The amount of stack in a segment available to Rust code
+static size_t
+user_stack_size(stk_seg *stk) {
+    return (size_t)(stk->end
+                    - (uintptr_t)&stk->data[0]
+                    - RED_ZONE_SIZE);
+}
+
+static void
+free_stk(rust_task *task, stk_seg *stk) {
+    LOGPTR(task->sched, "freeing stk segment", (uintptr_t)stk);
+    task->total_stack_sz -= user_stack_size(stk);
+    task->free(stk);
+}
+
 static stk_seg*
 new_stk(rust_scheduler *sched, rust_task *task, size_t requested_sz)
 {
@@ -152,9 +161,7 @@ new_stk(rust_scheduler *sched, rust_task *task, size_t requested_sz)
 
     // Try to reuse an existing stack segment
     if (task->stk != NULL && task->stk->prev != NULL) {
-        size_t prev_sz = (size_t)(task->stk->prev->end
-                                  - (uintptr_t)&task->stk->prev->data[0]
-                                  - RED_ZONE_SIZE);
+        size_t prev_sz = user_stack_size(task->stk->prev);
         if (min_sz <= prev_sz && requested_sz <= prev_sz) {
             LOG(task, mem, "reusing existing stack");
             task->stk = task->stk->prev;
@@ -171,14 +178,17 @@ new_stk(rust_scheduler *sched, rust_task *task, size_t requested_sz)
     // The size of the current stack segment, excluding red zone
     size_t current_sz = 0;
     if (task->stk != NULL) {
-        current_sz = (size_t)(task->stk->end
-                              - (uintptr_t)&task->stk->data[0]
-                              - RED_ZONE_SIZE);
+        current_sz = user_stack_size(task->stk);
     }
     // The calculated size of the new stack, excluding red zone
     size_t rust_stk_sz = get_next_stk_size(sched, task, min_sz,
                                            current_sz, requested_sz);
 
+    if (task->total_stack_sz + rust_stk_sz > sched->env->max_stack_size) {
+        LOG_ERR(task, task, "task %" PRIxPTR " ran out of stack", task);
+        task->fail();
+    }
+
     size_t sz = sizeof(stk_seg) + rust_stk_sz + RED_ZONE_SIZE;
     stk_seg *stk = (stk_seg *)task->malloc(sz, "stack");
     LOGPTR(task->sched, "new stk", (uintptr_t)stk);
@@ -191,6 +201,7 @@ new_stk(rust_scheduler *sched, rust_task *task, size_t requested_sz)
 
     task->stk = stk;
     config_valgrind_stack(task->stk);
+    task->total_stack_sz += user_stack_size(stk);
     return stk;
 }
 
@@ -222,6 +233,7 @@ del_stk(rust_task *task, stk_seg *stk)
     unconfig_valgrind_stack(stk);
     if (delete_stack) {
         free_stk(task, stk);
+        A(task->sched, task->total_stack_sz == 0, "Stack size should be 0");
     }
 }
 
@@ -249,7 +261,8 @@ rust_task::rust_task(rust_scheduler *sched, rust_task_list *state,
     killed(false),
     propagate_failure(true),
     dynastack(this),
-    cc_counter(0)
+    cc_counter(0),
+    total_stack_sz(0)
 {
     LOGPTR(sched, "new task", (uintptr_t)this);
     DLOG(sched, task, "sizeof(task) = %d (0x%x)", sizeof *this, sizeof *this);
diff --git a/src/rt/rust_task.h b/src/rt/rust_task.h
index b47b62e2d14..6ae663d2066 100644
--- a/src/rt/rust_task.h
+++ b/src/rt/rust_task.h
@@ -126,6 +126,9 @@ rust_task : public kernel_owned<rust_task>, rust_cond
 
     debug::task_debug_info debug;
 
+    // The amount of stack we're using, excluding red zones
+    size_t total_stack_sz;
+
     // Only a pointer to 'name' is kept, so it must live as long as this task.
     rust_task(rust_scheduler *sched,
               rust_task_list *state,
diff --git a/src/test/run-fail/too-much-recursion.rs b/src/test/run-fail/too-much-recursion.rs
new file mode 100644
index 00000000000..db244bd0435
--- /dev/null
+++ b/src/test/run-fail/too-much-recursion.rs
@@ -0,0 +1,7 @@
+// error-pattern:ran out of stack
+
+// Test that the task fails after hiting the recursion limit
+
+fn main() {
+    main();
+}
\ No newline at end of file
diff --git a/src/test/run-pass/morestack2.rs b/src/test/run-pass/morestack2.rs
index ad46265fb93..a7e2c11cc01 100644
--- a/src/test/run-pass/morestack2.rs
+++ b/src/test/run-pass/morestack2.rs
@@ -10,5 +10,5 @@ fn getbig(i: int) -> int {
 }
 
 fn main() {
-    getbig(100000);
+    getbig(10000);
 }
\ No newline at end of file
diff --git a/src/test/run-pass/morestack3.rs b/src/test/run-pass/morestack3.rs
index 7f38c34347b..2ba7945894a 100644
--- a/src/test/run-pass/morestack3.rs
+++ b/src/test/run-pass/morestack3.rs
@@ -37,6 +37,6 @@ fn getbig(a0: int,
 }
 
 fn main() {
-    let a = 100000;
+    let a = 10000;
     getbig(a, a+1, a+2, a+3, a+4, a+5, a+6, a+7, a+8, a+9);
 }
\ No newline at end of file
diff --git a/src/test/run-pass/morestack4.rs b/src/test/run-pass/morestack4.rs
index 1883f0e3ea9..13200ee4ec9 100644
--- a/src/test/run-pass/morestack4.rs
+++ b/src/test/run-pass/morestack4.rs
@@ -53,45 +53,45 @@ fn getbig(i: biggy) {
 
 fn main() {
     getbig({
-        a00: 100000u64,
-        a01: 100000u64,
-        a02: 100000u64,
-        a03: 100000u64,
-        a04: 100000u64,
-        a05: 100000u64,
-        a06: 100000u64,
-        a07: 100000u64,
-        a08: 100000u64,
-        a09: 100000u64,
-        a10: 100000u64,
-        a11: 100000u64,
-        a12: 100000u64,
-        a13: 100000u64,
-        a14: 100000u64,
-        a15: 100000u64,
-        a16: 100000u64,
-        a17: 100000u64,
-        a18: 100000u64,
-        a19: 100000u64,
-        a20: 100000u64,
-        a21: 100000u64,
-        a22: 100000u64,
-        a23: 100000u64,
-        a24: 100000u64,
-        a25: 100000u64,
-        a26: 100000u64,
-        a27: 100000u64,
-        a28: 100000u64,
-        a29: 100000u64,
-        a30: 100000u64,
-        a31: 100000u64,
-        a32: 100000u64,
-        a33: 100000u64,
-        a34: 100000u64,
-        a35: 100000u64,
-        a36: 100000u64,
-        a37: 100000u64,
-        a38: 100000u64,
-        a39: 100000u64,
+        a00: 10000u64,
+        a01: 10000u64,
+        a02: 10000u64,
+        a03: 10000u64,
+        a04: 10000u64,
+        a05: 10000u64,
+        a06: 10000u64,
+        a07: 10000u64,
+        a08: 10000u64,
+        a09: 10000u64,
+        a10: 10000u64,
+        a11: 10000u64,
+        a12: 10000u64,
+        a13: 10000u64,
+        a14: 10000u64,
+        a15: 10000u64,
+        a16: 10000u64,
+        a17: 10000u64,
+        a18: 10000u64,
+        a19: 10000u64,
+        a20: 10000u64,
+        a21: 10000u64,
+        a22: 10000u64,
+        a23: 10000u64,
+        a24: 10000u64,
+        a25: 10000u64,
+        a26: 10000u64,
+        a27: 10000u64,
+        a28: 10000u64,
+        a29: 10000u64,
+        a30: 10000u64,
+        a31: 10000u64,
+        a32: 10000u64,
+        a33: 10000u64,
+        a34: 10000u64,
+        a35: 10000u64,
+        a36: 10000u64,
+        a37: 10000u64,
+        a38: 10000u64,
+        a39: 10000u64,
     });
 }
\ No newline at end of file
diff --git a/src/test/run-pass/too-much-recursion.rs b/src/test/run-pass/too-much-recursion.rs
new file mode 100644
index 00000000000..a10774d8f2b
--- /dev/null
+++ b/src/test/run-pass/too-much-recursion.rs
@@ -0,0 +1,12 @@
+// error-pattern:ran out of stack
+
+// Test that the task fails after hiting the recursion limit, but
+// that it doesn't bring down the whole proc
+
+fn main() {
+    task::spawn {||
+        task::unsupervise();
+        fn f() { f() };
+        f();
+    };
+}
\ No newline at end of file