Rewrite the exchange allocator to work without an active scheduler. #4457

10 files changed, 136 insertions, 92 deletions

diff --git a/src/rt/rust_exchange_alloc.cpp b/src/rt/rust_exchange_alloc.cpp
new file mode 100644
index 00000000000..6c0204ca736
--- /dev/null
+++ b/src/rt/rust_exchange_alloc.cpp
@@ -0,0 +1,63 @@
+// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#include "rust_exchange_alloc.h"
+#include "sync/sync.h"
+#include <stdlib.h>
+#include <assert.h>
+#include <string.h>
+#include <stdio.h>
+
+uintptr_t exchange_count = 0;
+
+void *
+rust_exchange_alloc::malloc(size_t size, bool zero) {
+  void *value = ::malloc(size);
+  assert(value);
+  if (zero) {
+    memset(value, 0, size);
+  }
+
+  sync::increment(exchange_count);
+
+  return value;
+}
+
+void *
+rust_exchange_alloc::calloc(size_t size) {
+  return this->malloc(size);
+}
+
+void *
+rust_exchange_alloc::realloc(void *ptr, size_t size) {
+  void *new_ptr = ::realloc(ptr, size);
+  assert(new_ptr);
+  return new_ptr;
+}
+
+void
+rust_exchange_alloc::free(void *ptr) {
+  sync::decrement(exchange_count);
+  ::free(ptr);
+}
+
+extern "C" uintptr_t *
+rust_get_exchange_count_ptr() {
+  return &exchange_count;
+}
+
+void
+rust_check_exchange_count_on_exit() {
+  if (exchange_count != 0) {
+    printf("exchange heap not empty on on exit");
+    printf("%d dangling allocations", (int)exchange_count);
+    abort();
+  }
+}
diff --git a/src/rt/rust_exchange_alloc.h b/src/rt/rust_exchange_alloc.h
new file mode 100644
index 00000000000..1b52929acf1
--- /dev/null
+++ b/src/rt/rust_exchange_alloc.h
@@ -0,0 +1,31 @@
+// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#ifndef RUST_EXCHANGE_ALLOC_H
+#define RUST_EXCHANGE_ALLOC_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+class rust_exchange_alloc {
+ public:
+    void *malloc(size_t size, bool zero = true);
+    void *calloc(size_t size);
+    void *realloc(void *mem, size_t size);
+    void free(void *mem);
+};
+
+extern "C" uintptr_t *
+rust_get_exchange_count_ptr();
+
+void
+rust_check_exchange_count_on_exit();
+
+#endif
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index 99a6ed94944..e0494c9300b 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -22,7 +22,6 @@
     KLOG_LVL(this, field, log_err, __VA_ARGS__)
 
 rust_kernel::rust_kernel(rust_env *env) :
-    _region(env, true),
     _log(NULL),
     max_task_id(INIT_TASK_ID-1), // sync_add_and_fetch increments first
     rval(0),
