Process name instrumentation infrastructure

This is special feature that changes the process name of the running puppet
entity to display its current activity.

It is disabled by default, and can be enabled by sending the QUIT signal
to the process in question (or calling enable through the code).

This system can work only if some "probes" are integrated in the core puppet
codebase. Since tools to visualize process names have a large refresh time
(ie more than 1s) it only makes sense to track long activities (like compilation,
transaction or file serving).
Those probes are the subject of a subsequent patch.

This system tracks every thread activity and form a strings which will
be used as the process name. Due to the way it is implemented it is
possible that it doesn't work on all platforms (I tested successfully
on osx and linux). On some systems the space available is dependent on
the original size of the full command. That's why if this string is longer
than a 50 characters, the string is scrolled (like stock market tickers).

Note: This is not intended to be a generic instrumentation system. Also, being
block based means that it can reduce performance if the instrumentation
probes are used in tight inner loops.

Signed-off-by: Brice Figureau <brice-puppet@daysofwonder.com>

3 files changed, 348 insertions, 0 deletions

diff --git a/lib/puppet/util/instrumentation.rb b/lib/puppet/util/instrumentation.rb
new file mode 100644
index 000000000..5981bea59
--- /dev/null
+++ b/lib/puppet/util/instrumentation.rb
@@ -0,0 +1,12 @@
+require 'puppet/util/instrumentation/process_name'
+
+module Puppet::Util::Instrumentation
+
+  def instrument(title)
+    Puppet::Util::Instrumentation::ProcessName.instrument(title) do
+      yield
+    end
+  end
+  module_function :instrument
+
+end
\ No newline at end of file
diff --git a/lib/puppet/util/instrumentation/process_name.rb b/lib/puppet/util/instrumentation/process_name.rb
new file mode 100644
index 000000000..370d29e2e
--- /dev/null
+++ b/lib/puppet/util/instrumentation/process_name.rb
@@ -0,0 +1,129 @@
+require 'puppet'
+require 'puppet/util/instrumentation'
+
+module Puppet::Util::Instrumentation
+  class ProcessName
+
+    # start scrolling when process name is longer than
+    SCROLL_LENGTH = 50
+
+    @active = false
+    class << self
+      attr_accessor :active, :reason
+    end
+
+    trap(:QUIT) do
+      active? ? disable : enable
+    end
+
+    def self.active?
+      !! @active
+    end
+
+    def self.enable
+      mutex.synchronize do
+        Puppet.info("Process Name instrumentation is enabled")
+        @active = true
+        @x = 0
+        setproctitle
+      end
+    end
+
+    def self.disable
+      mutex.synchronize do
+        Puppet.info("Process Name instrumentation is disabled")
+        @active = false
+        $0 = @oldname
+      end
+    end
+
+    def self.instrument(activity)
+      # inconditionnally start the scroller thread here
+      # because it doesn't seem possible to start a new thrad
+      # from the USR2 signal handler
+      @scroller ||= Thread.new do
+        loop do
+          scroll if active?
+          sleep 1
+        end
+      end
+
+      push_activity(Thread.current, activity)
+      yield
+    ensure
+      pop_activity(Thread.current)
+    end
+
+    def self.setproctitle
+      @oldname ||= $0
+      $0 = "#{base}: " + rotate(process_name,@x) if active?
+    end
+
+    def self.push_activity(thread, activity)
+      mutex.synchronize do
+        @reason ||= {}
+        @reason[thread] ||= []
+        @reason[thread].push(activity)
+        setproctitle
+      end
+    end
+
+    def self.pop_activity(thread)
+      mutex.synchronize do
+        @reason[thread].pop
+        if @reason[thread].empty?
+          @reason.delete(thread)
+        end
+        setproctitle
+      end
+    end
+
+    def self.process_name
+      out = (@reason || {}).inject([]) do |out, reason|
+        out << "#{thread_id(reason[0])} #{reason[1].join(',')}"
+      end
+      out.join(' | ')
+    end
+
+    # certainly non-portable
+    def self.thread_id(thread)
+      thread.inspect.gsub(/^#<.*:0x([a-f0-9]+) .*>$/, '\1')
+    end
+
+    def self.rotate(string, steps)
+      steps ||= 0
+      if string.length > 0 && steps > 0
+        steps = steps % string.length
+        return string[steps..string.length].concat " -- #{string[0..(steps-1)]}"
+      end
+      string
+    end
+
+    def self.base
+      basename = case Puppet.run_mode.name
+      when :master
+        "master"
+      when :agent
+        "agent"
+      else
+        "puppet"
+      end
+    end
+
+    def self.mutex
+      #Thread.exclusive {
+        @mutex ||= Sync.new
+      #}
+      @mutex
+    end
+
+    def self.scroll
+      return if process_name.length < SCROLL_LENGTH
+      mutex.synchronize do
+        setproctitle
+        @x += 1
+      end
+    end
+
+  end
+end
\ No newline at end of file
diff --git a/spec/unit/util/instrumentation/process_name_spec.rb b/spec/unit/util/instrumentation/process_name_spec.rb
new file mode 100644
index 000000000..9cbedf2d2
--- /dev/null
+++ b/spec/unit/util/instrumentation/process_name_spec.rb
@@ -0,0 +1,207 @@
+#!/usr/bin/env ruby
+
+require File.dirname(__FILE__) + '/../../../spec_helper'
+
+describe Puppet::Util::Instrumentation::ProcessName do
+
+  ProcessName = Puppet::Util::Instrumentation::ProcessName
+
+  after(:each) do
+    ProcessName.reason = {}
+  end
+
+  it "should be disabled by default" do
+    ProcessName.should_not be_active
+  end
+
+  describe "when managing thread activity" do
+    before(:each) do
+      ProcessName.stubs(:setproctitle)
+      ProcessName.stubs(:base).returns("base")
+    end
+
+    it "should be able to append activity" do
+      thread1 = stub 'thread1'
+      ProcessName.push_activity(:thread1,"activity1")
+      ProcessName.push_activity(:thread1,"activity2")
+
+      ProcessName.reason[:thread1].should == ["activity1", "activity2"]
+    end
+
+    it "should be able to remove activity" do
+      ProcessName.push_activity(:thread1,"activity1")
+      ProcessName.push_activity(:thread1,"activity1")
+      ProcessName.pop_activity(:thread1)
+
+      ProcessName.reason[:thread1].should == ["activity1"]
+    end
+
+    it "should maintain activity thread by thread" do
+      ProcessName.push_activity(:thread1,"activity1")
+      ProcessName.push_activity(:thread2,"activity2")
+
+      ProcessName.reason[:thread1].should == ["activity1"]
+      ProcessName.reason[:thread2].should == ["activity2"]
+    end
+
+    it "should set process title" do
+      ProcessName.expects(:setproctitle)
+
+      ProcessName.push_activity("thread1","activity1")
+    end
+  end
+
+  describe "when computing the current process name" do
+      before(:each) do
+        ProcessName.stubs(:setproctitle)
+        ProcessName.stubs(:base).returns("base")
+      end
+
+      it "should include every running thread activity" do
+        thread1 = stub 'thread1', :inspect => "\#<Thread:0xdeadbeef run>", :hash => 1
+        thread2 = stub 'thread2', :inspect => "\#<Thread:0x12344321 run>", :hash => 0
+
+        ProcessName.push_activity(thread1,"Compiling node1.domain.com")
+        ProcessName.push_activity(thread2,"Compiling node4.domain.com")
+        ProcessName.push_activity(thread1,"Parsing file site.pp")
+        ProcessName.push_activity(thread2,"Parsing file node.pp")
+
+        ProcessName.process_name.should == "12344321 Compiling node4.domain.com,Parsing file node.pp | deadbeef Compiling node1.domain.com,Parsing file site.pp"
+      end
+  end
+
+  describe "when finding base process name" do
+      {:master => "master", :agent => "agent", :user => "puppet"}.each do |program,base|
+        it "should return #{base} for #{program}" do
+          Puppet.run_mode.stubs(:name).returns(program)
+          ProcessName.base.should == base
+        end
+      end
+  end
+
+  describe "when finding a thread id" do
+      it "should return the id from the thread inspect string" do
+        thread = stub 'thread', :inspect => "\#<Thread:0x1234abdc run>"
+        ProcessName.thread_id(thread).should == "1234abdc"
+      end
+  end
+
+  describe "when scrolling the instrumentation string" do
+      it "should rotate the string of various step" do
+        ProcessName.rotate("this is a rotation", 10).should == "rotation -- this is a "
+      end
+
+      it "should not rotate the string for the 0 offset" do
+        ProcessName.rotate("this is a rotation", 0).should == "this is a rotation"
+      end
+  end
+
+  describe "when setting process name" do
+    before(:each) do
+      ProcessName.stubs(:process_name).returns("12345 activity")
+      ProcessName.stubs(:base).returns("base")
+      @oldname = $0
+    end
+
+    after(:each) do
+      $0 = @oldname
+    end
+
+    it "should not do it if the feature is disabled" do
+      ProcessName.setproctitle
+
+      $0.should_not == "base: 12345 activity"
+    end
+
+    it "should do it if the feature is enabled" do
+      ProcessName.active = true
+      ProcessName.setproctitle
+
+      $0.should == "base: 12345 activity"
+    end
+  end
+
+  describe "when setting a probe" do
+    before(:each) do
+      thread = stub 'thread', :inspect => "\#<Thread:0x1234abdc run>"
+      Thread.stubs(:current).returns(thread)
+      Thread.stubs(:new)
+      ProcessName.active = true
+    end
+
+    it "should start the scroller thread" do
+      Thread.expects(:new)
+      ProcessName.instrument("doing something") do
+      end
+    end
+
+    it "should push current thread activity and execute the block" do
+      ProcessName.instrument("doing something") do
+        $0.should == "puppet: 1234abdc doing something"
+      end
+    end
+
+    it "should finally pop the activity" do
+      ProcessName.instrument("doing something") do
+      end
+      $0.should == "puppet: "
+    end
+  end
+
+  describe "when enabling" do
+    before do
+      Thread.stubs(:new)
+      ProcessName.stubs(:setproctitle)
+    end
+
+    it "should be active" do
+      ProcessName.enable
+      ProcessName.should be_active
+    end
+
+    it "should set the new process name" do
+      ProcessName.expects(:setproctitle)
+      ProcessName.enable
+    end
+  end
+
+  describe "when disabling" do
+    it "should set active to false" do
+      ProcessName.active = true
+      ProcessName.disable
+      ProcessName.should_not be_active
+    end
+
+    it "should restore the old process name" do
+      oldname = $0
+      ProcessName.active = true
+      ProcessName.setproctitle
+      ProcessName.disable
+      $0.should == oldname
+    end
+  end
+
+  describe "when scrolling" do
+    it "should do nothing for shorter process names" do
+      ProcessName.expects(:setproctitle).never
+      ProcessName.scroll
+    end
+
+    it "should call setproctitle" do
+      ProcessName.stubs(:process_name).returns("x" * 60)
+      ProcessName.expects(:setproctitle)
+      ProcessName.scroll
+    end
+
+    it "should increment rotation offset" do
+      name = "x" * 60
+      ProcessName.active = true
+      ProcessName.stubs(:process_name).returns(name)
+      ProcessName.expects(:rotate).once.with(name,1).returns("")
+      ProcessName.expects(:rotate).once.with(name,2).returns("")
+      ProcessName.scroll
+      ProcessName.scroll
+    end
+  end
+
+end
\ No newline at end of file