jruby · Aug 31, 2015 · Aug 31, 2015 · Aug 31, 2015 · Aug 31, 2015 · Aug 31, 2015
diff --git a/spec/ruby/core/objectspace/each_object_spec.rb b/spec/ruby/core/objectspace/each_object_spec.rb
@@ -1,4 +1,7 @@
 require File.expand_path('../../../spec_helper', __FILE__)
+require File.expand_path('../fixtures', __FILE__)
+
+require 'weakref'
 
 describe "ObjectSpace.each_object" do
   it "calls the block once for each living, non-immediate object in the Ruby process" do
@@ -43,4 +46,127 @@ class ObjectSpaceSpecEachOtherObject; end
     # this is needed to prevent the new_obj from being GC'd too early
     new_obj.should_not == nil
   end
+
+  it "finds an object stored in a global variable" do
+    $object_space_global_variable = ObjectSpaceFixtures::ObjectToBeFound.new(:global)
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:global)
+  end
+
+  it "finds an object stored in a top-level constant" do
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:top_level_constant)
+  end
+
+  it "finds an object stored in a top-level constant" do
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:second_level_constant)
+  end
+
+  it "finds an object stored in a local variable" do
+    local = ObjectSpaceFixtures::ObjectToBeFound.new(:local)
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:local)
+  end
+
+  it "finds an object stored in a local variable captured in a block explicitly" do
+    proc = Proc.new {
+      local_in_block = ObjectSpaceFixtures::ObjectToBeFound.new(:local_in_block_explicit)
+      Proc.new { local_in_block }
+    }.call
+
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:local_in_block_explicit)
+  end
+
+  it "finds an object stored in a local variable captured in a block implicitly" do
+    proc = Proc.new {
+      local_in_block = ObjectSpaceFixtures::ObjectToBeFound.new(:local_in_block_implicit)
+      Proc.new { }
+    }.call
+
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:local_in_block_implicit)
+  end
+
+  it "finds an object stored in a local variable captured in a Proc#binding" do
+    binding = Proc.new {
+      local_in_proc_binding = ObjectSpaceFixtures::ObjectToBeFound.new(:local_in_proc_binding)
+      Proc.new { }.binding
+    }.call
+
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:local_in_proc_binding)
+  end
+
+  it "finds an object stored in a local variable captured in a Kernel#binding" do
+    b = Proc.new {
+      local_in_kernel_binding = ObjectSpaceFixtures::ObjectToBeFound.new(:local_in_kernel_binding)
+      binding
+    }.call
+
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:local_in_kernel_binding)
+  end
+
+  it "finds an object stored in a local variable set in a binding manually" do
+    b = binding
+    b.local_variable_set(:local, ObjectSpaceFixtures::ObjectToBeFound.new(:local_in_manual_binding))
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:local_in_manual_binding)
+  end
+
+  it "finds an object stored in an array" do
+    array = [ObjectSpaceFixtures::ObjectToBeFound.new(:array)]
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:array)
+  end
+
+  it "finds an object stored in a hash key" do
+    hash = {ObjectSpaceFixtures::ObjectToBeFound.new(:hash_key) => :value}
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:hash_key)
+  end
+
+  it "finds an object stored in a hash value" do
+    hash = {a: ObjectSpaceFixtures::ObjectToBeFound.new(:hash_value)}
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:hash_value)
+  end
+
+  it "finds an object stored in an instance variable" do
+    local = ObjectSpaceFixtures::ObjectWithInstanceVariable.new
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:instance_variable)
+  end
+
+  it "doesn't find an object stored in a WeakRef that should have been cleared" do
+    weak_ref = WeakRef.new(ObjectSpaceFixtures::ObjectToBeFound.new(:weakref))
+    ObjectSpaceFixtures::to_be_found_symbols.should_not include(:weakref)
+  end
+
+  it "finds an object stored in a thread local" do
+    Thread.current.thread_variable_set(:object_space_thread_local, ObjectSpaceFixtures::ObjectToBeFound.new(:thread_local))
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:thread_local)
+  end
+
+  it "finds an object stored in a fiber local" do
+    Thread.current[:object_space_fiber_local] = ObjectSpaceFixtures::ObjectToBeFound.new(:fiber_local)
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:fiber_local)
+  end
+
+  it "finds an object captured in an at_exit handler" do
+    Proc.new {
+      local = ObjectSpaceFixtures::ObjectToBeFound.new(:at_exit)
+
+      at_exit do
+        local
+      end
+    }.call
+
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:at_exit)
+  end
+
+  it "finds an object captured in finalizer" do
+    alive = Object.new
+
+    Proc.new {
+      local = ObjectSpaceFixtures::ObjectToBeFound.new(:at_exit)
+
+      ObjectSpace.define_finalizer(alive, Proc.new {
+        local
+      })
+    }.call
+
+    ObjectSpaceFixtures::to_be_found_symbols.should include(:at_exit)
+
+    alive.should_not be_nil
+  end
 end
diff --git a/spec/ruby/core/objectspace/fixtures.rb b/spec/ruby/core/objectspace/fixtures.rb
@@ -32,4 +32,29 @@ def self.scoped(wr)
     return Proc.new { wr.write "finalized"; wr.close }
   end
 
+  class ObjectToBeFound
+    attr_reader :id
+
+    def initialize(id)
+      @id = id
+    end
+  end
+
+  class ObjectWithInstanceVariable
+    def initialize
+      @instance_variable = ObjectToBeFound.new(:instance_variable)
+    end
+  end
+
+  def self.to_be_found_symbols
+    GC.start
+    ObjectSpace.each_object(ObjectToBeFound).map do |o|
+      o.id
+    end
+  end
+
+  SECOND_LEVEL_CONSTANT = ObjectToBeFound.new(:second_level_constant)
+
 end