@@ -77,21 +76,21 @@ rust_kernel::fatal(char const *fmt, ...) {
 
 void *
 rust_kernel::malloc(size_t size, const char *tag) {
-    return _region.malloc(size, tag);
+    return exchange_alloc.malloc(size);
 }
 
 void *
 rust_kernel::calloc(size_t size, const char *tag) {
-    return _region.calloc(size, tag);
+    return exchange_alloc.calloc(size);
 }
 
 void *
 rust_kernel::realloc(void *mem, size_t size) {
-    return _region.realloc(mem, size);
+    return exchange_alloc.realloc(mem, size);
 }
 
 void rust_kernel::free(void *mem) {
-    _region.free(mem);
+    exchange_alloc.free(mem);
 }
 
 rust_sched_id
@@ -217,6 +216,7 @@ rust_kernel::run() {
     assert(osmain_driver != NULL);
     osmain_driver->start_main_loop();
     sched_reaper.join();
+    rust_check_exchange_count_on_exit();
     return rval;
 }
 
diff --git a/src/rt/rust_kernel.h b/src/rt/rust_kernel.h
index c3d5a5a19bb..11af02dace4 100644
--- a/src/rt/rust_kernel.h
+++ b/src/rt/rust_kernel.h
@@ -45,11 +45,12 @@
 #include <map>
 #include <vector>
 
-#include "memory_region.h"
+#include "rust_exchange_alloc.h"
 #include "rust_log.h"
 #include "rust_sched_reaper.h"
 #include "rust_type.h"
 #include "util/hash_map.h"
+#include "sync/lock_and_signal.h"
 
 class rust_scheduler;
 class rust_sched_driver;
@@ -71,7 +72,7 @@ struct exit_functions {
 };
 
 class rust_kernel {
-    memory_region _region;
+    rust_exchange_alloc exchange_alloc;
     rust_log _log;
 
     // The next task id
@@ -135,7 +136,7 @@ public:
     void *calloc(size_t size, const char *tag);
     void *realloc(void *mem, size_t size);
     void free(void *mem);
-    memory_region *region() { return &_region; }
+    rust_exchange_alloc *region() { return &exchange_alloc; }
 
     void fail();
 
diff --git a/src/rt/rust_sched_loop.cpp b/src/rt/rust_sched_loop.cpp
index 5ddfd88d4b4..0d0eaaee962 100644
--- a/src/rt/rust_sched_loop.cpp
+++ b/src/rt/rust_sched_loop.cpp
@@ -260,7 +260,7 @@ rust_sched_loop::run_single_turn() {
 
         assert(!extra_c_stack);
         if (cached_c_stack) {
-            destroy_stack(kernel->region(), cached_c_stack);
+            destroy_exchange_stack(kernel->region(), cached_c_stack);
             cached_c_stack = NULL;
         }
 
@@ -389,14 +389,15 @@ void
 rust_sched_loop::prepare_c_stack(rust_task *task) {
     assert(!extra_c_stack);
     if (!cached_c_stack && !task->have_c_stack()) {
-        cached_c_stack = create_stack(kernel->region(), C_STACK_SIZE);
+        cached_c_stack = create_exchange_stack(kernel->region(),
+                                               C_STACK_SIZE);
     }
 }
 
 void
 rust_sched_loop::unprepare_c_stack() {
     if (extra_c_stack) {
-        destroy_stack(kernel->region(), extra_c_stack);
+        destroy_exchange_stack(kernel->region(), extra_c_stack);
         extra_c_stack = NULL;
     }
 }
diff --git a/src/rt/rust_stack.cpp b/src/rt/rust_stack.cpp
index 466399bd5b5..3bcda8adf40 100644
--- a/src/rt/rust_stack.cpp
+++ b/src/rt/rust_stack.cpp
@@ -53,6 +53,8 @@ check_stack_canary(stk_seg *stk) {
     assert(stk->canary == canary_value && "Somebody killed the canary");
 }
 
+// XXX: Duplication here between the local and exchange heap constructors
+
 stk_seg *
 create_stack(memory_region *region, size_t sz) {
     size_t total_sz = sizeof(stk_seg) + sz;
@@ -69,3 +71,20 @@ destroy_stack(memory_region *region, stk_seg *stk) {
     deregister_valgrind_stack(stk);
     region->free(stk);
 }
+
+stk_seg *
+create_exchange_stack(rust_exchange_alloc *exchange, size_t sz) {
+    size_t total_sz = sizeof(stk_seg) + sz;
+    stk_seg *stk = (stk_seg *)exchange->malloc(total_sz, false);
+    memset(stk, 0, sizeof(stk_seg));
+    stk->end = (uintptr_t) &stk->data[sz];
+    add_stack_canary(stk);
+    register_valgrind_stack(stk);
+    return stk;
+}
+
+void
+destroy_exchange_stack(rust_exchange_alloc *exchange, stk_seg *stk) {
+    deregister_valgrind_stack(stk);
+    exchange->free(stk);
+}
diff --git a/src/rt/rust_stack.h b/src/rt/rust_stack.h
index 91a6f8b256a..51b884e47b1 100644
--- a/src/rt/rust_stack.h
+++ b/src/rt/rust_stack.h
@@ -12,6 +12,7 @@
 #define RUST_STACK_H
 
 #include "rust_globals.h"
+#include "rust_exchange_alloc.h"
 #include "memory_region.h"
 
 struct rust_task;
@@ -37,6 +38,12 @@ create_stack(memory_region *region, size_t sz);
 void
 destroy_stack(memory_region *region, stk_seg *stk);
 
+stk_seg *
+create_exchange_stack(rust_exchange_alloc *exchange, size_t sz);
+
+void
+destroy_exchange_stack(rust_exchange_alloc *exchange, stk_seg *stk);
+
 // Must be called before each time a stack is reused to tell valgrind
 // that the stack is accessible.
 void
diff --git a/src/rt/rust_task.cpp b/src/rt/rust_task.cpp
index a9246963ca4..e51af464e48 100644
--- a/src/rt/rust_task.cpp
+++ b/src/rt/rust_task.cpp
@@ -181,7 +181,7 @@ void task_start_wrapper(spawn_args *a)
         // free the environment (which should be a unique closure).
         const type_desc *td = env->td;
         td->drop_glue(NULL, NULL, NULL, box_body(env));
-        upcall_exchange_free(env);
+        task->kernel->region()->free(env);
     }
 
     // The cleanup work needs lots of stack
diff --git a/src/rt/rust_upcall.cpp b/src/rt/rust_upcall.cpp
index ea396146755..008b470fede 100644
--- a/src/rt/rust_upcall.cpp
+++ b/src/rt/rust_upcall.cpp
@@ -141,81 +141,6 @@ struct s_trace_args {
 };
 
 /**********************************************************************
- * Allocate an object in the exchange heap
- */
-
-struct s_exchange_malloc_args {
-    rust_task *task;
-    uintptr_t retval;
-    type_desc *td;
-    uintptr_t size;
-};
-
-extern "C" CDECL void
-upcall_s_exchange_malloc(s_exchange_malloc_args *args) {
-    rust_task *task = args->task;
-    LOG_UPCALL_ENTRY(task);
-
-    size_t total_size = get_box_size(args->size, args->td->align);
-    void *p = task->kernel->malloc(total_size, "exchange malloc");
-
-    rust_opaque_box *header = static_cast<rust_opaque_box*>(p);
-    header->ref_count = -1; // This is not ref counted
-    header->td = args->td;
-    header->prev = 0;
-    header->next = 0;
-
-    LOG(task, mem, "exchange malloced %p of size %" PRIuPTR,
-        header, args->size);
-
-    args->retval = (uintptr_t)header;
-}
-
-extern "C" CDECL uintptr_t
-upcall_exchange_malloc(type_desc *td, uintptr_t size) {
-    rust_task *task = rust_get_current_task();
-    s_exchange_malloc_args args = {task, 0, td, size};
-    UPCALL_SWITCH_STACK(task, &args, upcall_s_exchange_malloc);
-    return args.retval;
-}
-
-// FIXME (#2861): Alias used by libcore/rt.rs to avoid naming conflicts with
-// autogenerated wrappers for upcall_exchange_malloc. Remove this when we
-// fully move away away from the C upcall path.
-extern "C" CDECL uintptr_t
-rust_upcall_exchange_malloc(type_desc *td, uintptr_t size) {
-    return upcall_exchange_malloc(td, size);
-}
-
-struct s_exchange_free_args {
-    rust_task *task;
-    void *ptr;
-};
-
-extern "C" CDECL void
-upcall_s_exchange_free(s_exchange_free_args *args) {
-    rust_task *task = args->task;
-    LOG_UPCALL_ENTRY(task);
-    LOG(task, mem, "exchange freed %p", args->ptr);
-    task->kernel->free(args->ptr);
-}
-
-extern "C" CDECL void
-upcall_exchange_free(void *ptr) {
-    rust_task *task = rust_get_current_task();
-    s_exchange_free_args args = {task,ptr};
-    UPCALL_SWITCH_STACK(task, &args, upcall_s_exchange_free);
-}
-
-// FIXME (#2861): Alias used by libcore/rt.rs to avoid naming conflicts with
-// autogenerated wrappers for upcall_exchange_free. Remove this when we fully
-// move away away from the C upcall path.
-extern "C" CDECL void
-rust_upcall_exchange_free(void *ptr) {
-    return upcall_exchange_free(ptr);
-}
-
-/**********************************************************************
  * Allocate an object in the task-local heap.
  */
 
diff --git a/src/rt/rustrt.def.in b/src/rt/rustrt.def.in
index dcc02341e76..9076670392a 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -68,10 +68,6 @@ upcall_call_shim_on_rust_stack
 upcall_new_stack
 upcall_del_stack
 upcall_reset_stack_limit
-upcall_exchange_malloc
-upcall_exchange_free
-rust_upcall_exchange_free
-rust_upcall_exchange_malloc
 rust_upcall_fail
 rust_upcall_free
 rust_upcall_malloc
@@ -194,4 +190,5 @@ rust_raw_thread_join_delete
 rust_register_exit_function
 rust_get_global_data_ptr
 rust_inc_weak_task_count
-rust_dec_weak_task_count
\ No newline at end of file
+rust_dec_weak_task_count
+rust_get_exchange_count_ptr