From a08ee8330d3d739467bfa34deeb797d83e59ed3c Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Fri, 16 Oct 2020 14:25:58 +1300
Subject: Rename to `Fiber#set_scheduler`.

---
 scheduler.c | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

(limited to 'scheduler.c')

diff --git a/scheduler.c b/scheduler.c
index d7e713a710..326cb0068f 100644
--- a/scheduler.c
+++ b/scheduler.c
@@ -8,6 +8,7 @@
 
 **********************************************************************/
 
+#include "vm_core.h"
 #include "internal/scheduler.h"
 #include "ruby/io.h"
 
@@ -37,6 +38,54 @@ Init_Scheduler(void)
     id_io_wait = rb_intern_const("io_wait");
 }
 
+VALUE
+rb_scheduler_get()
+{
+    rb_thread_t *thread = GET_THREAD();
+    VM_ASSERT(thread);
+
+    return thread->scheduler;
+}
+
+VALUE
+rb_scheduler_set(VALUE scheduler)
+{
+    rb_thread_t *thread = GET_THREAD();
+    VM_ASSERT(thread);
+
+    // We invoke Scheduler#close when setting it to something else, to ensure the previous scheduler runs to completion before changing the scheduler. That way, we do not need to consider interactions, e.g., of a Fiber from the previous scheduler with the new scheduler.
+    if (thread->scheduler != Qnil) {
+        rb_scheduler_close(thread->scheduler);
+    }
+
+    thread->scheduler = scheduler;
+
+    return thread->scheduler;
+}
+
+static VALUE
+rb_threadptr_scheduler_current(rb_thread_t *thread)
+{
+    VM_ASSERT(thread);
+
+    if (thread->blocking == 0) {
+        return thread->scheduler;
+    } else {
+        return Qnil;
+    }
+}
+
+VALUE
+rb_scheduler_current()
+{
+    return rb_threadptr_scheduler_current(GET_THREAD());
+}
+
+VALUE rb_thread_scheduler_current(VALUE thread)
+{
+    return rb_threadptr_scheduler_current(rb_thread_ptr(thread));
+}
+
 VALUE
 rb_scheduler_close(VALUE scheduler)
 {
-- 
cgit v1.2.1