Add "multi" methods instrumentation for Rails cache integration

lib/ddtrace/contrib/active_support/cache/instrumentation.rb

@@ -18,9 +18,13 @@ def start_trace_cache(payload)
             # NOTE: the ``finish_trace_cache()`` is fired but it already has a safe-guard
             # to avoid any kind of issue.
             current_span = tracer.active_span
-            return if payload[:action] == Ext::RESOURCE_CACHE_GET &&
-                      current_span.try(:name) == Ext::SPAN_CACHE &&
-                      current_span.try(:resource) == Ext::RESOURCE_CACHE_GET
+            return if current_span.try(:name) == Ext::SPAN_CACHE &&
+                      (
+                        payload[:action] == Ext::RESOURCE_CACHE_GET &&
+                        current_span.try(:resource) == Ext::RESOURCE_CACHE_GET ||
+                        payload[:action] == Ext::RESOURCE_CACHE_MGET &&
+                        current_span.try(:resource) == Ext::RESOURCE_CACHE_MGET
+                      )
 
             tracing_context = payload.fetch(:tracing_context)
 
@@ -59,6 +63,32 @@ def finish_trace_cache(payload)
             Datadog.logger.debug(e.message)
           end
 
+          def finish_trace_cache_multi(payload)
+            # retrieve the tracing context and continue the trace
+            tracing_context = payload.fetch(:tracing_context)
+            span = tracing_context[:dd_cache_span]
+            return unless span && !span.finished?
+
+            begin
+              # discard parameters from the cache_store configuration
+              if defined?(::Rails)
+                store, = *Array.wrap(::Rails.configuration.cache_store).flatten
+                span.set_tag(Ext::TAG_CACHE_BACKEND, store)
+              end
+              normalized_keys = payload.fetch(:keys, []).map do |key|
+                ::ActiveSupport::Cache.expand_cache_key(key)
+              end
+              cache_keys = Datadog::Utils.truncate(normalized_keys, Ext::QUANTIZE_CACHE_MAX_KEY_SIZE)
+              span.set_tag(Ext::TAG_CACHE_KEY_MULTI, cache_keys)
+
+              span.set_error(payload[:exception]) if payload[:exception]
+            ensure
+              span.finish
+            end
+          rescue StandardError => e
+            Datadog.logger.debug(e.message)
+          end
+
           # Defines instrumentation for ActiveSupport cache reading
           module Read
             def read(*args, &block)
@@ -82,6 +112,29 @@ def read(*args, &block)
             end
           end
 
+          # Defines instrumentation for ActiveSupport cache reading of multiple keys
+          module ReadMulti
+            def read_multi(*keys, &block)
+              payload = {
+                action: Ext::RESOURCE_CACHE_MGET,
+                keys: keys,
+                tracing_context: {}
+              }
+
+              begin
+                # process and catch cache exceptions
+                Instrumentation.start_trace_cache(payload)
+                super
+              rescue Exception => e
+                payload[:exception] = [e.class.name, e.message]
+                payload[:exception_object] = e
+                raise e
+              end
+            ensure
+              Instrumentation.finish_trace_cache_multi(payload)
+            end
+          end
+
           # Defines instrumentation for ActiveSupport cache fetching
           module Fetch
             def fetch(*args, &block)
@@ -105,6 +158,31 @@ def fetch(*args, &block)
             end
           end
 
+          # Defines instrumentation for ActiveSupport cache fetching of multiple keys
+          module FetchMulti
+            def fetch_multi(*args, &block)
+              # extract options hash
+              keys = args[-1].instance_of?(Hash) ? args[0..-2] : args
+              payload = {
+                action: Ext::RESOURCE_CACHE_MGET,
+                keys: keys,
+                tracing_context: {}
+              }
+
+              begin
+                # process and catch cache exceptions
+                Instrumentation.start_trace_cache(payload)
+                super
+              rescue Exception => e
+                payload[:exception] = [e.class.name, e.message]
+                payload[:exception_object] = e
+                raise e
+              end
+            ensure
+              Instrumentation.finish_trace_cache_multi(payload)
+            end
+          end
+
           # Defines instrumentation for ActiveSupport cache writing
           module Write
             def write(*args, &block)
@@ -128,6 +206,29 @@ def write(*args, &block)
             end
           end
 
