From 3b10b72141845f53c52d20c15015516cc6e5c066 Mon Sep 17 00:00:00 2001
From: Burke Libbey <burke@burkelibbey.org>
Date: Mon, 5 Mar 2012 16:11:10 -0500
Subject: [PATCH] Added ability to give an instrumentation callback

---
 README.md                            |  3 ++-
 lib/wolverine/configuration.rb       | 23 ++++++++++++++++++++-
 lib/wolverine/script.rb              | 23 ++++++++++++++++++---
 lib/wolverine/version.rb             |  2 +-
 test/wolverine/configuration_test.rb | 11 ++++++++++
 test/wolverine/script_test.rb        | 30 +++++++++++++++++++++++-----
 6 files changed, 81 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index bb2108d..b5eed88 100644
--- a/README.md
+++ b/README.md
@@ -51,8 +51,9 @@ Available configuration options:
 
 * `Wolverine.config.redis` (default `Redis.new`)
 * `Wolverine.config.script_path` (default `Rails.root + 'app/wolverine'`)
+* `Wolverine.config.instrumentation` (default none)
 
-If you want to override one or both of these, doing so in an initializer is recommended but not required.
+If you want to override one or more of these, doing so in an initializer is recommended but not required. See the [full documentation](http://shopify.github.com/wolverine/Wolverine/Configuration.html) for more details.
 
 ## More information
 
diff --git a/lib/wolverine/configuration.rb b/lib/wolverine/configuration.rb
index 4806f22..e9f81b7 100644
--- a/lib/wolverine/configuration.rb
+++ b/lib/wolverine/configuration.rb
@@ -1,9 +1,30 @@
 module Wolverine
-  class Configuration < Struct.new(:redis, :script_path)
+  class Configuration < Struct.new(:redis, :script_path, :instrumentation)
+
+    # @return [Redis] the redis connection actively in use by Wolverine
     def redis
       super || @redis ||= Redis.new
     end
 
+    # Wolverine.config.instrumentation can be used to specify a callback to
+    # fire with the runtime of each script. This can be useful for analyzing
+    # scripts to make sure they aren't running for an unreasonable amount of
+    # time.
+    # 
+    # The proc will receive three parameters: 
+    # 
+    # * +script_name+: A unique identifier for the script, based on its
+    #   location in the file system
+    # * +runtime+: A float, the total execution time of the script
+    # * +eval_type+: Either +eval+ or +evalsha+, the method used to run
+    #   the script
+    # @return [#call] the proc or other callable to be triggered on completion
+    #   of a script.
+    def instrumentation
+      super || @instrumentation ||= proc { |script_name, runtime, eval_type| nil }
+    end
+
+    # @return [Pathname] the path wolverine will check for scripts
     def script_path
       super || @script_path ||= Rails.root + 'app/wolverine'
     end
diff --git a/lib/wolverine/script.rb b/lib/wolverine/script.rb
index 92b51aa..7721ecc 100644
--- a/lib/wolverine/script.rb
+++ b/lib/wolverine/script.rb
@@ -1,3 +1,5 @@
+require 'pathname'
+require 'benchmark'
 require 'digest/sha1'
 
 module Wolverine
@@ -10,7 +12,7 @@ class Script
     #
     # @param file [Pathname] the full path to the indicated file
     def initialize file
-      @file = file
+      @file = Pathname.new(file)
       @content = load_lua file
       @digest = Digest::SHA1.hexdigest @content
     end
@@ -43,11 +45,26 @@ def call redis, *args
     private
 
     def run_evalsha redis, *args
-      redis.evalsha @digest, args.size, *args
+      instrument :evalsha do
+        redis.evalsha @digest, args.size, *args
+      end
     end
 
     def run_eval redis, *args
-      redis.eval @content, args.size, *args
+      instrument :eval do
+        redis.eval @content, args.size, *args
+      end
+    end
+
+    def instrument eval_type
+      ret = nil
+      runtime = Benchmark.realtime { ret = yield }
+      Wolverine.config.instrumentation.call relative_path.to_s, runtime, eval_type
+      ret
+    end
+
+    def relative_path
+      path = @file.relative_path_from(Wolverine.config.script_path)
     end
 
     def load_lua file
diff --git a/lib/wolverine/version.rb b/lib/wolverine/version.rb
index 0e148d5..3f20f06 100644
--- a/lib/wolverine/version.rb
+++ b/lib/wolverine/version.rb
@@ -1,3 +1,3 @@
 module Wolverine
-  VERSION = "0.2.2"
+  VERSION = "0.2.3"
 end
diff --git a/test/wolverine/configuration_test.rb b/test/wolverine/configuration_test.rb
index ca7e0fd..f43612a 100644
--- a/test/wolverine/configuration_test.rb
+++ b/test/wolverine/configuration_test.rb
@@ -19,6 +19,11 @@ def test_default_script_path
       assert_equal Pathname.new('foo/app/wolverine'), actual
     end
 
+    def test_default_instrumentation
+      config = Wolverine::Configuration.new
+      assert_equal nil, config.instrumentation.call(1, 2, 3)
+    end
+
     def test_setting_redis
       config = Wolverine::Configuration.new
       config.redis = :foo
@@ -31,5 +36,11 @@ def test_setting_script_path
       assert_equal :foo, config.script_path
     end
 
+    def test_setting_instrumentation
+      config = Wolverine::Configuration.new
+      config.instrumentation = proc { |a, b, c| :omg }
+      assert_equal :omg, config.instrumentation.call(1,2,3)
+    end
+
   end
 end
diff --git a/test/wolverine/script_test.rb b/test/wolverine/script_test.rb
index 8ed582f..fe193fb 100644
--- a/test/wolverine/script_test.rb
+++ b/test/wolverine/script_test.rb
@@ -7,21 +7,23 @@ class ScriptTest < MiniTest::Unit::TestCase
     DIGEST = Digest::SHA1.hexdigest(CONTENT)
 
     def setup
+      base = Pathname.new('/a/b/c/d')
+      Wolverine.config.script_path = base
       Wolverine::Script.any_instance.stubs(:load_lua => CONTENT)
     end
 
+    def teardown
+      Wolverine.config.instrumentation = proc{}
+    end
+
     def script
-      @script ||= Wolverine::Script.new('file1')
+      @script ||= Wolverine::Script.new('/a/b/c/d/e/file1.lua')
     end
 
     def test_error
-      base = Pathname.new('/a/b/c/d')
-      file = Pathname.new('/a/b/c/d/e/file1.lua')
-      Wolverine.config.script_path = base
       redis = stub
       redis.expects(:evalsha).raises(%q{ERR Error running script (call to f_178d75adaa46af3d8237cfd067c9fdff7b9d504f): [string "func definition"]:1: attempt to compare nil with number})
       begin
-        script = Wolverine::Script.new(file)
         script.call(redis)
       rescue Wolverine::LuaError => e
         assert_equal "attempt to compare nil with number", e.message
@@ -30,6 +32,24 @@ def test_error
       end
     end
 
+    def test_instrumentation
+      callback = Object.new
+      tc = self
+      meta = class << callback ; self ; end
+      meta.send(:define_method, :call) { |a, b, c|
+        tc.assert_equal "e/file1.lua", a
+        tc.assert_operator b, :<, 1
+        tc.assert_equal :evalsha, c
+      }
+      Wolverine.config.instrumentation = callback
+      redis = Class.new do
+        define_method(:evalsha) do |digest, size, *args|
+          nil
+        end
+      end
+      script.call(redis.new, :a, :b)
+    end
+
     def test_call_with_cache_hit
       tc = self
       redis = Class.new do