Fixes #3522. END not working as in MRI

jruby · May 23, 2017 · 154275c · 154275c
1 parent 800dee2
commit 154275c
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/core/src/main/java/org/jruby/Ruby.java b/core/src/main/java/org/jruby/Ruby.java
@@ -3121,6 +3121,26 @@ public IRubyObject pushExitBlock(RubyProc proc) {
         return proc;
     }
 
+    /**
+     * It is possible for looping or repeated execution to encounter the same END
+     * block multiple times.  Rather than store extra runtime state we will just
+     * make sure it is not already registered.  at_exit by contrast can push the
+     * same block many times (and should use pushExistBlock).
+     */
+    public void pushEndBlock(RubyProc proc) {
+        if (alreadyRegisteredEndBlock(proc) != null) return;
+        pushExitBlock(proc);
+    }
+
+    private RubyProc alreadyRegisteredEndBlock(RubyProc newProc) {
+        Block block = newProc.getBlock();
+
+        for (RubyProc proc: atExitBlocks) {
+            if (block.equals(proc.getBlock())) return proc;
+        }
+        return null;
+    }
+
     // use this for JRuby-internal finalizers
     public void addInternalFinalizer(Finalizable finalizer) {
         synchronized (internalFinalizersMutex) {

diff --git a/core/src/main/java/org/jruby/ir/runtime/IRRuntimeHelpers.java b/core/src/main/java/org/jruby/ir/runtime/IRRuntimeHelpers.java
@@ -1555,7 +1555,7 @@ public static RaiseException newRequiredKeywordArgumentError(ThreadContext conte
 
     @JIT
     public static void pushExitBlock(ThreadContext context, Block blk) {
-        context.runtime.pushExitBlock(context.runtime.newProc(Block.Type.LAMBDA, blk));
+        context.runtime.pushEndBlock(context.runtime.newProc(Block.Type.LAMBDA, blk));
     }
 
     @JIT