+          # Defines instrumentation for ActiveSupport cache writing of multiple keys
+          module WriteMulti
+            def write_multi(hash, options = nil)
+              payload = {
+                action: Ext::RESOURCE_CACHE_MSET,
+                keys: hash.keys,
+                tracing_context: {}
+              }
+
+              begin
+                # process and catch cache exceptions
+                Instrumentation.start_trace_cache(payload)
+                super
+              rescue Exception => e
+                payload[:exception] = [e.class.name, e.message]
+                payload[:exception_object] = e
+                raise e
+              end
+            ensure
+              Instrumentation.finish_trace_cache_multi(payload)
+            end
+          end
+
           # Defines instrumentation for ActiveSupport cache deleting
           module Delete
             def delete(*args, &block)

lib/ddtrace/contrib/active_support/cache/patcher.rb

@@ -17,8 +17,11 @@ def target_version
 
           def patch
             patch_cache_store_read
+            patch_cache_store_read_multi
             patch_cache_store_fetch
+            patch_cache_store_fetch_multi
             patch_cache_store_write
+            patch_cache_store_write_multi
             patch_cache_store_delete
           end
 
@@ -30,14 +33,32 @@ def patch_cache_store_read
             cache_store_class(:read).send(:prepend, Cache::Instrumentation::Read)
           end
 
+          def patch_cache_store_read_multi
+            cache_store_class(:read_multi).send(:prepend, Cache::Instrumentation::ReadMulti)
+          end
+
           def patch_cache_store_fetch
             cache_store_class(:fetch).send(:prepend, Cache::Instrumentation::Fetch)
           end
 
+          def patch_cache_store_fetch_multi
+            klass = cache_store_class(:fetch_multi)
+            return unless klass.public_method_defined?(:fetch_multi)
+
+            klass.send(:prepend, Cache::Instrumentation::FetchMulti)
+          end
+
           def patch_cache_store_write
             cache_store_class(:write).send(:prepend, Cache::Instrumentation::Write)
           end
 
+          def patch_cache_store_write_multi
+            klass = cache_store_class(:write_multi)
+            return unless klass.public_method_defined?(:write_multi)
+
+            klass.send(:prepend, Cache::Instrumentation::WriteMulti)
+          end
+
           def patch_cache_store_delete
             cache_store_class(:delete).send(:prepend, Cache::Instrumentation::Delete)
           end

lib/ddtrace/contrib/active_support/ext.rb

@@ -12,12 +12,15 @@ module Ext
         QUANTIZE_CACHE_MAX_KEY_SIZE = 300
         RESOURCE_CACHE_DELETE = 'DELETE'.freeze
         RESOURCE_CACHE_GET = 'GET'.freeze
+        RESOURCE_CACHE_MGET = 'MGET'.freeze
         RESOURCE_CACHE_SET = 'SET'.freeze
+        RESOURCE_CACHE_MSET = 'MSET'.freeze
         SERVICE_CACHE = 'active_support-cache'.freeze
         SPAN_CACHE = 'rails.cache'.freeze
         SPAN_TYPE_CACHE = 'cache'.freeze
         TAG_CACHE_BACKEND = 'rails.cache.backend'.freeze
         TAG_CACHE_KEY = 'rails.cache.key'.freeze
+        TAG_CACHE_KEY_MULTI = 'rails.cache.keys'.freeze
       end
     end
   end

spec/ddtrace/contrib/rails/cache_spec.rb

@@ -18,6 +18,7 @@
   let(:cache) { Rails.cache }
 
   let(:key) { 'custom-key' }
+  let(:multi_keys) { %w[custom-key-1 custom-key-2 custom-key-3] }
 
   context '#read' do
     subject(:read) { cache.read(key) }
@@ -44,6 +45,32 @@
     end
   end
 
+  context '#read_multi' do
+    subject(:read_multi) { cache.read_multi(*multi_keys) }
+
+    before { multi_keys.each { |key| cache.write(key, 50 + key[-1].to_i) } }
+
+    it_behaves_like 'measured span for integration', false do
+      before { read_multi }
+      let(:span) { spans[0] }
+    end
+
+    it do
+      expect(read_multi).to eq(Hash[multi_keys.zip([51, 52, 53])])
+      expect(spans).to have(1 + multi_keys.size).items
+      get = spans[0]
+      expect(get.name).to eq('rails.cache')
+      expect(get.span_type).to eq('cache')
+      expect(get.resource).to eq('MGET')
+      expect(get.service).to eq('rails-cache')
+      expect(get.get_tag('rails.cache.backend').to_s).to eq('file_store')
+      expect(JSON.parse(get.get_tag('rails.cache.keys'))).to eq(multi_keys)
+      spans[1..-1].each do |set|
+        expect(set.name).to eq('rails.cache')
+      end
+    end
+  end
+
   context '#write' do
     subject(:write) { cache.write(key, 50) }
 
