rt: Store the task in TLS

Author: pcwalton

src/rt/rust_scheduler.cpp

@@ -1,8 +1,17 @@
 
 #include <stdarg.h>
+#include <cassert>
 #include "rust_internal.h"
 #include "globals.h"
 
+#ifndef _WIN32
+pthread_key_t rust_scheduler::task_key;
+#else
+DWORD rust_scheduler::task_key;
+#endif
+
+bool rust_scheduler::tls_initialized = false;
+
 rust_scheduler::rust_scheduler(rust_kernel *kernel,
                                rust_srv *srv,
                                int id) :
@@ -30,6 +39,9 @@ rust_scheduler::rust_scheduler(rust_kernel *kernel,
     pthread_attr_setstacksize(&attr, 1024 * 1024);
     pthread_attr_setdetachstate(&attr, true);
 #endif
+
+    if (!tls_initialized)
+        init_tls();
 }
 
 rust_scheduler::~rust_scheduler() {
@@ -275,6 +287,8 @@ rust_scheduler::start_main_loop() {
              scheduled_task->user.rust_sp,
              scheduled_task->state->name);
 
+        place_task_in_tls(scheduled_task);
+
         interrupt_flag = 0;
 
         DLOG(this, task,
@@ -347,6 +361,49 @@ void rust_scheduler::run() {
     this->start_main_loop();
 }
 
+#ifndef _WIN32
+void
+rust_scheduler::init_tls() {
+    int result = pthread_key_create(&task_key, NULL);
+    assert(!result && "Couldn't create the TLS key!");
+    tls_initialized = true;
+}
+
+void
+rust_scheduler::place_task_in_tls(rust_task *task) {
+    int result = pthread_setspecific(task_key, task);
+    assert(!result && "Couldn't place the task in TLS!");
+}
+
+rust_task *
+rust_scheduler::get_task() {
+    rust_task *task = reinterpret_cast<rust_task *>
+        (pthread_getspecific(task_key));
+    assert(task && "Couldn't get the task from TLS!");
+    return task;
+}
+#else
+void
+rust_scheduler::init_tls() {
+    task_key = TlsAlloc();
+    assert(task_key != TLS_OUT_OF_INDEXES && "Couldn't create the TLS key!");
+    tls_initialized = true;
+}
+
+void
+rust_scheduler::place_task_in_tls(rust_task *task) {
+    BOOL result = TlsSetValue(task_key, task);
+    assert(result && "Couldn't place the task in TLS!");
+}
+
+rust_task *
+rust_scheduler::get_task() {
+    rust_task *task = reinterpret_cast<rust_task *>(TlsGetValue(task_key));
+    assert(task && "Couldn't get the task from TLS!");
+    return task;
+}
+#endif
+
 //
 // Local Variables:
 // mode: C++

src/rt/rust_scheduler.h

@@ -1,6 +1,12 @@
 #ifndef RUST_SCHEDULER_H
 #define RUST_SCHEDULER_H
 
+#ifndef _WIN32
+#include <pthread.h>
+#else
+#include <windows.h>
+#endif
+
 struct rust_scheduler;
 
 class rust_crate_cache {
@@ -58,8 +64,13 @@ struct rust_scheduler : public kernel_owned<rust_scheduler>,
 
 #ifndef __WIN32__
     pthread_attr_t attr;
+    static pthread_key_t task_key;
+#else
+    static DWORD task_key;
 #endif
 
+    static bool tls_initialized;
+
     rust_env *env;
 
     // Only a pointer to 'name' is kept, so it must live as long as this
@@ -92,6 +103,11 @@ struct rust_scheduler : public kernel_owned<rust_scheduler>,
         kernel->win32_require(fn, ok);
     }
 #endif
+
+    void init_tls();
+    void place_task_in_tls(rust_task *task);
+
+    static rust_task *get_task();
 };
 
 inline rust_log &

src/rt/rust_upcall.cpp

@@ -1,6 +1,7 @@
 #include "rust_cc.h"
 #include "rust_gc.h"
 #include "rust_internal.h"
+#include "rust_scheduler.h"
 #include "rust_unwind.h"
 #include "rust_upcall.h"
 #include <stdint.h>
@@ -54,7 +55,8 @@ upcall_fail(rust_task *task,
 }
 
 extern "C" CDECL uintptr_t
-upcall_malloc(rust_task *task, size_t nbytes, type_desc *td) {
+upcall_malloc(rust_task *unused_task, size_t nbytes, type_desc *td) {
+    rust_task *task = rust_scheduler::get_task();
     LOG_UPCALL_ENTRY(task);
 
     LOG(task, mem,