Implement forwarding a safepoint to all Fibers of a Thread for SafepointPredicate.ALL_THREADS_AND_FIBERS

* Will also be useful for Fiber#backtrace added in Ruby 3.

Author: eregon

src/main/java/org/truffleruby/core/fiber/FiberManager.java

@@ -24,6 +24,7 @@
 import org.truffleruby.core.proc.RubyProc;
 import org.truffleruby.core.thread.ThreadManager;
 import org.truffleruby.core.thread.ThreadManager.BlockingAction;
+import org.truffleruby.language.SafepointAction;
 import org.truffleruby.language.control.BreakException;
 import org.truffleruby.language.control.DynamicReturnException;
 import org.truffleruby.language.control.ExitException;
@@ -42,6 +43,7 @@
 public class FiberManager {
 
     public static final String NAME_PREFIX = "Ruby Fiber";
+    public static final Object[] SAFEPOINT_ARGS = new Object[]{ FiberSafepointMessage.class };
 
     private final RubyLanguage language;
     private final RubyContext context;
@@ -106,6 +108,7 @@ private void fiberMain(RubyContext context, RubyFiber fiber, RubyProc block, Nod
 
         final Object prev = truffleContext.enter(currentNode);
         language.setupCurrentThread(thread, fiber.rubyThread);
+        fiber.rubyThread.setCurrentFiber(fiber);
 
         FiberMessage lastMessage = null;
         try {
@@ -201,7 +204,13 @@ private void assertNotEntered(String reason) {
 
     @TruffleBoundary
     private Object[] handleMessage(RubyFiber fiber, FiberMessage message, Node currentNode) {
-        fiber.rubyThread.setCurrentFiber(fiber);
+        // Written as a loop to not grow the stack when processing guest safepoints
+        while (message instanceof FiberSafepointMessage) {
+            final FiberSafepointMessage safepointMessage = (FiberSafepointMessage) message;
+            safepointMessage.action.run(fiber.rubyThread, currentNode);
+            final RubyFiber sendingFiber = safepointMessage.sendingFiber;
+            message = resumeAndWait(fiber, sendingFiber, FiberOperation.TRANSFER, SAFEPOINT_ARGS, currentNode);
+        }
 
         if (message instanceof FiberShutdownMessage) {
             throw new FiberShutdownException(currentNode);
@@ -210,13 +219,16 @@ private Object[] handleMessage(RubyFiber fiber, FiberMessage message, Node curre
         } else if (message instanceof FiberResumeMessage) {
             final FiberResumeMessage resumeMessage = (FiberResumeMessage) message;
             assert language.getCurrentThread() == resumeMessage.getSendingFiber().rubyThread;
-            if (resumeMessage.getOperation() == FiberOperation.RESUME ||
-                    resumeMessage.getOperation() == FiberOperation.RAISE) {
+            final FiberOperation operation = resumeMessage.getOperation();
+
+            if (operation == FiberOperation.RESUME || operation == FiberOperation.RAISE) {
                 fiber.lastResumedByFiber = resumeMessage.getSendingFiber();
             }
-            if (resumeMessage.getOperation() == FiberOperation.RAISE) {
+
+            if (operation == FiberOperation.RAISE) {
                 throw new RaiseException(context, (RubyException) resumeMessage.getArgs()[0]);
             }
+
             return resumeMessage.getArgs();
         } else {
             throw CompilerDirectives.shouldNotReachHere();
@@ -232,16 +244,38 @@ private void resume(RubyFiber fromFiber, RubyFiber fiber, FiberOperation operati
     @TruffleBoundary
     public Object[] transferControlTo(RubyFiber fromFiber, RubyFiber fiber, FiberOperation operation, Object[] args,
             Node currentNode) {
-        final TruffleContext truffleContext = context.getEnv().getContext();
+        final FiberMessage message = resumeAndWait(fromFiber, fiber, operation, args, currentNode);
+        return handleMessage(fromFiber, message, currentNode);
+    }
 
+    @TruffleBoundary
+    private FiberMessage resumeAndWait(RubyFiber fromFiber, RubyFiber fiber, FiberOperation operation, Object[] args,
+            Node currentNode) {
+        final TruffleContext truffleContext = context.getEnv().getContext();
         final FiberMessage message = context
                 .getThreadManager()
                 .leaveAndEnter(truffleContext, currentNode, () -> {
                     resume(fromFiber, fiber, operation, args);
                     return waitMessage(fromFiber, currentNode);
                 });
+        fromFiber.rubyThread.setCurrentFiber(fromFiber);
+        return message;
+    }
 
-        return handleMessage(fromFiber, message, currentNode);
+    @TruffleBoundary
+    public void safepoint(RubyFiber fromFiber, RubyFiber fiber, SafepointAction action, Node currentNode) {
+        final TruffleContext truffleContext = context.getEnv().getContext();
+        final FiberResumeMessage returnMessage = (FiberResumeMessage) context
+                .getThreadManager()
+                .leaveAndEnter(truffleContext, currentNode, () -> {
+                    addToMessageQueue(fiber, new FiberSafepointMessage(fromFiber, action));
+                    return waitMessage(fromFiber, currentNode);
+                });
+        fromFiber.rubyThread.setCurrentFiber(fromFiber);
+
+        if (returnMessage.getArgs() != SAFEPOINT_ARGS) {
+            throw CompilerDirectives.shouldNotReachHere();
+        }
     }
 
     public void start(RubyFiber fiber, Thread javaThread) {
@@ -385,6 +419,16 @@ public Object[] getArgs() {
 
     }
 
+    private static class FiberSafepointMessage implements FiberMessage {
+        private final RubyFiber sendingFiber;
+        private final SafepointAction action;
+
+        private FiberSafepointMessage(RubyFiber sendingFiber, SafepointAction action) {
+            this.sendingFiber = sendingFiber;
+            this.action = action;
+        }
+    }
+
     /** Used to cleanup and terminate Fibers when the parent Thread dies. */
     private static class FiberShutdownException extends TerminationException {
         private static final long serialVersionUID = 1522270454305076317L;

src/main/java/org/truffleruby/language/SafepointAction.java

@@ -14,6 +14,7 @@
 import com.oracle.truffle.api.nodes.Node;
 import org.truffleruby.RubyContext;
 import org.truffleruby.RubyLanguage;
+import org.truffleruby.core.fiber.RubyFiber;
 import org.truffleruby.core.thread.RubyThread;
 
 import java.util.Objects;
@@ -62,7 +63,17 @@ protected final void perform(Access access) {
         final RubyThread rubyThread = context.getThreadManager().getRubyThreadForJavaThread(access.getThread());
         if (filter.test(context, rubyThread, this)) {
             run(rubyThread, access.getLocation());
+
+            if (filter == SafepointPredicate.ALL_THREADS_AND_FIBERS) {
+                final RubyFiber currentFiber = rubyThread.getCurrentFiber();
+                for (RubyFiber fiber : rubyThread.runningFibers) {
+                    if (fiber != currentFiber) {
+                        context.fiberManager.safepoint(currentFiber, fiber, this, access.getLocation());
+                    }
+                }
+            }
         }
+
     }
 
     public boolean isSynchronous() {

src/main/ruby/truffleruby/core/main.rb

@@ -56,7 +56,7 @@ def to_s
 end
 
 show_backtraces = -> {
-  $stderr.puts 'All Thread backtraces:'
+  $stderr.puts 'All Thread and Fiber backtraces:'
   Primitive.all_fibers_backtraces.each do |fiber, backtrace|
     $stderr.puts "#{fiber} of #{Primitive.fiber_thread(fiber)}", backtrace, nil
   end
@@ -68,7 +68,7 @@ def to_s
   if Truffle::Boot.get_option('platform-handle-interrupt')
     Primitive.vm_watch_signal 'INT', true, -> _signo do
       if Truffle::Boot.get_option('backtraces-on-interrupt')
-        puts 'Interrupting...'
+        $stderr.puts 'Interrupting...'
         show_backtraces.call
       end