+
+OBJECT_SPACE_TOP_LEVEL_CONSTANT = ObjectSpaceFixtures::ObjectToBeFound.new(:top_level_constant)
diff --git a/truffle/src/main/java/org/jruby/truffle/runtime/object/ObjectGraph.java b/truffle/src/main/java/org/jruby/truffle/runtime/object/ObjectGraph.java
@@ -14,9 +14,14 @@
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.object.DynamicObject;
 import com.oracle.truffle.api.object.Property;
+import org.jruby.truffle.nodes.RubyGuards;
+import org.jruby.truffle.runtime.RubyArguments;
+import org.jruby.truffle.runtime.RubyConstant;
 import org.jruby.truffle.runtime.RubyContext;
+import org.jruby.truffle.runtime.core.ModuleFields;
 import org.jruby.truffle.runtime.hash.Entry;
 import org.jruby.truffle.runtime.layouts.Layouts;
+import org.jruby.truffle.runtime.methods.InternalMethod;
 import org.jruby.truffle.runtime.subsystems.SafepointAction;
 
 import java.util.HashSet;
@@ -74,13 +79,22 @@ public void run(DynamicObject thread, Node currentNode) {
                     }
 
                     try {
-                        // We only visit the global variables from the root thread
-                        visitObject(context.getCoreLibrary().getGlobalVariablesObject(), visitor);
+                        // We only visit the global variables and other global state from the root thread
+
+                        if (thread == context.getThreadManager().getRootThread()) {
+                            visitObject(context.getCoreLibrary().getGlobalVariablesObject(), visitor);
+
+                            for (DynamicObject handler : context.getAtExitManager().getHandlers()) {
+                                visitObject(handler, visitor);
+                            }
+                        }
+
+                        // All threads visit their thread object
 
-                        // All threads visit the thread object
                         visitObject(thread, visitor);
 
                         // All threads visit their call stack
+
                         if (Truffle.getRuntime().getCurrentFrame() != null) {
                             visitFrameInstance(Truffle.getRuntime().getCurrentFrame(), visitor);
                         }
@@ -115,9 +129,23 @@ private void visitObject(DynamicObject object, ObjectGraphVisitor visitor) throw
 
             // Visit all properties
 
-            for (Property property : object.getShape().getProperties()) {
+            for (Property property : object.getShape().getPropertyListInternal(false)) {
                 visitObject(property.get(object, object.getShape()), visitor);
             }
+
+            // Visit specific objects that we're managing
+
+            if (RubyGuards.isRubyModule(object)) {
+                final ModuleFields module = Layouts.MODULE.getFields(object);
+
+                for (DynamicObject ancestor : module.ancestors()) {
+                    visitObject(ancestor, visitor);
+                }
+
+                for (RubyConstant constant : module.getConstants().values()) {
+                    visitObject(constant.getValue(), visitor);
+                }
+            }
         }
     }
 
@@ -133,8 +161,11 @@ private void visitObject(Object object, ObjectGraphVisitor visitor) throws StopV
             visitObject(entry.getKey(), visitor);
             visitObject(entry.getValue(), visitor);
             visitObject(entry.getNextInLookup(), visitor);
-        } else if (object instanceof MaterializedFrame) {
-            visitFrame((MaterializedFrame) object, visitor);
+        } else if (object instanceof Frame) {
+            visitFrame((Frame) object, visitor);
+        } else if (object instanceof InternalMethod) {
+            final InternalMethod method = (InternalMethod) object;
+            visitObject(method.getDeclarationFrame(), visitor);
         }
     }
 
@@ -146,5 +177,17 @@ private void visitFrame(Frame frame, ObjectGraphVisitor visitor) throws StopVisi
         for (FrameSlot slot : frame.getFrameDescriptor().getSlots()) {
             visitObject(frame.getValue(slot), visitor);
         }
+
+        Frame declarationFrame;
+
+        try {
+            declarationFrame = RubyArguments.getDeclarationFrame(frame.getArguments());
+        } catch (Exception e) {
+            declarationFrame = null;
+        }
+
+        if (declarationFrame != null) {
+            visitFrame(declarationFrame, visitor);
+        }
     }
 }
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
@@ -17,7 +17,9 @@
 import org.jruby.truffle.runtime.control.RaiseException;
 import org.jruby.truffle.runtime.layouts.Layouts;
 
+import java.util.ArrayList;
 import java.util.Deque;
+import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.concurrent.ConcurrentLinkedDeque;
 
@@ -72,4 +74,11 @@ private void runExitHooks(Deque<DynamicObject> stack) {
         }
     }
 
+    public List<DynamicObject> getHandlers() {
+        final List<DynamicObject> handlers = new ArrayList<>();
+        handlers.addAll(runOnExit);
+        handlers.addAll(runOnExitAlways);
+        return handlers;
+    }
+
 }
diff --git a/truffle/src/main/ruby/core/thread.rb b/truffle/src/main/ruby/core/thread.rb
@@ -16,6 +16,13 @@ def []=(symbol, value)
     __thread_local_variables[symbol] = value
   end
 
+  def thread_variable?(symbol)
+    __thread_local_variables.has_key? symbol
+  end
+
+  alias_method :thread_variable_get, :[]
+  alias_method :thread_variable_set, :[]=
+
   def __thread_local_variables
     @__thread_local_variables ||= {}
   end