diff --git a/lib/puppet/util/instrumentation/Instrumentable.rb b/lib/puppet/util/instrumentation/Instrumentable.rb new file mode 100644 index 000000000..5789dcbe2 --- /dev/null +++ b/lib/puppet/util/instrumentation/Instrumentable.rb @@ -0,0 +1,143 @@ +require 'monitor' +require 'puppet/util/instrumentation' + +# This is the central point of all declared probes. +# Every class needed to declare probes should include this module +# and declare the methods that are subject to instrumentation: +# +# class MyClass +# extend Puppet::Util::Instrumentation::Instrumentable +# +# probe :mymethod +# +# def mymethod +# ... this is code to be instrumented ... +# end +# end +module Puppet::Util::Instrumentation::Instrumentable + INSTRUMENTED_CLASSES = {}.extend(MonitorMixin) + + attr_reader :probes + + class Probe + attr_reader :klass, :method, :label, :data + + def initialize(method, klass, options = {}) + @method = method + @klass = klass + + @label = options[:label] || method + @data = options[:data] || {} + end + + def enable + raise "Probe already enabled" if enabled? + + # We're forced to perform this copy because in the class_eval'uated + # block below @method would be evaluated in the class context. It's better + # to close on locally-scoped variables than to resort to complex namespacing + # to get access to the probe instance variables. + method = @method; label = @label; data = @data + klass.class_eval { + alias_method("instrumented_#{method}", method) + define_method(method) do |*args| + id = nil + instrumentation_data = nil + begin + instrumentation_label = label.respond_to?(:call) ? label.call(self, args) : label + instrumentation_data = data.respond_to?(:call) ? data.call(self, args) : data + id = Puppet::Util::Instrumentation.start(instrumentation_label, instrumentation_data) + send("instrumented_#{method}".to_sym, *args) + ensure + Puppet::Util::Instrumentation.stop(instrumentation_label, id, instrumentation_data || {}) + end + end + } + @enabled = true + end + + def disable + raise "Probe is not enabled" unless enabled? + + # For the same reason as in #enable, we're forced to do a local + # copy + method = @method + klass.class_eval do + alias_method(method, "instrumented_#{method}") + remove_method("instrumented_#{method}".to_sym) + end + @enabled = false + end + + def enabled? + !!@enabled + end + end + + # Declares a new probe + # + # It is possible to pass several options that will be later on evaluated + # and sent to the instrumentation layer. + # + # label:: + # this can either be a static symbol/string or a block. If it's a block + # this one will be evaluated on every call of the instrumented method and + # should return a string or a symbol + # + # data:: + # this can be a hash or a block. If it's a block this one will be evaluated + # on every call of the instrumented method and should return a hash. + # + #Example: + # + # class MyClass + # extend Instrumentable + # + # probe :mymethod, :data => Proc.new { |args| { :data => args[1] } }, :label => Proc.new { |args| args[0] } + # + # def mymethod(name, options) + # end + # + # end + # + def probe(method, options = {}) + INSTRUMENTED_CLASSES.synchronize { + (@probes ||= []) << Probe.new(method, self, options) + INSTRUMENTED_CLASSES[self] = @probes + } + end + + def self.probes + @probes + end + + def self.probe_names + probe_names = [] + each_probe { |probe| probe_names << "#{probe.klass}.#{probe.method}" } + probe_names + end + + def self.enable_probes + each_probe { |probe| probe.enable } + end + + def self.disable_probes + each_probe { |probe| probe.disable } + end + + def self.clear_probes + INSTRUMENTED_CLASSES.synchronize { + INSTRUMENTED_CLASSES.clear + } + nil # do not leak our probes to the exterior world + end + + def self.each_probe + INSTRUMENTED_CLASSES.synchronize { + INSTRUMENTED_CLASSES.each_key do |klass| + klass.probes.each { |probe| yield probe } + end + } + nil # do not leak our probes to the exterior world + end +end \ No newline at end of file diff --git a/spec/unit/util/instrumentation/instrumentable_spec.rb b/spec/unit/util/instrumentation/instrumentable_spec.rb new file mode 100755 index 000000000..f1906c49d --- /dev/null +++ b/spec/unit/util/instrumentation/instrumentable_spec.rb @@ -0,0 +1,177 @@ +#!/usr/bin/env rspec + +require 'spec_helper' + +require 'puppet/util/instrumentation' +require 'puppet/util/instrumentation/instrumentable' + +describe Puppet::Util::Instrumentation::Instrumentable::Probe do + + before(:each) do + Puppet::Util::Instrumentation.stubs(:start) + Puppet::Util::Instrumentation.stubs(:stop) + + class ProbeTest + def mymethod(arg1, arg2, arg3) + :it_worked + end + end + end + + after(:each) do + if ProbeTest.method_defined?(:instrumented_mymethod) + ProbeTest.class_eval { + remove_method(:mymethod) + alias_method(:mymethod, :instrumented_mymethod) + } + end + Puppet::Util::Instrumentation::Instrumentable.clear_probes + end + + describe "when enabling a probe" do + it "should raise an error if the probe is already enabled" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + lambda { probe.enable }.should raise_error + end + + it "should rename the original method name" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + ProbeTest.new.should respond_to(:instrumented_mymethod) + end + + it "should create a new method of the original name" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + ProbeTest.new.should respond_to(:mymethod) + end + end + + describe "when disabling a probe" do + it "should raise an error if the probe is already enabled" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + lambda { probe.disable }.should raise_error + end + + it "should rename the original method name" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + probe.disable + + Puppet::Util::Instrumentation.expects(:start).never + Puppet::Util::Instrumentation.expects(:stop).never + ProbeTest.new.mymethod(1,2,3).should == :it_worked + end + + it "should remove the created method" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + probe.disable + ProbeTest.new.should_not respond_to(:instrumented_mymethod) + end + end + + describe "when a probe is called" do + it "should call the original method" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + test = ProbeTest.new + test.expects(:instrumented_mymethod).with(1,2,3) + test.mymethod(1,2,3) + end + + it "should start the instrumentation" do + Puppet::Util::Instrumentation.expects(:start) + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + test = ProbeTest.new + test.mymethod(1,2,3) + end + + it "should stop the instrumentation" do + Puppet::Util::Instrumentation.expects(:stop) + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + test = ProbeTest.new + test.mymethod(1,2,3) + end + + describe "and the original method raises an exception" do + it "should propagate the exception" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest) + probe.enable + test = ProbeTest.new + test.expects(:instrumented_mymethod).with(1,2,3).raises + lambda { test.mymethod(1,2,3) }.should raise_error + end + end + + describe "with a static label" do + it "should send the label to the instrumentation layer" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest, :label => :mylabel) + probe.enable + test = ProbeTest.new + Puppet::Util::Instrumentation.expects(:start).with { |label,data| label == :mylabel }.returns(42) + Puppet::Util::Instrumentation.expects(:stop).with(:mylabel, 42, {}) + test.mymethod(1,2,3) + end + end + + describe "with a dynamic label" do + it "should send the evaluated label to the instrumentation layer" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest, :label => Proc.new { |parent,args| "dynamic#{args[0]}" } ) + probe.enable + test = ProbeTest.new + Puppet::Util::Instrumentation.expects(:start).with { |label,data| label == "dynamic1" }.returns(42) + Puppet::Util::Instrumentation.expects(:stop).with("dynamic1",42,{}) + test.mymethod(1,2,3) + end + end + + describe "with static data" do + it "should send the data to the instrumentation layer" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest, :data => { :static_data => "nothing" }) + probe.enable + test = ProbeTest.new + Puppet::Util::Instrumentation.expects(:start).with { |label,data| data == { :static_data => "nothing" }} + test.mymethod(1,2,3) + end + end + + describe "with dynamic data" do + it "should send the evaluated label to the instrumentation layer" do + probe = Puppet::Util::Instrumentation::Instrumentable::Probe.new(:mymethod, ProbeTest, :data => Proc.new { |parent, args| { :key => args[0] } } ) + probe.enable + test = ProbeTest.new + Puppet::Util::Instrumentation.expects(:start).with { |label,data| data == { :key => 1 } } + Puppet::Util::Instrumentation.expects(:stop) + test.mymethod(1,2,3) + end + end + end +end + +describe Puppet::Util::Instrumentation::Instrumentable do + before(:each) do + class ProbeTest2 + extend Puppet::Util::Instrumentation::Instrumentable + probe :mymethod + def mymethod(arg1,arg2,arg3) + end + end + end + + after do + Puppet::Util::Instrumentation::Instrumentable.clear_probes + end + + it "should allow probe definition" do + Puppet::Util::Instrumentation::Instrumentable.probe_names.should be_include("ProbeTest2.mymethod") + end + + it "should be able to enable all probes" do + Puppet::Util::Instrumentation::Instrumentable.enable_probes + ProbeTest2.new.should respond_to(:instrumented_mymethod) + end +end \ No newline at end of file