Make the top-level exception handler more compatible with MRI

* Fixes #2047

Author: eregon

CHANGELOG.md

@@ -65,6 +65,7 @@ Compatibility:
 * Implemented `ONIGENC_MBC_CASE_FOLD`.
 * Fixed `Thread#raise` to call the exception class' constructor with no arguments when given no message (#2045).
 * Fixed `refine + super` compatibility (#2039, @ssnickolay)
+* Make the top-level exception handler more compatible with MRI (#2047).
 
 Performance:
 

src/main/java/org/truffleruby/core/kernel/AtExitManager.java

@@ -17,6 +17,7 @@
 import org.truffleruby.Layouts;
 import org.truffleruby.RubyContext;
 import org.truffleruby.core.proc.ProcOperations;
+import org.truffleruby.language.control.ExitException;
 import org.truffleruby.language.control.RaiseException;
 
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
@@ -42,15 +43,15 @@ public void add(DynamicObject block, boolean always) {
     }
 
     public DynamicObject runAtExitHooks() {
-        return runExitHooks(atExitHooks);
+        return runExitHooks(atExitHooks, "at_exit");
     }
 
     public void runSystemExitHooks() {
-        runExitHooks(systemExitHooks);
+        runExitHooks(systemExitHooks, "system at_exit");
     }
 
     @TruffleBoundary
-    private DynamicObject runExitHooks(Deque<DynamicObject> stack) {
+    private DynamicObject runExitHooks(Deque<DynamicObject> stack, String name) {
         DynamicObject lastException = null;
 
         while (true) {
@@ -62,8 +63,12 @@ private DynamicObject runExitHooks(Deque<DynamicObject> stack) {
             try {
                 ProcOperations.rootCall(block);
             } catch (RaiseException e) {
-                lastException = handleAtExitException(context, e);
+                handleAtExitException(context, e.getException());
+                lastException = e.getException();
+            } catch (ExitException e) {
+                throw e;
             } catch (Exception e) {
+                System.err.println("Unexpected internal exception in " + name + ":");
                 e.printStackTrace();
             }
         }
@@ -77,16 +82,14 @@ public List<DynamicObject> getHandlers() {
     }
 
     @TruffleBoundary
-    public static DynamicObject handleAtExitException(RubyContext context, RaiseException raiseException) {
-        final DynamicObject rubyException = raiseException.getException();
+    public static void handleAtExitException(RubyContext context, DynamicObject rubyException) {
         DynamicObject logicalClass = Layouts.BASIC_OBJECT.getLogicalClass(rubyException);
         if (logicalClass == context.getCoreLibrary().systemExitClass ||
                 logicalClass == context.getCoreLibrary().signalExceptionClass) {
             // Do not show the backtrace for these
         } else {
             context.getDefaultBacktraceFormatter().printRubyExceptionOnEnvStderr("", rubyException);
         }
-        return rubyException;
     }
 
 }

src/main/java/org/truffleruby/language/exceptions/TopLevelRaiseHandler.java

@@ -38,25 +38,34 @@ public Object execute(VirtualFrame frame) {
         int exitCode = 0;
         DynamicObject caughtException = null;
 
+        // Execute the main script
         try {
             body.execute(frame);
         } catch (RaiseException e) {
-            DynamicObject rubyException = AtExitManager.handleAtExitException(getContext(), e);
-            caughtException = rubyException;
-            setLastException(rubyException);
-            exitCode = statusFromException(rubyException);
+            caughtException = e.getException();
+            exitCode = statusFromException(caughtException);
+            setLastException(caughtException); // Set $! for at_exit
+            // printing the main script exception is delayed after at_exit hooks
         } catch (ExitException e) {
-            exitCode = e.getCode();
-        } finally {
-            final DynamicObject atExitException = getContext().getAtExitManager().runAtExitHooks();
+            // hard #exit!, return immediately, skip at_exit hooks
+            return e.getCode();
+        }
 
+        // Execute at_exit hooks (except if hard #exit!)
+        try {
+            DynamicObject atExitException = getContext().getAtExitManager().runAtExitHooks();
             if (atExitException != null) {
                 exitCode = statusFromException(atExitException);
             }
 
             if (caughtException != null) {
+                // print the main script exception now
+                AtExitManager.handleAtExitException(getContext(), caughtException);
                 handleSignalException(caughtException);
             }
+        } catch (ExitException e) {
+            // hard #exit! during at_exit: ignore the main script exception
+            exitCode = e.getCode();
         }
 
         return exitCode;