Extract ThreadManager#leaveAndEnter for convenience

Author: eregon

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

@@ -246,22 +246,12 @@ private void resume(RubyFiber fromFiber, RubyFiber fiber, FiberOperation operati
     @TruffleBoundary
     public Object[] transferControlTo(RubyFiber fromFiber, RubyFiber fiber, FiberOperation operation, Object[] args) {
         final TruffleContext truffleContext = context.getEnv().getContext();
-
-        final FiberMessage message;
         final boolean isRubyManagedThread = context.getThreadManager().isRubyManagedThread(Thread.currentThread());
-        if (isRubyManagedThread) {
-            context.getSafepointManager().leaveThread();
-        }
-        try {
-            message = truffleContext.leaveAndEnter(null, () -> {
-                resume(fromFiber, fiber, operation, args);
-                return waitMessage(fromFiber);
-            });
-        } finally {
-            if (isRubyManagedThread) {
-                context.getSafepointManager().enterThread();
-            }
-        }
+
+        final FiberMessage message = context.getThreadManager().leaveAndEnter(truffleContext, null, () -> {
+            resume(fromFiber, fiber, operation, args);
+            return waitMessage(fromFiber);
+        }, isRubyManagedThread);
 
         return handleMessage(fromFiber, message);
     }
@@ -325,18 +315,10 @@ public void killOtherFibers() {
         // This method might not be executed on the rootFiber Java Thread but possibly on another Java Thread.
 
         final TruffleContext truffleContext = context.getEnv().getContext();
-        assert truffleContext.isEntered();
-        assert context.getThreadManager().isRubyManagedThread(Thread.currentThread());
-
-        context.getSafepointManager().leaveThread();
-        try {
-            truffleContext.leaveAndEnter(null, () -> {
-                doKillOtherFibers();
-                return null;
-            });
-        } finally {
-            context.getSafepointManager().enterThread();
-        }
+        context.getThreadManager().leaveAndEnter(truffleContext, null, () -> {
+            doKillOtherFibers();
+            return null;
+        }, true);
     }
 
     private void doKillOtherFibers() {

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

@@ -19,6 +19,7 @@
 import java.util.concurrent.locks.Lock;
 import java.util.function.Supplier;
 
+import com.oracle.truffle.api.TruffleContext;
 import com.oracle.truffle.api.object.DynamicObjectLibrary;
 import org.truffleruby.RubyContext;
 import org.truffleruby.RubyLanguage;
@@ -217,6 +218,8 @@ public void spawnFiber(RubyFiber fiber, Runnable task) {
         });
     }
 
+    /** Whether the thread was created by TruffleRuby. Also decides whether we use the {@link SafepointManager} on the
+     * thread. */
     @TruffleBoundary
     public boolean isRubyManagedThread(Thread thread) {
         return rubyManagedThreads.contains(thread);
@@ -451,6 +454,23 @@ public interface UnblockingAction {
         void unblock();
     }
 
+    public <T> T leaveAndEnter(TruffleContext truffleContext, Node currentNode, Supplier<T> runWhileOutsideContext,
+            boolean isRubyManagedThread) {
+        assert truffleContext.isEntered();
+        assert isRubyManagedThread == isRubyManagedThread(Thread.currentThread());
+
+        if (isRubyManagedThread) {
+            context.getSafepointManager().leaveThread();
+        }
+        try {
+            return truffleContext.leaveAndEnter(currentNode, runWhileOutsideContext);
+        } finally {
+            if (isRubyManagedThread) {
+                context.getSafepointManager().enterThread();
+            }
+        }
+    }
+
     /** Only use when no context is available. */
     @TruffleBoundary
     public static <T> T retryWhileInterrupted(BlockingAction<T> action) {