[Truffle] Add an method to add at_exit handlers that are always run.

jruby · May 12, 2015 · 6b59f27 · eregon · May 13, 2015 · 6b59f27
1 parent b89e4a8
commit 6b59f27
Showing 7 changed files with 64 additions and 38 deletions.
diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/core/KernelNodes.java b/truffle/src/main/java/org/jruby/truffle/nodes/core/KernelNodes.java
@@ -25,6 +25,7 @@
 import org.jcodings.specific.USASCIIEncoding;
 import org.jcodings.specific.UTF8Encoding;
 import org.jruby.RubyThread.Status;
+import org.jruby.exceptions.MainExitException;
 import org.jruby.truffle.nodes.RubyNode;
 import org.jruby.truffle.nodes.cast.BooleanCastNode;
 import org.jruby.truffle.nodes.cast.BooleanCastNodeGen;
@@ -229,27 +230,11 @@ public AbortNode(RubyContext context, SourceSection sourceSection) {
             super(context, sourceSection);
         }
 
+        @TruffleBoundary
         @Specialization
         public RubyBasicObject abort() {
-            CompilerDirectives.transferToInterpreter();
-            System.exit(1);
-            return nil();
-        }
-    }
-
-    @CoreMethod(names = "at_exit", isModuleFunction = true, needsBlock = true)
-    public abstract static class AtExitNode extends CoreMethodArrayArgumentsNode {
-
-        public AtExitNode(RubyContext context, SourceSection sourceSection) {
-            super(context, sourceSection);
-        }
-
-        @Specialization
-        public Object atExit(RubyProc block) {
-            CompilerDirectives.transferToInterpreter();
-
-            getContext().getAtExitManager().add(block);
-            return nil();
+            getContext().innerShutdown(false);
+            throw new MainExitException(1, true);
         }
     }
 

diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/core/TrufflePrimitiveNodes.java b/truffle/src/main/java/org/jruby/truffle/nodes/core/TrufflePrimitiveNodes.java
@@ -454,4 +454,19 @@ public RubyString hostOS() {
 
     }
 
+    @CoreMethod(names = "at_exit", isModuleFunction = true, needsBlock = true, required = 1)
+    public abstract static class AtExitSystemNode extends CoreMethodArrayArgumentsNode {
+
+        public AtExitSystemNode(RubyContext context, SourceSection sourceSection) {
+            super(context, sourceSection);
+        }
+
+        @CompilerDirectives.TruffleBoundary
+        @Specialization
+        public Object atExit(boolean always, RubyProc block) {
+            getContext().getAtExitManager().add(block, always);
+            return nil();
+        }
+    }
+
 }
diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/methods/ExceptionTranslatingNode.java b/truffle/src/main/java/org/jruby/truffle/nodes/methods/ExceptionTranslatingNode.java
@@ -15,6 +15,8 @@
 import com.oracle.truffle.api.nodes.ControlFlowException;
 import com.oracle.truffle.api.source.SourceSection;
 import com.oracle.truffle.api.utilities.BranchProfile;
+import org.jruby.exceptions.MainExitException;
+import org.jruby.exceptions.Unrescuable;
 import org.jruby.truffle.nodes.RubyNode;
 import org.jruby.truffle.runtime.DebugOperations;
 import org.jruby.truffle.runtime.RubyContext;
