Merge pull request #1183 from herwinw/fiber_scheduler

natalie-lang · Aug 31, 2023 · d587cbd · d587cbd
2 parents 00b3bd8 + 70a471e
commit d587cbd
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 0 deletions.
diff --git a/include/natalie/fiber_object.hpp b/include/natalie/fiber_object.hpp
@@ -147,6 +147,8 @@ class FiberObject : public Object {
     static Value ref(Env *env, Value);
     static Value refeq(Env *env, Value, Value);
     Value resume(Env *env, Args args);
+    static Value scheduler();
+    static Value set_scheduler(Env *, Value);
     Value set_storage(Env *, Value);
     Value storage(Env *) const;
     void yield_back(Env *env, Args args);
@@ -211,6 +213,7 @@ class FiberObject : public Object {
 
     inline static FiberObject *s_current = nullptr;
     inline static FiberObject *s_main = nullptr;
+    inline static Value s_scheduler { nullptr };
 };
 
 }
diff --git a/lib/natalie/compiler/binding_gen.rb b/lib/natalie/compiler/binding_gen.rb
@@ -554,6 +554,8 @@ def generate_name
 gen.static_binding('Fiber', '[]=', 'FiberObject', 'refeq', argc: 2, pass_env: true, pass_block: false, return_type: :Object)
 gen.static_binding('Fiber', 'blocking?', 'FiberObject', 'is_blocking_current', argc: 0, pass_env: false, pass_block: false, return_type: :Object)
 gen.static_binding('Fiber', 'current', 'FiberObject', 'current', argc: 0, pass_env: false, pass_block: false, return_type: :Object)
+gen.static_binding('Fiber', 'scheduler', 'FiberObject', 'scheduler', argc: 0, pass_env: false, pass_block: false, return_type: :Object)
+gen.static_binding('Fiber', 'set_scheduler', 'FiberObject', 'set_scheduler', argc: 1, pass_env: true, pass_block: false, return_type: :Object)
 gen.static_binding('Fiber', 'yield', 'FiberObject', 'yield', argc: :any, pass_env: true, pass_block: false, return_type: :Object)
 gen.binding('Fiber', 'alive?', 'FiberObject', 'is_alive', argc: 0, pass_env: false, pass_block: false, return_type: :bool)
 gen.binding('Fiber', 'blocking?', 'FiberObject', 'is_blocking', argc: 0, pass_env: false, pass_block: false, return_type: :bool)

diff --git a/src/fiber_object.cpp b/src/fiber_object.cpp
@@ -116,6 +116,24 @@ Value FiberObject::resume(Env *env, Args args) {
     }
 }
 
+Value FiberObject::scheduler() {
+    return s_scheduler;
+}
+
+Value FiberObject::set_scheduler(Env *env, Value scheduler) {
+    if (scheduler->is_nil()) {
+        s_scheduler = nullptr;
+    } else {
+        TM::Vector<TM::String> required_methods { "block", "unblock", "kernel_sleep", "io_wait" };
+        for (const auto &required_method : required_methods) {
+            if (!scheduler->respond_to(env, SymbolObject::intern(required_method)))
+                env->raise("ArgumentError", "Scheduler must implement #{}", required_method);
+        }
+        s_scheduler = scheduler;
+    }
+    return scheduler;
+}
+
 Value FiberObject::set_storage(Env *env, Value storage) {
     if (storage == nullptr || storage->is_nil()) {
         m_storage = nullptr;

diff --git a/test/natalie/fiber_test.rb b/test/natalie/fiber_test.rb
@@ -54,3 +54,45 @@
                                     )
   end
 end
+
+describe 'Scheduler' do
+  it 'validates the scheduler for required methods' do
+    # We should not be testing for ordering of these error messages
+    required_methods = [:block, :unblock, :kernel_sleep, :io_wait]
+    required_methods.each do |missing_method|
+      scheduler = Object.new
+      required_methods.difference([missing_method]).each do |method|
+        scheduler.define_singleton_method(method) {}
+      end
+
+      -> { Fiber.set_scheduler(scheduler) }.should raise_error(ArgumentError, /Scheduler must implement ##{missing_method}/)
+    end
+  end
+
+  it 'can set and get the scheduler' do
+    required_methods = [:block, :unblock, :kernel_sleep, :io_wait]
+    scheduler = Object.new
+    required_methods.each do |method|
+      scheduler.define_singleton_method(method) {}
+    end
+    Fiber.set_scheduler(scheduler)
+    Fiber.scheduler.should == scheduler
+  end
+
+  it 'can remove the scheduler' do
+    required_methods = [:block, :unblock, :kernel_sleep, :io_wait]
+    scheduler = Object.new
+    required_methods.each do |method|
+      scheduler.define_singleton_method(method) {}
+    end
+    Fiber.set_scheduler(scheduler)
+    Fiber.set_scheduler(nil)
+    Fiber.scheduler.should be_nil
+  end
+
+  it 'can assign a nil scheduler multiple times' do
+    Fiber.set_scheduler(nil)
+    Fiber.set_scheduler(nil)
+    Fiber.scheduler.should be_nil
+  end
+end