From 606b68d8172ad561355ebcfb33d6e3e31d0105e8 Mon Sep 17 00:00:00 2001
From: Julien Portalier <julien@portalier.com>
Date: Thu, 7 Dec 2023 14:38:19 +0100
Subject: [PATCH] MT: move set_current_thread to Fiber (from
 Crystal::Scheduler)

This avoids manipulating `fiber.@current_thread` which ain't very
pretty.

Crystal::Scheduler is now entirely responsible for setting the current
thread. It's also responsible to verify that a making sure a fiber will
always be enqueued or resumed on the thread it's been associated to,
which is only affecting the `#resume(fiber)` and `#yield(fiber)` methods
that are barely used.

Lastly, we remove the current_thread store that was always replacing any
previous value on context swap. Sadly, this doesn't seem to have any
noticeable impact on performance.
---
 src/concurrent.cr            |  6 ++----
 src/crystal/scheduler.cr     | 27 +++++++++++++++++++++------
 src/crystal/system/thread.cr |  2 +-
 src/fiber.cr                 | 16 ++++++++++++++--
 4 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/src/concurrent.cr b/src/concurrent.cr
index 5b2747d0c581..af2f0aecf736 100644
--- a/src/concurrent.cr
+++ b/src/concurrent.cr
@@ -59,10 +59,8 @@ end
 # ```
 def spawn(*, name : String? = nil, same_thread = false, &block)
   fiber = Fiber.new(name, &block)
-  if same_thread
-    fiber.@current_thread.set(Thread.current)
-  end
-  Crystal::Scheduler.enqueue fiber
+  {% if flag?(:preview_mt) %} fiber.set_current_thread if same_thread {% end %}
+  fiber.enqueue
   fiber
 end
 
diff --git a/src/crystal/scheduler.cr b/src/crystal/scheduler.cr
index 0667626f787e..130ebc7280f8 100644
--- a/src/crystal/scheduler.cr
+++ b/src/crystal/scheduler.cr
@@ -51,6 +51,7 @@ class Crystal::Scheduler
   end
 
   def self.resume(fiber : Fiber) : Nil
+    validate_running_thread(fiber)
     Thread.current.scheduler.resume(fiber)
   end
 
@@ -63,9 +64,22 @@ class Crystal::Scheduler
   end
 
   def self.yield(fiber : Fiber) : Nil
+    validate_running_thread(fiber)
     Thread.current.scheduler.yield(fiber)
   end
 
+  private def self.validate_running_thread(fiber : Fiber) : Nil
+    {% if flag?(:preview_mt) %}
+      if th = fiber.get_current_thread
+        unless th == Thread.current
+          raise "BUG: tried to manually resume #{fiber} on #{Thread.current} instead of #{th}"
+        end
+      else
+        fiber.set_current_thread
+      end
+    {% end %}
+  end
+
   {% if flag?(:preview_mt) %}
     def self.enqueue_free_stack(stack : Void*) : Nil
       Thread.current.scheduler.enqueue_free_stack(stack)
@@ -76,11 +90,15 @@ class Crystal::Scheduler
     private getter(fiber_channel : Crystal::FiberChannel) { Crystal::FiberChannel.new }
     @free_stacks = Deque(Void*).new
   {% end %}
+
+  @main : Fiber
   @lock = Crystal::SpinLock.new
   @sleeping = false
 
   # :nodoc:
-  def initialize(@main : Fiber)
+  def initialize(thread : Thread)
+    @main = thread.main_fiber
+    {% if flag?(:preview_mt) %} @main.set_current_thread(thread) {% end %}
     @current = @main
     @runnables = Deque(Fiber).new
   end
@@ -95,8 +113,8 @@ class Crystal::Scheduler
 
   protected def resume(fiber : Fiber) : Nil
     validate_resumable(fiber)
+
     {% if flag?(:preview_mt) %}
-      set_current_thread(fiber)
       GC.lock_read
     {% elsif flag?(:interpreted) %}
       # No need to change the stack bottom!
@@ -130,10 +148,6 @@ class Crystal::Scheduler
     end
   end
 
-  private def set_current_thread(fiber)
-    fiber.@current_thread.set(Thread.current)
-  end
-
   private def fatal_resume_error(fiber, message)
     Crystal::System.print_error "\nFATAL: #{message}: #{fiber}\n"
     caller.each { |line| Crystal::System.print_error "  from #{line}\n" }
@@ -236,6 +250,7 @@ class Crystal::Scheduler
       @@workers = Array(Thread).new(count) do |i|
         if i == 0
           worker_loop = Fiber.new(name: "Worker Loop") { Thread.current.scheduler.run_loop }
+          worker_loop.set_current_thread
           Thread.current.scheduler.enqueue worker_loop
           Thread.current
         else
diff --git a/src/crystal/system/thread.cr b/src/crystal/system/thread.cr
index 878be639a7a3..88784ed68330 100644
--- a/src/crystal/system/thread.cr
+++ b/src/crystal/system/thread.cr
@@ -105,7 +105,7 @@ class Thread
   end
 
   # :nodoc:
-  getter scheduler : Crystal::Scheduler { Crystal::Scheduler.new(main_fiber) }
+  getter scheduler : Crystal::Scheduler { Crystal::Scheduler.new(self) }
 
   protected def start
     Thread.threads.push(self)
diff --git a/src/fiber.cr b/src/fiber.cr
index f89489c2bd13..aa2af7bf2229 100644
--- a/src/fiber.cr
+++ b/src/fiber.cr
@@ -62,7 +62,7 @@ class Fiber
   property name : String?
 
   @alive = true
-  @current_thread = Atomic(Thread?).new(nil)
+  {% if flag?(:preview_mt) %} @current_thread = Atomic(Thread?).new(nil) {% end %}
 
   # :nodoc:
   property next : Fiber?
@@ -136,7 +136,7 @@ class Fiber
       {% end %}
     thread.gc_thread_handler, @stack_bottom = GC.current_thread_stack_bottom
     @name = "main"
-    @current_thread.set(thread)
+    {% if flag?(:preview_mt) %} @current_thread.set(thread) {% end %}
     Fiber.fibers.push(self)
   end
 
@@ -305,4 +305,16 @@ class Fiber
     # Push the used section of the stack
     GC.push_stack @context.stack_top, @stack_bottom
   end
+
+  {% if flag?(:preview_mt) %}
+    # :nodoc:
+    def set_current_thread(thread = Thread.current) : Thread
+      @current_thread.set(thread)
+    end
+
+    # :nodoc:
+    def get_current_thread : Thread?
+      @current_thread.lazy_get
+    end
+  {% end %}
 end