@@ -66,6 +68,9 @@ public Object execute(VirtualFrame frame) {
         } catch (RaiseException exception) {
             rethrowProfile.enter();
             throw exception;
+        } catch (MainExitException exception) {
+            CompilerDirectives.transferToInterpreter();
+            throw exception;
         } catch (ArithmeticException exception) {
             CompilerDirectives.transferToInterpreter();
             throw new RaiseException(translate(exception));

diff --git a/truffle/src/main/java/org/jruby/truffle/runtime/RubyContext.java b/truffle/src/main/java/org/jruby/truffle/runtime/RubyContext.java
@@ -16,7 +16,6 @@
 import com.oracle.truffle.api.instrument.Probe;
 import com.oracle.truffle.api.interop.TruffleObject;
 import com.oracle.truffle.api.nodes.Node;
-import com.oracle.truffle.api.object.Shape;
 import com.oracle.truffle.api.source.BytesDecoder;
 import com.oracle.truffle.api.source.Source;
 import com.oracle.truffle.api.source.SourceSection;
@@ -46,7 +45,6 @@
 import org.jruby.truffle.runtime.core.*;
 import org.jruby.truffle.runtime.methods.InternalMethod;
 import org.jruby.truffle.runtime.object.ObjectIDOperations;
-import org.jruby.truffle.runtime.object.RubyObjectType;
 import org.jruby.truffle.runtime.subsystems.*;
 import org.jruby.truffle.translator.NodeWrapper;
 import org.jruby.truffle.translator.TranslatorDriver;
@@ -360,8 +358,8 @@ public long getNextObjectID() {
         return id;
     }
 
-    public void innertShutdown() {
-        atExitManager.run();
+    public void innerShutdown(boolean normalExit) {
+        atExitManager.run(normalExit);
 
         if (instrumentationServerManager != null) {
             instrumentationServerManager.shutdown();
@@ -668,7 +666,7 @@ public RubyNode wrap(RubyNode node) {
     @Override
     public void shutdown() {
         try {
-            innertShutdown();
+            innerShutdown(true);
         } catch (RaiseException e) {
             final RubyException rubyException = e.getRubyException();
 

diff --git a/truffle/src/main/java/org/jruby/truffle/runtime/subsystems/AtExitManager.java b/truffle/src/main/java/org/jruby/truffle/runtime/subsystems/AtExitManager.java
@@ -11,28 +11,44 @@
 
 import org.jruby.truffle.runtime.core.RubyProc;
 
+import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.ListIterator;
 
-/**
- * Manages at_exit callbacks.
- */
 public class AtExitManager {
 
-    private final List<RubyProc> blocks = new LinkedList<>();
+    private final List<RubyProc> runOnExit = new LinkedList<>();
+    private final List<RubyProc> runOnExitAlways = new LinkedList<>();
 
-    public synchronized void add(RubyProc block) {
-        blocks.add(block);
+    public synchronized void add(RubyProc block, boolean always) {
+        if (always) {
+            runOnExitAlways.add(block);
+        } else {
+            runOnExit.add(block);
+        }
     }
 
-    public synchronized void run() {
-        final ListIterator<RubyProc> iterator = blocks.listIterator(blocks.size());
+    public synchronized void run(boolean normalExit) {
+        if (normalExit) {
+            Collections.reverse(runOnExit);
+
+            for (RubyProc block : runOnExit) {
+                try {
+                    block.rootCall();
+                } catch (Exception e) {
+                }
+            }
+        }
+
+        Collections.reverse(runOnExitAlways);
 
-        while (iterator.hasPrevious()) {
-            RubyProc hook = iterator.previous();
-            iterator.remove();
-            hook.rootCall();
+        for (RubyProc block : runOnExitAlways) {
+            try {
+                block.rootCall();
+            } catch (Exception e) {
+            }
         }
     }
+
 }
diff --git a/truffle/src/main/ruby/core/kernel.rb b/truffle/src/main/ruby/core/kernel.rb
@@ -28,4 +28,11 @@ def caller(start = 1, limit = nil)
   end
   module_function :caller
 
+  def at_exit(&block)
+    Truffle::Primitive.at_exit false, &block
+  end
+  private :at_exit
+  module_function :caller
+
+
 end
diff --git a/truffle/src/main/ruby/core/shims.rb b/truffle/src/main/ruby/core/shims.rb
@@ -37,15 +37,15 @@ def external_encoding
 if STDOUT.tty?
   STDOUT.sync = true
 else
-  at_exit do
+  Truffle::Primitive.at_exit true do
     STDOUT.flush
   end
 end
 
 if STDERR.tty?
   STDERR.sync = true
 else
-  at_exit do
+  Truffle::Primitive.at_exit true do
     STDERR.flush
   end
 end