@@ -81,6 +108,70 @@
     end
   end
 
+  context '#write_multi' do
+    let(:values) { multi_keys.map { |k| 50 + k[-1].to_i } }
+
+    subject(:write_multi) { cache.write_multi(Hash[multi_keys.zip(values)], opt_name: :opt_value) }
+
+    context 'when the method is defined' do
+      before do
+        unless ::ActiveSupport::Cache::Store.public_method_defined?(:write_multi)
+          skip 'Test is not applicable to this Rails version'
+        end
+      end
+
+      it_behaves_like 'measured span for integration', false do
+        before { write_multi }
+      end
+
+      it do
+        write_multi
+        expect(span.name).to eq('rails.cache')
+        expect(span.span_type).to eq('cache')
+        expect(span.resource).to eq('MSET')
+        expect(span.service).to eq('rails-cache')
+        expect(span.get_tag('rails.cache.backend').to_s).to eq('file_store')
+        expect(JSON.parse(span.get_tag('rails.cache.keys'))).to eq(multi_keys)
+      end
+
+      context 'with custom cache_service' do
+        before { Datadog.configuration[:rails][:cache_service] = 'service-cache' }
+
+        it 'uses the proper service name' do
+          write_multi
+          expect(span.service).to eq('service-cache')
+        end
+      end
+
+      context 'with complex cache key' do
+        let(:key) { ['custom-key', %w[x y], user] }
+        let(:user) { double('User', cache_key: 'User:3') }
+
+        it 'expands key using ActiveSupport' do
+          cache.write_multi(key => 0)
+          expect(span.get_tag('rails.cache.keys')).to eq('["custom-key/x/y/User:3"]')
+        end
+      end
+    end
+
+    context 'when the method is not defined' do
+      before do
+        if ::ActiveSupport::Cache::Store.public_method_defined?(:write_multi)
+          skip 'Test is not applicable to this Rails version'
+        end
+      end
+      it do
+        expect(::ActiveSupport::Cache::Store.ancestors).not_to(
+          include(::Datadog::Contrib::ActiveSupport::Cache::Instrumentation::WriteMulti)
+        )
+      end
+
+      it do
+        expect { subject }.to raise_error NoMethodError
+      end
+    end
+  end
+
   context '#delete' do
     subject!(:delete) { cache.delete(key) }
 
@@ -123,6 +214,58 @@
     end
   end
 
+  context '#fetch_multi' do
+    subject(:fetch_multi) { cache.fetch_multi(*multi_keys, expires_in: 42) { |key| 50 + key[-1].to_i } }
+
+    context 'when the method is defined' do
+      before do
+        unless ::ActiveSupport::Cache::Store.public_method_defined?(:fetch_multi)
+          skip 'Test is not applicable to this Rails version'
+        end
+      end
+
+      it_behaves_like 'measured span for integration', false do
+        before { fetch_multi }
+        # Choose either GET or SET span
+        let(:span) { spans.sample }
+      end
+
+      context 'with exception' do
+        subject(:fetch_multi) { cache.fetch_multi('exception', 'another', 'one') { raise 'oops' } }
+
+        it do
+          expect { fetch_multi }.to raise_error(StandardError)
+          expect(span.name).to eq('rails.cache')
+          expect(span.span_type).to eq('cache')
+          expect(span.resource).to eq('MGET')
+          expect(span.service).to eq('rails-cache')
+          expect(span.get_tag('rails.cache.backend').to_s).to eq('file_store')
+          expect(span.get_tag('rails.cache.keys')).to eq('["exception", "another", "one"]')
+          expect(span.get_tag('error.type')).to eq('RuntimeError')
+          expect(span.get_tag('error.msg')).to eq('oops')
+        end
+      end
+    end
+
+    context 'when the method is not defined' do
+      before do
+        if ::ActiveSupport::Cache::Store.public_method_defined?(:fetch_multi)
+          skip 'Test is not applicable to this Rails version'
+        end
+      end
+
+      it do
+        expect(::ActiveSupport::Cache::Store.ancestors).not_to(
+          include(::Datadog::Contrib::ActiveSupport::Cache::Instrumentation::FetchMulti)
+        )
+      end
+
+      it do
+        expect { subject }.to raise_error NoMethodError
+      end
+    end
+  end
+
   context 'with very large cache key' do
     it 'truncates key too large' do
       max_key_size = Datadog::Contrib::ActiveSupport::Ext::QUANTIZE_CACHE_MAX_KEY_SIZE