memcached.rb in oboe-1.4.0.2

- old
+ new
@@ -1,32 +1,91 @@
 # Copyright (c) 2012 by Tracelytics, Inc.
 # All rights reserved.
 #
 
-if defined?(Memcached)
-  class Memcached
-    include Oboe::API::Memcache
+module Oboe
+  module Inst
+    module Memcached
+      include Oboe::API::Memcache
+      
+      def self.included(cls)
+        puts "[oboe/loading] Instrumenting memcached" if Oboe::Config[:verbose]
 
-    MEMCACHE_OPS.reject { |m| not method_defined?(m) }.each do |m|
-      define_method("#{m}_with_oboe") do |*args|
-        opts = { :KVOp => m }
-        if args.length and args[0].class != Array
-            opts[:KVKey] = args[0].to_s
-            rhost = remote_host(args[0].to_s)
-            opts[:RemoteHost] = rhost if rhost
+        cls.class_eval do
+          MEMCACHE_OPS.reject { |m| not method_defined?(m) }.each do |m|
+            define_method("#{m}_with_oboe") do |*args|
+              opts = { :KVOp => m }
+              if args.length and args[0].class != Array
+                opts[:KVKey] = args[0].to_s
+                rhost = remote_host(args[0].to_s)
+                opts[:RemoteHost] = rhost if rhost
+              end
+
+              Oboe::API.trace('memcache', opts) do
+                result = send("#{m}_without_oboe", *args)
+                if m == :get and args.length and args[0].class == String
+                  Oboe::API.log('memcache', 'info', { :KVHit => memcache_hit?(result) })
+                end
+                result
+              end
+            end
+
+            class_eval "alias #{m}_without_oboe #{m}"
+            class_eval "alias #{m} #{m}_with_oboe"
+          end
         end
+      end
 
-        Oboe::API.trace('memcache', opts) do
-          result = send("#{m}_without_oboe", *args)
-          if m == :get and args.length and args[0].class == String
-              Oboe::API.log('memcache', 'info', { :KVHit => memcache_hit?(result) })
+    end # module Memcached
+
+    module MemcachedRails
+      def self.included(cls)
+        cls.class_eval do
+          if ::Memcached::Rails.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 Memcached.  Partial traces may occur." 
           end
-          result
         end
       end
 
-      class_eval "alias #{m}_without_oboe #{m}"
-      class_eval "alias #{m} #{m}_with_oboe"
+      def get_multi_with_oboe(keys, raw=false)
+        if Oboe.tracing?
+          layer_kvs = {}
+          layer_kvs[:KVOp] = :get_multi
+
+          Oboe::API.trace('memcache', layer_kvs || {}, :get_multi) do
+            begin
+              info_kvs = {}
+              info_kvs[:KVKeyCount] = keys.flatten.length 
+
+              values = get_multi_without_oboe(keys, raw)
+              
+              info_kvs[:KVHitCount] = values.length
+              Oboe::API.log('memcache', 'info', info_kvs)
+            rescue
+              values = get_multi_without_oboe(keys, raw)
+            end
+            values 
+          end
+        else
+          get_multi_without_oboe(keys, raw)
+        end
+      end
+
+    end # module MemcachedRails
+  end # module Inst
+end # module Oboe
+
+if defined?(::Memcached) and Oboe::Config[:memcached][:enabled]
+  ::Memcached.class_eval do
+    include Oboe::Inst::Memcached
+  end
+  
+  if defined?(::Memcached::Rails)
+    ::Memcached::Rails.class_eval do
+      include Oboe::Inst::MemcachedRails
     end
   end
-  puts "[oboe/loading] Instrumenting memcached" if Oboe::Config[:verbose]
 end
+