[GR-33681] [GR-33546] Backport: Fix setting the current RubyThread for fibers before TruffleContext#enter() which polls

PullRequest: truffleruby/2962

Author: eregon

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

@@ -51,6 +51,7 @@
 import org.truffleruby.core.exception.RubySyntaxError;
 import org.truffleruby.core.exception.RubySystemCallError;
 import org.truffleruby.core.exception.RubySystemExit;
+import org.truffleruby.core.fiber.FiberPoolThread;
 import org.truffleruby.core.fiber.RubyFiber;
 import org.truffleruby.core.hash.RubyHash;
 import org.graalvm.options.OptionValues;
@@ -177,12 +178,14 @@ public final class RubyLanguage extends TruffleLanguage<RubyContext> {
 
     /** We need an extra indirection added to ContextThreadLocal due to multiple Fibers of different Ruby Threads
      * sharing the same Java Thread when using the fiber pool. */
-    private static final class ThreadLocalState {
-        private RubyThread rubyThread;
+    public static final class ThreadLocalState {
+        public RubyThread rubyThread;
     }
 
     private final ContextThreadLocal<ThreadLocalState> threadLocalState = createContextThreadLocal(
-            (context, thread) -> new ThreadLocalState());
+            (context, thread) -> thread instanceof FiberPoolThread
+                    ? ((FiberPoolThread) thread).threadLocalState
+                    : new ThreadLocalState());
 
     private final CyclicAssumption tracingCyclicAssumption = new CyclicAssumption("object-space-tracing");
     @CompilationFinal private volatile Assumption tracingAssumption = tracingCyclicAssumption.getAssumption();

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

@@ -93,7 +93,7 @@ private void fiberMain(RubyContext context, RubyFiber fiber, RubyProc block, Nod
         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 FiberPoolThread thread = (FiberPoolThread) Thread.currentThread();
         final SourceSection sourceSection = block.sharedMethodInfo.getSourceSection();
         final String oldName = thread.getName();
         thread.setName(NAME_PREFIX + " id=" + thread.getId() + " from " + RubyLanguage.fileLine(sourceSection));
@@ -106,9 +106,14 @@ private void fiberMain(RubyContext context, RubyFiber fiber, RubyProc block, Nod
         final FiberMessage message = waitMessage(fiber, currentNode);
         final TruffleContext truffleContext = context.getEnv().getContext();
 
-        final Object prev = truffleContext.enter(currentNode);
-        language.setupCurrentThread(thread, fiber.rubyThread);
+        /* We need to setup the ThreadLocalState before enter(), as that polls and the current thread is needed for
+         * guest safepoints. It can't just be set in RubyLanguage#initializeThread() as this java.lang.Thread is reused
+         * for multiple Fibers and multiple Ruby Threads due to the Fiber pool, but initializeThread() is only called
+         * once per thread. */
         fiber.rubyThread.setCurrentFiber(fiber);
+        thread.threadLocalState.rubyThread = fiber.rubyThread;
+
+        final Object prev = truffleContext.enter(currentNode);
 
         FiberMessage lastMessage = null;
         try {
@@ -141,8 +146,8 @@ private void fiberMain(RubyContext context, RubyFiber fiber, RubyProc block, Nod
             // Make sure that other fibers notice we are dead before they gain control back
             fiber.alive = false;
             // Leave context before addToMessageQueue() -> parent Fiber starts executing
-            language.setupCurrentThread(thread, null);
             truffleContext.leave(currentNode, prev);
+            thread.threadLocalState.rubyThread = null;
             cleanup(fiber, thread);
             thread.setName(oldName);
 

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2021 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.core.fiber;
+
+import org.truffleruby.RubyLanguage.ThreadLocalState;
+
+public final class FiberPoolThread extends Thread {
+
+    public final ThreadLocalState threadLocalState;
+
+    public FiberPoolThread(Runnable target) {
+        super(target);
+        threadLocalState = new ThreadLocalState();
+    }
+
+}

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

@@ -40,6 +40,7 @@
 import org.truffleruby.core.exception.RubyException;
 import org.truffleruby.core.exception.RubySystemExit;
 import org.truffleruby.core.fiber.FiberManager;
+import org.truffleruby.core.fiber.FiberPoolThread;
 import org.truffleruby.core.fiber.RubyFiber;
 import org.truffleruby.core.klass.RubyClass;
 import org.truffleruby.core.string.StringUtils;
@@ -155,7 +156,7 @@ private Thread createFiberJavaThread(Runnable runnable) {
             throw new UnsupportedOperationException("fibers should not be created while pre-initializing the context");
         }
 
-        final Thread thread = new Thread(runnable); // context.getEnv().createUnenteredThread(runnable);
+        final Thread thread = new FiberPoolThread(runnable); // context.getEnv().createUnenteredThread(runnable);
         thread.setName("Ruby-FiberPool-" + thread.getName());
         thread.setDaemon(true); // GR-33255
         rubyManagedThreads.add(thread); // need to be set before initializeThread()

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

@@ -58,17 +58,19 @@ protected final void perform(Access access) {
                             Thread.currentThread());
         }
 
+        final RubyLanguage language = RubyLanguage.getCurrentLanguage();
         final RubyContext context = RubyLanguage.getCurrentContext();
-        // Not using language.getCurrentThread() here because it might not be set yet, see fiberMain()
-        final RubyThread rubyThread = context.getThreadManager().getRubyThreadForJavaThread(access.getThread());
+
+        final RubyThread rubyThread = language.getCurrentThread();
+        final Node node = access.getLocation();
         if (filter.test(context, rubyThread, this)) {
-            run(rubyThread, access.getLocation());
+            run(rubyThread, node);
 
             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());
+                        context.fiberManager.safepoint(currentFiber, fiber, this, node);
                     }
                 }
             }