rt: Don't switch to the C stack on the upcall_new_stack fast path

src/rt/rust_task.cpp

@@ -509,6 +509,37 @@ rust_task::free_stack(stk_seg *stk) {
     destroy_stack(&local_region, stk);
 }
 
+struct new_stack_args {
+    rust_task *task;
+    size_t requested_sz;
+};
+
+void
+new_stack_slow(new_stack_args *args) {
+    args->task->new_stack(args->requested_sz);
+}
+
+// NB: This runs on the Rust stack
+// This is the new stack fast path, in which we
+// reuse the next cached stack segment
+void
+rust_task::new_stack_fast(size_t requested_sz) {
+    // The minimum stack size, in bytes, of a Rust stack, excluding red zone
+    size_t min_sz = thread->min_stack_size;
+
+    // Try to reuse an existing stack segment
+    if (stk != NULL && stk->prev != NULL) {
+        size_t prev_sz = user_stack_size(stk->prev);
+        if (min_sz <= prev_sz && requested_sz <= prev_sz) {
+            stk = stk->prev;
+            return;
+        }
+    }
+
+    new_stack_args args = {this, requested_sz};
+    call_on_c_stack(&args, (void*)new_stack_slow);
+}
+
 void
 rust_task::new_stack(size_t requested_sz) {
     LOG(this, mem, "creating new stack for task %" PRIxPTR, this);
@@ -572,7 +603,7 @@ rust_task::next_stack(size_t stk_sz, void *args_addr, size_t args_sz) {
         maybe_next_stack = stk->prev;
     }
 
-    new_stack(stk_sz + args_sz);
+    new_stack_fast(stk_sz + args_sz);
     A(thread, stk->end - (uintptr_t)stk->data >= stk_sz + args_sz,
       "Did not receive enough stack");
     uint8_t *new_sp = (uint8_t*)stk->end;

src/rt/rust_task.h

@@ -33,6 +33,7 @@ typedef unsigned long task_result;
 struct spawn_args;
 struct cleanup_args;
 struct reset_args;
+struct new_stack_args;
 
 // std::lib::task::task_notification
 //
@@ -114,7 +115,8 @@ private:
     // Called when the atomic refcount reaches zero
     void delete_this();
 
-    void new_stack(size_t sz);
+    void new_stack_fast(size_t requested_sz);
+    void new_stack(size_t requested_sz);
     void free_stack(stk_seg *stk);
     size_t get_next_stack_size(size_t min, size_t current, size_t requested);
 
@@ -128,6 +130,7 @@ private:
     friend void task_start_wrapper(spawn_args *a);
     friend void cleanup_task(cleanup_args *a);
     friend void reset_stack_limit_on_c_stack(reset_args *a);
+    friend void new_stack_slow(new_stack_args *a);
 
 public:
 

src/rt/rust_upcall.cpp

@@ -439,26 +439,13 @@ upcall_log_type(const type_desc *tydesc, uint8_t *data, uint32_t level) {
     UPCALL_SWITCH_STACK(&args, upcall_s_log_type);
 }
 
-struct s_new_stack_args {
+// NB: This needs to be blazing fast. Don't switch stacks
-    void *result;
-    size_t stk_sz;
-    void *args_addr;
-    size_t args_sz;
-};
-
-extern "C" CDECL void
-upcall_s_new_stack(struct s_new_stack_args *args) {
-    rust_task *task = rust_task_thread::get_task();
-    args->result = task->next_stack(args->stk_sz,
-                                    args->args_addr,
-                                    args->args_sz);
-}
-
 extern "C" CDECL void *
 upcall_new_stack(size_t stk_sz, void *args_addr, size_t args_sz) {
-    s_new_stack_args args = {NULL, stk_sz, args_addr, args_sz};
+    rust_task *task = rust_task_thread::get_task();
-    UPCALL_SWITCH_STACK(&args, upcall_s_new_stack);
+    return task->next_stack(stk_sz,
-    return args.result;
+                            args_addr,
+                            args_sz);
 }
 
 // NB: This needs to be blazing fast. Don't switch stacks