Add at_exit function rust-lang#4450

Author: brson

src/libcore/private.rs

@@ -28,6 +28,9 @@ use task;
 use task::{TaskBuilder, atomically};
 use uint;
 
+#[path = "private/at_exit.rs"]
+pub mod at_exit;
+
 extern mod rustrt {
     #[legacy_exports];
     unsafe fn rust_task_weaken(ch: rust_port_id);

src/libcore/private/at_exit.rs

@@ -0,0 +1,79 @@
+use sys;
+use cast;
+use ptr;
+use task;
+use uint;
+use vec;
+use libc::{c_void, size_t};
+
+/**
+Register a function to be run during runtime shutdown.
+
+After all non-weak tasks have exited, registered exit functions will
+execute, in random order, on the primary scheduler. Each function runs
+in its own unsupervised task.
+*/
+pub fn at_exit(f: ~fn()) unsafe {
+    let runner: &fn(*ExitFunctions) = exit_runner;
+    let runner_pair: sys::Closure = cast::transmute(runner);
+    let runner_ptr = runner_pair.code;
+    let runner_ptr = cast::transmute(runner_ptr);
+    rustrt::rust_register_exit_function(runner_ptr, ~f);
+}
+
+// NB: The double pointer indirection here is because ~fn() is a fat
+// pointer and due to FFI problems I am more comfortable making the
+// interface use a normal pointer
+extern mod rustrt {
+    fn rust_register_exit_function(runner: *c_void, f: ~~fn());
+}
+
+struct ExitFunctions {
+    // The number of exit functions
+    count: size_t,
+    // The buffer of exit functions
+    start: *~~fn()
+}
+
+fn exit_runner(exit_fns: *ExitFunctions) unsafe {
+    let exit_fns = &*exit_fns;
+    let count = (*exit_fns).count;
+    let start = (*exit_fns).start;
+    let exit_fns_vec = vec::from_buf(start, count as uint);
+    // Run in random order to prevent accidental dependencies on
+    // shutdown order
+    let exit_fns_vec = rand::Rng().shuffle(exit_fns_vec);
+    for exit_fns_vec.each |exit_fn| {
+        let f: ~~fn() = cast::reinterpret_cast(&*exit_fn);
+        match f {
+            ~f => {
+                task::task().supervised().spawn(f);
+            }
+        }
+    }
+}
+
+#[abi = "rust-intrinsic"]
+pub extern mod rusti {
+    fn move_val_init<T>(dst: &mut T, -src: T);
+    fn init<T>() -> T;
+}
+
+#[test]
+fn test_at_exit() {
+    let i = 10;
+    do at_exit {
+        assert i == 10;
+    }
+}
+
+#[test]
+fn test_at_exit_many() {
+    let i = 10;
+    for uint::range(20, 100) |j| {
+        do at_exit {
+            assert i == 10;
+            assert j > i;
+        }
+    }
+}

src/rt/rust_builtin.cpp

@@ -975,6 +975,12 @@ rust_log_str(uint32_t level, const char *str, size_t size) {
     task->sched_loop->get_log().log(task, level, "%.*s", (int)size, str);
 }
 
+extern "C" void
+rust_register_exit_function(spawn_fn runner, fn_env_pair *f) {
+    rust_task *task = rust_get_current_task();
+    task->kernel->register_exit_function(runner, f);
+}
+
 //
 // Local Variables:
 // mode: C++

src/rt/rust_kernel.cpp

@@ -36,6 +36,8 @@ rust_kernel::rust_kernel(rust_env *env) :
     non_weak_tasks(0),
     global_loop_chan(0),
     global_env_chan(0),
+    at_exit_runner(NULL),
+    at_exit_started(false),
     env(env)
 
 {
@@ -427,6 +429,7 @@ rust_kernel::begin_shutdown() {
         }
     }
 
+    run_exit_functions();
     allow_scheduler_exit();
     end_weak_tasks();
 }
@@ -446,6 +449,47 @@ rust_kernel::send_to_port(rust_port_id chan, void *sptr) {
     }
 }
 
+void
+rust_kernel::register_exit_function(spawn_fn runner, fn_env_pair *f) {
+    scoped_lock with(at_exit_lock);
+
+    assert(!at_exit_started && "registering at_exit function after exit");
+
+    if (at_exit_runner) {
+        assert(runner == at_exit_runner
+               && "there can be only one at_exit_runner");
+    }
+
+    at_exit_runner = runner;
+    at_exit_fns.push_back(f);
+}
+
+void
+rust_kernel::run_exit_functions() {
+    rust_task *task;
+
+    {
+        scoped_lock with(at_exit_lock);
+
+        assert(!at_exit_started && "running exit functions twice?");
+
+        at_exit_started = true;
+
+        if (at_exit_runner == NULL) {
+            return;
+        }
+
+        rust_scheduler *sched = get_scheduler_by_id(main_sched_id());
+        assert(sched);
+        task = sched->create_task(NULL, "at_exit");
+
+        final_exit_fns.count = at_exit_fns.size();
+        final_exit_fns.start = at_exit_fns.data();
+    }
+
+    task->start(at_exit_runner, NULL, &final_exit_fns);
+}
+
 //
 // Local Variables:
 // mode: C++

src/rt/rust_kernel.h

@@ -50,6 +50,7 @@
 #include "memory_region.h"
 #include "rust_log.h"
 #include "rust_sched_reaper.h"
+#include "rust_type.h"
 #include "util/hash_map.h"
 
 class rust_scheduler;
@@ -66,6 +67,13 @@ typedef intptr_t rust_port_id;
 
 typedef std::map<rust_sched_id, rust_scheduler*> sched_map;
 
+// This is defined as a struct only because we need a single pointer to pass
+// to the Rust function that runs the at_exit functions
+struct exit_functions {
+    size_t count;
+    fn_env_pair **start;
+};
+
 class rust_kernel {
     memory_region _region;
     rust_log _log;
@@ -126,6 +134,14 @@ class rust_kernel {
     // Used to serialize access to getenv/setenv
     uintptr_t global_env_chan;
 
+    lock_and_signal at_exit_lock;
+    spawn_fn at_exit_runner;
+    bool at_exit_started;
+    std::vector<fn_env_pair*> at_exit_fns;
+    exit_functions final_exit_fns;
+
+    void run_exit_functions();
+
 public:
     struct rust_env *env;
 
@@ -175,6 +191,8 @@ class rust_kernel {
 
     uintptr_t* get_global_loop() { return &global_loop_chan; }
     uintptr_t* get_global_env_chan() { return &global_env_chan; }
+
+    void register_exit_function(spawn_fn runner, fn_env_pair *f);
 };
 
 template <typename T> struct kernel_owned {

src/rt/rustrt.def.in

@@ -208,3 +208,4 @@ linenoiseHistoryAdd
 linenoiseHistorySetMaxLen
 linenoiseHistorySave
 linenoiseHistoryLoad
+rust_register_exit_function