Ensure leaveAndEnter() never returns null

Author: eregon

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

@@ -15,7 +15,6 @@
 import com.oracle.truffle.api.TruffleSafepoint.Interrupter;
 import org.truffleruby.core.fiber.RubyFiber.FiberStatus;
 
-import com.oracle.truffle.api.TruffleContext;
 import com.oracle.truffle.api.TruffleSafepoint;
 import org.truffleruby.RubyContext;
 import org.truffleruby.RubyLanguage;
@@ -75,9 +74,8 @@ private void createThreadToReceiveFirstMessage(RubyFiber fiber, Node currentNode
         fiber.initializeNode = null;
 
         var sourceSection = block.getSharedMethodInfo().getSourceSection();
-        final TruffleContext truffleContext = context.getEnv().getContext();
 
-        truffleContext.leaveAndEnter(currentNode, Interrupter.THREAD_INTERRUPT, (unused) -> {
+        context.getThreadManager().leaveAndEnter(currentNode, Interrupter.THREAD_INTERRUPT, (unused) -> {
             Thread thread = context.getThreadManager().createFiberJavaThread(fiber, sourceSection,
                     () -> beforeEnter(fiber, initializeNode),
                     () -> fiberMain(context, fiber, block, initializeNode),
@@ -226,7 +224,7 @@ private void addToMessageQueue(RubyFiber fiber, FiberMessage message) {
     @TruffleBoundary
     private FiberMessage waitMessage(RubyFiber fiber, Node currentNode) throws InterruptedException {
         assertNotEntered("should have left context while waiting fiber message");
-        return fiber.messageQueue.take();
+        return Objects.requireNonNull(fiber.messageQueue.take());
     }
 
     private void assertNotEntered(String reason) {
@@ -317,8 +315,7 @@ private FiberMessage resumeAndWait(RubyFiber fromFiber, RubyFiber toFiber, Fiber
             context.fiberManager.createThreadToReceiveFirstMessage(toFiber, currentNode);
         }
 
-        final TruffleContext truffleContext = context.getEnv().getContext();
-        final FiberMessage message = truffleContext.leaveAndEnter(currentNode, Interrupter.THREAD_INTERRUPT,
+        var message = context.getThreadManager().leaveAndEnter(currentNode, Interrupter.THREAD_INTERRUPT,
                 (unused) -> {
                     resume(fromFiber, toFiber, operation, descriptor, args);
                     return waitMessage(fromFiber, currentNode);
@@ -329,8 +326,7 @@ private FiberMessage resumeAndWait(RubyFiber fromFiber, RubyFiber toFiber, Fiber
 
     @TruffleBoundary
     public void safepoint(RubyFiber fromFiber, RubyFiber fiber, SafepointAction action, Node currentNode) {
-        final TruffleContext truffleContext = context.getEnv().getContext();
-        final FiberResumeMessage returnMessage = (FiberResumeMessage) truffleContext.leaveAndEnter(currentNode,
+        var returnMessage = (FiberResumeMessage) context.getThreadManager().leaveAndEnter(currentNode,
                 Interrupter.THREAD_INTERRUPT, (unused) -> {
                     addToMessageQueue(fiber, new FiberSafepointMessage(fromFiber, action));
                     return waitMessage(fromFiber, currentNode);
@@ -385,8 +381,7 @@ public void killOtherFibers(RubyThread thread) {
         final TruffleSafepoint safepoint = TruffleSafepoint.getCurrent();
         boolean allowSideEffects = safepoint.setAllowSideEffects(false);
         try {
-            final TruffleContext truffleContext = context.getEnv().getContext();
-            truffleContext.leaveAndEnter(DummyNode.INSTANCE, Interrupter.THREAD_INTERRUPT, (unused) -> {
+            context.getThreadManager().leaveAndEnter(DummyNode.INSTANCE, Interrupter.THREAD_INTERRUPT, (unused) -> {
                 doKillOtherFibers(thread);
                 return BlockingAction.SUCCESS;
             }, null);

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

@@ -11,6 +11,7 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.Objects;
 import java.util.Set;
 import java.util.Timer;
 import java.util.concurrent.ConcurrentHashMap;
@@ -23,6 +24,7 @@
 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;
 import com.oracle.truffle.api.TruffleSafepoint;
 import com.oracle.truffle.api.TruffleSafepoint.Interrupter;
@@ -607,6 +609,15 @@ public <T> T runUntilResult(Node currentNode, BlockingAction<T> action, Runnable
         }
     }
 
+    /** Same as {@link TruffleContext#leaveAndEnter} but ensures it never returns null. Also the interruptible must
+     * never return null. */
+    public <T, R> R leaveAndEnter(Node node, Interrupter interrupter, InterruptibleFunction<T, R> interruptible,
+            T object) {
+        final TruffleContext truffleContext = context.getEnv().getContext();
+        R result = truffleContext.leaveAndEnter(node, interrupter, interruptible, object);
+        return Objects.requireNonNull(result);
+    }
+
     @TruffleBoundary
     Interrupter getNativeCallInterrupter() {
         if (nativeInterrupt) {