memcache.rb in oboe-1.4.0.2

- old
+ new
@@ -1,40 +1,109 @@
 # Copyright (c) 2012 by Tracelytics, Inc.
 # All rights reserved.
 
-if defined?(::MemCache)
-  class ::MemCache
-    include Oboe::API::Memcache
+module Oboe
+  module Inst
+    module MemCache
+      include Oboe::API::Memcache
+      
+      def self.included(cls)
+        puts "[oboe/loading] Instrumenting memcache" if Oboe::Config[:verbose]
 
-    MEMCACHE_OPS.reject { |m| not method_defined?(m) }.each do |m|
-      opts = { :KVOp => m }
-      define_method("#{m}_with_oboe") do |*args|
-        Oboe::API.trace('memcache', opts) do
-          send("#{m}_without_oboe", *args) 
+        cls.class_eval do
+          MEMCACHE_OPS.reject { |m| not method_defined?(m) }.each do |m|
+
+            define_method("#{m}_with_oboe") do |*args|
+              report_kvs = { :KVOp => m }
+
+              if Oboe.tracing?
+                Oboe::API.trace('memcache', report_kvs) do
+                  result = send("#{m}_without_oboe", *args) 
+                end
+              else
+                result = send("#{m}_without_oboe", *args) 
+              end
+              result
+            end
+
+            class_eval "alias #{m}_without_oboe #{m}"
+            class_eval "alias #{m} #{m}_with_oboe"
+          end
+
+          if ::MemCache.method_defined? :request_setup
+            alias request_setup_without_oboe request_setup
+            alias request_setup request_setup_with_oboe
+          elsif Oboe::Config[:verbose]
+            puts "[oboe/loading] Couldn't properly instrument Memcache.  Partial traces may occur."
+          end
+
+          if ::MemCache.method_defined? :cache_get
+            alias cache_get_without_oboe cache_get
+            alias cache_get cache_get_with_oboe
+          elsif Oboe::Config[:verbose]
+            puts "[oboe/loading] Couldn't properly instrument Memcache.  Partial traces may occur." 
+          end
+          
+          if ::MemCache.method_defined? :get_multi
+            alias get_multi_without_oboe get_multi
+            alias get_multi get_multi_with_oboe
+          elsif Oboe::Config[:verbose]
+            puts "[oboe/loading] Couldn't properly instrument Memcache.  Partial traces may occur." 
+          end
         end
       end
 
-      class_eval "alias #{m}_without_oboe #{m}"
-      class_eval "alias #{m} #{m}_with_oboe"
-    end
+      def get_multi_with_oboe(*args)
+        if Oboe.tracing?
+          layer_kvs = {}
+          layer_kvs[:KVOp] = :get_multi
 
-    define_method(:request_setup_with_oboe) do |*args|
-      server, cache_key = request_setup_without_oboe(*args)
-      Oboe::API.log('memcache', 'info', { :KVKey => cache_key, :RemoteHost => server.host })
-      return [server, cache_key]
-    end
+          Oboe::API.trace('memcache', layer_kvs || {}, :get_multi) do
+            begin
+              info_kvs = {}
+               
+              if args.last.is_a?(Hash) || args.last.nil?
+                info_kvs[:KVKeyCount] = args.flatten.length - 1
+              else
+                info_kvs[:KVKeyCount] = args.flatten.length 
+              end
 
-    alias request_setup_without_oboe request_setup
-    alias request_setup request_setup_with_oboe
+              values = get_multi_without_oboe(args)
+              
+              info_kvs[:KVHitCount] = values.length
+              Oboe::API.log('memcache', 'info', info_kvs)
+            rescue
+              values = get_multi_without_oboe(args)
+            end
+            values 
+          end
+        else
+          get_multi_without_oboe(args)
+        end
+      end
+      
+      def request_setup_with_oboe(*args)
+        if Oboe.tracing? and not Oboe::Context.tracing_layer_op?(:get_multi)
+          server, cache_key = request_setup_without_oboe(*args)
+          Oboe::API.log('memcache', 'info', { :KVKey => cache_key, :RemoteHost => server.host })
+        else
+          server, cache_key = request_setup_without_oboe(*args)
+        end
+        return [server, cache_key]
+      end
 
+      def cache_get_with_oboe(server, cache_key)
+        result = cache_get_without_oboe(server, cache_key)
+        Oboe::API.log('memcache', 'info', { :KVHit => memcache_hit?(result) })
+        result
+      end
 
-    define_method(:cache_get_with_oboe) do |server, cache_key|
-      result = cache_get_without_oboe(server, cache_key)
-      Oboe::API.log('memcache', 'info', { :KVHit => memcache_hit?(result) })
-      result
-    end
+    end # module MemCache
+  end # module Inst
+end # module Oboe
 
-    alias cache_get_without_oboe cache_get
-    alias cache_get cache_get_with_oboe
+if defined?(::MemCache) and Oboe::Config[:memcache][:enabled]
+  ::MemCache.class_eval do
+    include Oboe::Inst::MemCache
   end
-  puts "[oboe/loading] Instrumenting memcache" if Oboe::Config[:verbose]
 end
+