[GR-44211] Use Env#newTruffleThreadBuilder for Ruby Fibers

Author: eregon

spec/tags/core/exception/top_level_tags.txt

@@ -3,4 +3,3 @@ slow:An Exception reaching the top level with a custom backtrace is printed on S
 slow:An Exception reaching the top level the Exception#cause is printed to STDERR with backtraces
 slow:An Exception reaching the top level kills all threads and fibers, ensure clauses are only run for threads current fibers, not for suspended fibers with ensure on the root fiber
 slow:An Exception reaching the top level kills all threads and fibers, ensure clauses are only run for threads current fibers, not for suspended fibers with ensure on non-root fiber
-fails(GR-44211):An Exception reaching the top level kills all threads and fibers, ensure clauses are only run for threads current fibers, not for suspended fibers with ensure on non-root fiber

src/main/java/org/truffleruby/RubyLanguage.java

@@ -406,6 +406,17 @@ public void invalidateTracingAssumption() {
         tracingAssumption = tracingCyclicAssumption.getAssumption();
     }
 
+    private boolean multiThreading = false;
+
+    public boolean isMultiThreaded() {
+        return multiThreading;
+    }
+
+    @Override
+    protected void initializeMultiThreading(RubyContext context) {
+        this.multiThreading = true;
+    }
+
     @Override
     protected void initializeMultipleContexts() {
         LOGGER.fine("initializeMultipleContexts()");

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

@@ -65,8 +65,12 @@ public void initialize(RubyFiber fiber, boolean blocking, RubyProc block, Node c
         final TruffleContext truffleContext = context.getEnv().getContext();
 
         context.getThreadManager().leaveAndEnter(truffleContext, currentNode, () -> {
-            context.getThreadManager().spawnFiber(fiber, sourceSection,
-                    () -> fiberMain(context, fiber, block, currentNode));
+            ThreadManager threadManager = context.getThreadManager();
+            Thread thread = threadManager.createFiberJavaThread(fiber, sourceSection,
+                    () -> beforeEnter(fiber, currentNode),
+                    () -> fiberMain(context, fiber, block, currentNode),
+                    () -> afterLeave(fiber), currentNode);
+            thread.start();
             waitForInitialization(context, fiber, currentNode);
             return BlockingAction.SUCCESS;
         });
@@ -90,23 +94,25 @@ public static void waitForInitialization(RubyContext context, RubyFiber fiber, N
 
     private static final BranchProfile UNPROFILED = BranchProfile.getUncached();
 
-    private void fiberMain(RubyContext context, RubyFiber fiber, RubyProc block, Node currentNode) {
+    private void beforeEnter(RubyFiber fiber, Node currentNode) {
         assert !fiber.isRootFiber() : "Root Fibers execute threadMain() and not fiberMain()";
         assertNotEntered("Fibers should start unentered to avoid triggering multithreading");
 
         final Thread thread = Thread.currentThread();
-        final TruffleContext truffleContext = context.getEnv().getContext();
-
         start(fiber, thread);
 
         // fully initialized
         fiber.initializedLatch.countDown();
 
-        final FiberMessage message = waitMessage(fiber, currentNode);
-        fiber.rubyThread.setCurrentFiber(fiber);
+        fiber.firstMessage = waitMessage(fiber, currentNode);
 
         // enter() polls so we need the current Fiber to be set before enter()
-        final Object prev = truffleContext.enter(currentNode);
+        fiber.rubyThread.setCurrentFiber(fiber);
+    }
+
+    private void fiberMain(RubyContext context, RubyFiber fiber, RubyProc block, Node currentNode) {
+        final FiberMessage message = fiber.firstMessage;
+        fiber.firstMessage = null;
 
         FiberMessage lastMessage = null;
         try {
@@ -134,18 +140,24 @@ private void fiberMain(RubyContext context, RubyFiber fiber, RubyProc block, Nod
             final RuntimeException exception = ThreadManager.printInternalError(e);
             lastMessage = new FiberExceptionMessage(exception);
         } finally {
-            final RubyFiber returnFiber = lastMessage == null ? null : getReturnFiber(fiber, currentNode, UNPROFILED);
+            fiber.lastMessage = lastMessage;
+            fiber.returnFiber = lastMessage == null ? null : getReturnFiber(fiber, currentNode, UNPROFILED);
 
             // Perform all cleanup before resuming the parent Fiber
             // Make sure that other fibers notice we are dead before they gain control back
             fiber.status = FiberStatus.TERMINATED;
             // Leave context before addToMessageQueue() -> parent Fiber starts executing
-            truffleContext.leave(currentNode, prev);
-            cleanup(fiber, thread);
+        }
+    }
 
-            if (lastMessage != null) {
-                addToMessageQueue(returnFiber, lastMessage);
-            }
+    private void afterLeave(RubyFiber fiber) {
+        final Thread thread = Thread.currentThread();
+        cleanup(fiber, thread);
+
+        if (fiber.lastMessage != null) {
+            addToMessageQueue(fiber.returnFiber, fiber.lastMessage);
+            fiber.returnFiber = null;
+            fiber.lastMessage = null;
         }
     }
 
@@ -161,7 +173,7 @@ public RubyFiber getReturnFiber(RubyFiber currentFiber, Node currentNode, Branch
             return previousFiber;
         } else {
 
-            if (currentFiber == rootFiber) {
+            if (currentFiber == rootFiber) { // Note: this is always false for the fiberMain() caller
                 errorProfile.enter();
                 throw new RaiseException(context, context.getCoreExceptions().yieldFromRootFiberError(currentNode));
             }

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

@@ -121,6 +121,13 @@ public abstract static class InitializeNode extends PrimitiveArrayArgumentsNode
         @TruffleBoundary
         @Specialization
         protected Object initialize(RubyFiber fiber, boolean blocking, RubyProc block) {
+            if (!getContext().getEnv().isCreateThreadAllowed()) {
+                // Because TruffleThreadBuilder#build denies it already, before the thread is even started.
+                // The permission is called allowCreateThread, so it kind of makes sense.
+                throw new RaiseException(getContext(),
+                        coreExceptions().securityError("fibers not allowed with allowCreateThread(false)", this));
+            }
+
             getContext().fiberManager.initialize(fiber, blocking, block, this);
             return nil;
         }

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

@@ -92,6 +92,11 @@ public enum FiberStatus {
     boolean blocking = true;
     public RubyArray cGlobalVariablesDuringInitFunction;
 
+    // To pass state between beforeEnter(), fiberMain() and afterLeave()
+    FiberManager.FiberMessage firstMessage;
+    RubyFiber returnFiber;
+    FiberManager.FiberMessage lastMessage;
+
     public RubyFiber(
             RubyClass rubyClass,
             Shape shape,

src/main/java/org/truffleruby/core/thread/ThreadManager.java

@@ -21,6 +21,7 @@
 import java.util.function.Supplier;
 
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.ValueType;
 import com.oracle.truffle.api.TruffleContext;
 import com.oracle.truffle.api.TruffleOptions;
@@ -164,16 +165,29 @@ private static ThreadFactory getVirtualThreadFactory() {
 
     @CompilationFinal static ThreadFactory VIRTUAL_THREAD_FACTORY = getVirtualThreadFactory();
 
-    private Thread createFiberJavaThread(RubyFiber fiber, SourceSection sourceSection, Runnable runnable) {
+    public Thread createFiberJavaThread(RubyFiber fiber, SourceSection sourceSection, Runnable beforeEnter,
+            Runnable body, Runnable afterLeave, Node node) {
         if (context.isPreInitializing()) {
-            throw new UnsupportedOperationException("fibers should not be created while pre-initializing the context");
+            throw CompilerDirectives
+                    .shouldNotReachHere("fibers should not be created while pre-initializing the context");
         }
 
         final Thread thread;
         if (context.getOptions().VIRTUAL_THREAD_FIBERS) {
-            thread = VIRTUAL_THREAD_FACTORY.newThread(runnable);
+            thread = VIRTUAL_THREAD_FACTORY.newThread(() -> {
+                var truffleContext = context.getEnv().getContext();
+                beforeEnter.run();
+                Object prev = truffleContext.enter(node);
+                try {
+                    body.run();
+                } finally {
+                    truffleContext.leave(node, prev);
+                    afterLeave.run();
+                }
+            });
         } else {
-            thread = new Thread(runnable); // context.getEnv().createUnenteredThread(runnable);
+            thread = context.getEnv().newTruffleThreadBuilder(body).beforeEnter(beforeEnter).afterLeave(afterLeave)
+                    .build();
         }
 
         language.rubyThreadInitMap.put(thread, fiber.rubyThread);
@@ -186,15 +200,16 @@ private Thread createFiberJavaThread(RubyFiber fiber, SourceSection sourceSectio
         return thread;
     }
 
-    private Thread createJavaThread(Runnable runnable, RubyThread rubyThread, String info) {
+    private Thread createJavaThread(Runnable runnable, RubyThread rubyThread, String info, Node node) {
         if (context.getOptions().SINGLE_THREADED) {
             throw new RaiseException(
                     context,
-                    context.getCoreExceptions().securityError("threads not allowed in single-threaded mode", null));
+                    context.getCoreExceptions().securityError("threads not allowed in single-threaded mode", node));
         }
 
         if (context.isPreInitializing()) {
-            throw new UnsupportedOperationException("threads should not be created while pre-initializing the context");
+            throw CompilerDirectives
+                    .shouldNotReachHere("threads should not be created while pre-initializing the context");
         }
 
         final Thread thread = context.getEnv().newTruffleThreadBuilder(runnable).build();
@@ -235,10 +250,6 @@ private static Thread.UncaughtExceptionHandler uncaughtExceptionHandler(RubyFibe
         };
     }
 
-    public void spawnFiber(RubyFiber fiber, SourceSection sourceSection, Runnable task) {
-        createFiberJavaThread(fiber, sourceSection, task).start();
-    }
-
     /** Whether the thread was created by TruffleRuby. */
     @TruffleBoundary
     public boolean isRubyManagedThread(Thread thread) {
@@ -303,7 +314,8 @@ public void initialize(RubyThread rubyThread, Node currentNode, String info, Str
         rubyThread.sourceLocation = info;
         final RubyFiber rootFiber = rubyThread.getRootFiber();
 
-        final Thread thread = createJavaThread(() -> threadMain(rubyThread, currentNode, task), rubyThread, info);
+        final Thread thread = createJavaThread(() -> threadMain(rubyThread, currentNode, task), rubyThread, info,
+                currentNode);
         rubyThread.thread = thread;
 
         thread.start();

src/main/java/org/truffleruby/debug/TruffleDebugNodes.java

@@ -1382,6 +1382,7 @@ protected RubyString parseName(InternalMethod method) {
     /** Creates a Truffle thread which is not {@link ThreadManager#isRubyManagedThread(java.lang.Thread)}}. */
     @CoreMethod(names = "create_polyglot_thread", onSingleton = true, required = 1)
     public abstract static class CreatePolyglotThread extends CoreMethodArrayArgumentsNode {
+        @TruffleBoundary
         @Specialization
         protected Object parseName(Object hostRunnable) {
             Runnable runnable = (Runnable) getContext().getEnv().asHostObject(hostRunnable);
@@ -1411,4 +1412,13 @@ protected long handleCount() {
         }
     }
 
+    @CoreMethod(names = "multithreaded?", onSingleton = true)
+    public abstract static class IsMultiThreadedNode extends CoreMethodArrayArgumentsNode {
+
+        @Specialization
+        protected boolean isMultiThreaded() {
+            return getLanguage().isMultiThreaded();
+        }
+    }
+
 }

src/test/java/org/truffleruby/ContextPermissionsTest.java

@@ -16,6 +16,7 @@
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 public class ContextPermissionsTest {
@@ -88,6 +89,7 @@ public void testThreadsNoNative() throws Throwable {
             Assert.assertEquals(3, context.eval("ruby", "1 + 2").asInt());
 
             Assert.assertEquals(7, context.eval("ruby", "Thread.new { 3 + 4 }.value").asInt());
+            assertTrue(context.eval("ruby", "Truffle::Debug.multithreaded?").asBoolean());
 
             RubyTest.assertThrows(
                     () -> context.eval("ruby", "File.stat('.')"),
@@ -118,8 +120,29 @@ public void testNoThreadsEnforcesSingleThreadedOption() throws Throwable {
     }
 
     @Test
-    public void testFiberDoesNotTriggerMultiThreading() {
+    public void testNoThreads() {
         try (Context context = Context.newBuilder("ruby").allowCreateThread(false).build()) {
+            RubyTest.assertThrows(
+                    () -> context.eval("ruby", "Thread.new {}.join"),
+                    e -> {
+                        assertEquals("threads not allowed in single-threaded mode", e.getMessage());
+                        assertEquals("SecurityError", e.getGuestObject().getMetaObject().getMetaQualifiedName());
+                    });
+
+            RubyTest.assertThrows(
+                    () -> context.eval("ruby", "Fiber.new {}.resume"),
+                    e -> {
+                        assertEquals("fibers not allowed with allowCreateThread(false)", e.getMessage());
+                        assertEquals("SecurityError", e.getGuestObject().getMetaObject().getMetaQualifiedName());
+                    });
+
+            assertFalse(context.eval("ruby", "Truffle::Debug.multithreaded?").asBoolean());
+        }
+    }
+
+    @Test
+    public void testFiberDoesNotTriggerMultiThreading() {
+        try (Context context = Context.newBuilder("ruby").allowCreateThread(true).build()) {
             final Value array = context.eval(
                     "ruby",
                     "a = [1]; f = Fiber.new { a << 3; Fiber.yield; a << 5 }; a << 2; f.resume; a << 4; f.resume");
@@ -128,12 +151,14 @@ public void testFiberDoesNotTriggerMultiThreading() {
             for (int i = 0; i < 5; i++) {
                 assertEquals(i + 1, array.getArrayElement(i).asInt());
             }
+
+            assertFalse(context.eval("ruby", "Truffle::Debug.multithreaded?").asBoolean());
         }
     }
 
     @Test
     public void testNestedFiberAndTerminateFiber() {
-        try (Context context = Context.newBuilder("ruby").allowCreateThread(false).build()) {
+        try (Context context = Context.newBuilder("ruby").allowCreateThread(true).build()) {
             final Value array = context.eval(
                     "ruby",
                     "a = []; Fiber.new { a << 1; Fiber.new { a << 2; Fiber.yield; unreachable }.resume; a << 3 }.resume");
@@ -142,6 +167,8 @@ public void testNestedFiberAndTerminateFiber() {
             for (int i = 0; i < 3; i++) {
                 assertEquals(i + 1, array.getArrayElement(i).asInt());
             }
+
+            assertFalse(context.eval("ruby", "Truffle::Debug.multithreaded?").asBoolean());
         }
     }