lib/sqreen/instrumentation.rb in sqreen-alt-1.10.3 vs lib/sqreen/instrumentation.rb in sqreen-alt-1.10.4

@@ -174,14 +174,33 @@
         end
         returns
       end
     end
 
+    def self.guard_multi_call(instance, method, original_method, args, block)
+      @sqreen_multi_instr ||= nil
+      key = [method]
+      Instrumentation.guard_call(nil, :guard_multi_call) do
+        args.each{|e| key.push(e.object_id)}
+      end
+      if key && @sqreen_multi_instr && @sqreen_multi_instr[instance.object_id].member?(key)
+        return instance.send(original_method, *args, &block)
+      end
+      @sqreen_multi_instr ||= Hash.new {|h, k| h[k]=Set.new } # TODO this should probably be a thread local
+      @sqreen_multi_instr[instance.object_id].add(key)
+      r = yield
+      return r
+    ensure
+      if @sqreen_multi_instr && @sqreen_multi_instr[instance.object_id] && @sqreen_multi_instr[instance.object_id].delete(key).empty?
+        @sqreen_multi_instr.delete(instance.object_id)
+      end
+    end
+
     def self.guard_call(method, retval)
       @sqreen_in_instr ||= nil
       return retval if @sqreen_in_instr && @sqreen_in_instr.member?(method)
-      @sqreen_in_instr ||= Set.new
+      @sqreen_in_instr ||= Set.new # TODO this should probably be a thread local
       @sqreen_in_instr.add(method)
       r = yield
       @sqreen_in_instr.delete(method)
       return r
     rescue Exception => e
@@ -195,10 +214,11 @@
           Sqreen.log.debug do
             "Instrumented #{Instrumentation.instrumented_pid} != PID #{Process.pid}"
           end
           return send(original_meth, *args, &block)
         end
+        Instrumentation.guard_multi_call(self, meth, original_meth, args, block) do
         Sqreen.stats.callbacks_calls += 1
 
         skip = false
         result = nil
 
@@ -269,10 +289,11 @@
             end
           end
           result
         end
       end
+      end
     end
 
     def override_class_method(klass, meth)
       # FIXME: This is somehow ugly. We should reduce the amount of
       # `evaled` code.
@@ -323,12 +344,12 @@
         alias_method meth, saved_meth_name
         send(method_kind, meth)
       end
     end
 
-    def get_saved_method_name(meth)
-      "#{meth}_not_modified".to_sym
+    def get_saved_method_name(meth, suffix=nil)
+      "#{meth}_sq#{suffix}_not_modified".to_sym
     end
 
     def override_instance_method(klass_name, meth)
       saved_meth_name = get_saved_method_name(meth)
       new_method      = "#{meth}_modified".to_sym
@@ -401,10 +422,34 @@
       return true if is_class_method?(obj, method)
       return false unless obj.respond_to?(:instance_methods)
       is_instance_method?(obj, method)
     end
 
+    # Override a singleton method on an instance
+    def override_singleton_method(instance, klass_name, meth)
+      saved_meth_name = get_saved_method_name(meth, 'singleton')
+      if instance.respond_to?(saved_meth_name, true)
+        Sqreen.log.debug { "#{saved_meth_name} found #{instance.class}##{instance.object_id} already instrumented" }
+        return nil
+      elsif instance.frozen?
+        Sqreen.log.debug { "#{instance.class}##{instance.object_id} is frozen, not reinstrumenting" }
+        return nil
+      end
+      raise Sqreen::NotImplementedYet, "#{instance.inspect} doesn't respond to define_singleton_method" unless instance.respond_to?(:define_singleton_method)
+      p = Instrumentation.define_callback_method(meth, saved_meth_name,
+                                                 klass_name)
+      instance.define_singleton_method(saved_meth_name, instance.method(meth))
+      instance.define_singleton_method(meth, p)
+      # Hide saved method (its only available in this syntax)
+      eval <<-RUBY, binding, __FILE__, __LINE__ + 1
+      class << instance
+        private :#{saved_meth_name}
+      end
+      saved_meth_name
+      RUBY
+    end
+
     def add_callback(cb)
       @@override_semaphore.synchronize do
         klass = cb.klass
         method = cb.method
         key = [klass, method]
@@ -425,16 +470,28 @@
             #  - define_method
             #  - method_added
             #  - method_missing
             #  ...
             #
-            msg = "#{cb} is neither singleton or instance"
+            msg = "#{cb} is neither class or instance"
             raise Sqreen::NotImplementedYet, msg
           end
 
           @@overriden_methods += [key] if success
         else
           Sqreen.log.debug "#{key} was already overriden"
+        end
+
+        if klass != Object && klass != Kernel && !Sqreen.features['instrument_all_instances'] && !defined?(::JRUBY_VERSION)
+          insts = 0
+          ObjectSpace.each_object(klass) do |e|
+            next if e.is_a?(Class) || e.is_a?(Module)
+            next unless e.singleton_methods.include?(method.to_sym)
+            insts += 1 if override_singleton_method(e, klass, method)
+          end
+          if insts > 0
+            Sqreen.log.debug { "Reinstrumented #{insts} instances of #{klass}" }
+          end
         end
 
         @@registered_callbacks.add(cb)
         @@instrumented_pid = Process.pid
       end