Use ActiveSupport events to instrument ActiveSupport::Cache

Signed-off-by: Marco Costa <marco.costa@datadoghq.com>

lib/datadog/tracing/contrib/active_support/cache/event.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+require_relative '../notifications/event'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module ActiveSupport
+        module Cache
+          # Defines basic behaviors for an ActiveSupport event.
+          module Event
+            def self.included(base)
+              base.include(ActiveSupport::Notifications::Event)
+              base.extend(ClassMethods)
+            end
+
+            # Class methods for ActiveRecord events.
+            module ClassMethods
+              def span_options
+                {}
+              end
+
+              def configuration
+                Datadog.configuration.tracing[:active_support]
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/active_support/cache/events.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require_relative 'events/cache'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module ActiveSupport
+        module Cache
+          # Defines collection of instrumented ActiveSupport events
+          module Events
+            ALL = [
+              Events::Cache,
+            ].freeze
+
+            module_function
+
+            def all
+              self::ALL
+            end
+
+            def subscriptions
+              all.collect(&:subscriptions).collect(&:to_a).flatten
+            end
+
+            def subscribe!
+              all.each(&:subscribe!)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/active_support/cache/events/cache.rb

@@ -0,0 +1,156 @@
+# frozen_string_literal: true
+
+require_relative '../../ext'
+require_relative '../event'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module ActiveSupport
+        module Cache
+          module Events
+            # Defines instrumentation for instantiation.active_record event
+            module Cache
+              include ActiveSupport::Cache::Event
+
+              module_function
+
+              # Acts as this module's initializer.
+              def subscribe!
+                @cache_backend = {}
+                super
+              end
+
+              def event_name
+                /\Acache_(?:delete|read|read_multi|write|write_multi)\.active_support\z/
+              end
+
+              def span_name
+                Ext::SPAN_CACHE
+              end
+
+              def span_options
+                {
+                  type: Ext::SPAN_TYPE_CACHE
+                }
+              end
+
+              # DEV: Look for other uses of `ActiveSupport::Cache::Store#instrument`, to find other useful event keys.
+              MAPPING = {
+                'cache_delete.active_support' => { resource: Ext::RESOURCE_CACHE_DELETE },
+                'cache_read.active_support' => { resource: Ext::RESOURCE_CACHE_GET },
+                'cache_read_multi.active_support' => { resource: Ext::RESOURCE_CACHE_MGET, multi_key: true },
+                'cache_write.active_support' => { resource: Ext::RESOURCE_CACHE_SET },
+                'cache_write_multi.active_support' => { resource: Ext::RESOURCE_CACHE_MSET, multi_key: true }
+              }.freeze
+
+              def trace?(event, _payload)
+                return false if !Tracing.enabled? || !configuration.enabled
+
+                # DEV-3.0: Backwards compatibility code for the 2.x gem series.
+                # DEV-3.0: See documentation at {Datadog::Tracing::Contrib::ActiveSupport::Cache::Instrumentation}
+                # DEV-3.0: for the complete information about this backwards compatibility code.
+                case event
+                when 'cache_read.active_support'
+                  !ActiveSupport::Cache::Instrumentation.nested_read?
+                when 'cache_read_multi.active_support'
+                  !ActiveSupport::Cache::Instrumentation.nested_multiread?
+                else
+                  true
+                end
+              end
+
+              def on_start(span, event, _id, payload)
+                key = payload[:key]
+                store = payload[:store]
+
+                mapping = MAPPING[event]
+
+                span.service = configuration[:cache_service]
+                span.resource = mapping[:resource]
+
+                span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
+                span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_CACHE)
+
+                if span.service != Datadog.configuration.service
+                  span.set_tag(Tracing::Contrib::Ext::Metadata::TAG_BASE_SERVICE, Datadog.configuration.service)
+                end
+
+                span.set_tag(Ext::TAG_CACHE_BACKEND, cache_backend(store))
+
+                span.set_tag('EVENT', event)
+
+                set_cache_key(span, key, mapping[:multi_key])
+              end
+
+              def set_cache_key(span, key, multi_key)
+                if multi_key
+                  keys = key.is_a?(Hash) ? key.keys : key # `write`s use Hashes, while `read`s use Arrays
+                  resolved_key = keys.map { |k| ::ActiveSupport::Cache.expand_cache_key(k) }
+                  cache_key = Core::Utils.truncate(resolved_key, Ext::QUANTIZE_CACHE_MAX_KEY_SIZE)
+                  span.set_tag(Ext::TAG_CACHE_KEY_MULTI, cache_key)
+                else
+                  resolved_key = ::ActiveSupport::Cache.expand_cache_key(key)
+                  cache_key = Core::Utils.truncate(resolved_key, Ext::QUANTIZE_CACHE_MAX_KEY_SIZE)
+                  span.set_tag(Ext::TAG_CACHE_KEY, cache_key)
+                end
+              end
+
+              # The name of the `store` is never saved by Rails.
+              # ActiveSupport looks up stores by converting a symbol into a 'require' path,
+              # then "camelizing" it for a `const_get` call:
+              # ```
+              # require "active_support/cache/#{store}"
+              # ActiveSupport::Cache.const_get(store.to_s.camelize)
+              # ```
+              # @see https://github.com/rails/rails/blob/261975dbef77731d2c76f907f1076c5132ebc0e4/activesupport/lib/active_support/cache.rb#L139-L149
+              #
+              # We can reverse engineer
+              # the original symbol by converting the class name to snake case:
+              # e.g. ActiveSupport::Cache::RedisStore -> active_support/cache/redis_store
+              # In this case, `redis_store` is the store name.
+              #
+              # Because there's no API retrieve only the class name
+              # (only `RedisStore`, and not `ActiveSupport::Cache::RedisStore`)
+              # the easiest way to retrieve the store symbol is to convert the fully qualified
+              # name using the Rails-provided method `#underscore`, which is the reverse of `#camelize`,
+              # then extracting the last part of it.
+              #
+              # Also, this method caches the store name, given this value will be retrieve
+              # multiple times and involves string manipulation.
+              def cache_backend(store)
+                # Cache the backend name to avoid the expensive string manipulation required to calculate it.
+                # DEV: We can't store it directly in the `store` object because it is a frozen String.
+                if (name = @cache_backend[store])
+                  return name
+                end
+
+                # DEV: #underscore is available through ActiveSupport, and is
+                # DEV: the exact reverse operation to `#camelize`.
+                # DEV: #demodulize is available through ActiveSupport, and is
+                # DEV: used to remove the module ('*::') part of a constant name.
+                name = ::ActiveSupport::Inflector.demodulize(store)
+                name = ::ActiveSupport::Inflector.underscore(name)
+
+                # Despite a given application only ever having 1-3 store types,
+                # we limit the size of the `@cache_backend` just in case, because
+                # the user can create custom Cache store classes themselves.
+                @cache_backend[store] = name if @cache_backend.size < 50
+
+                name
+              end
+
+              # DEV: There are two possibly interesting fields in the `on_finish` payload:
+              # | `:hit`             | If this read is a hit   |
+              # | `:super_operation` | `:fetch` if a read is done with [`fetch`][ActiveSupport::Cache::Store#fetch] |
+              # @see https://github.com/rails/rails/blob/b9d6759401c3d50a51e0a7650cb2331f4218d11f/guides/source/active_support_instrumentation.md?plain=1#L528-L529
+              # def on_finish(span, event, id, payload)
+              #   super
+              # end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/active_support/cache/instrumentation.rb

@@ -9,7 +9,28 @@ module Tracing
     module Contrib
       module ActiveSupport
         module Cache
-          # Defines instrumentation for ActiveSupport caching
+          # DEV-3.0: Backwards compatibility code for the 2.x gem series.
+          # DEV-3.0:
+          # DEV-3.0: `ActiveSupport::Cache` is now instrumented by subscribing to ActiveSupport::Notifications events.
+          # DEV-3.0: The implementation is located at {Datadog::Tracing::Contrib::ActiveSupport::Cache::Events::Cache}.
+          # DEV-3.0: The events emitted provide richer introspection points (e.g. events for cache misses on `#fetch`) while
+          # DEV-3.0: also ensuring we are using Rails' public API for improved compatibility.
+          # DEV-3.0:
+          # DEV-3.0: But one operation holds us back: `ActiveSupport::Cache::Store#fetch`.
+          # DEV-3.0: This method does not have an event that could produce an equivalent span to today's 2.x implementation.
+          # DEV-3.0: In 2.x, `#fetch` produces two separate, *nested* spans: one for the `#read` operation and
+          # DEV-3.0: another for the `#write` operation that is called internally by `#fetch` when the cache key needs
+          # DEV-3.0: to be populated on a cache miss.
+          # DEV-3.0: But the ActiveSupport events emitted by `#fetch` provide two *sibling* events for the`#read` and
+          # DEV-3.0: `#write` operations.
+          # DEV-3.0: Moving from nested spans to sibling spans would be a breaking change. One notable difference is
+          # DEV-3.0: that if the nested `#write` operation fails 2.x, the `#read` span is marked as an error. This would
+          # DEV-3.0: not be the case with sibling spans, and would be a very visible change.
+          # DEV-3.0:
+          # DEV-3.0: At the end of the day, moving to ActiveSupport events is the better approach, but we have to retain
+          # DEV-3.0: this last monkey patch for `#fetch` (and all its supporting code) to avoid a breaking change for now.
+          #
+          # Defines the deprecate  monkey-patch instrumentation for `ActiveSupport::Cache::Store#fetch`
           module Instrumentation
             module_function
 
@@ -43,6 +64,10 @@ def trace(action, store, key: nil, multi_key: nil)
             # In most of the cases, `#fetch()` and `#read()` calls are nested.
             # Instrument both does not add any value.
             # This method checks if these two operations are nested.
+            #
+            # DEV-3.0: We should not have these checks in the 3.x series because ActiveSupport events provide more
+            # DEV-3.0: legible nested spans. While using ActiveSupport events, the nested spans actually provide meaningful
+            # DEV-3.0: information.
             def nested_read?
               current_span = Tracing.active_span
               current_span && current_span.name == Ext::SPAN_CACHE && current_span.resource == Ext::RESOURCE_CACHE_GET
@@ -107,29 +132,7 @@ def dd_store_name
               end
             end
 
-            # Defines instrumentation for ActiveSupport cache reading
-            module Read
-              include InstanceMethods
-
-              def read(*args, &block)
-                return super if Instrumentation.nested_read?
-
-                Instrumentation.trace(Ext::RESOURCE_CACHE_GET, dd_store_name, key: args[0]) { super }
-              end
-            end
-
-            # Defines instrumentation for ActiveSupport cache reading of multiple keys
-            module ReadMulti
-              include InstanceMethods
-
-              def read_multi(*keys, &block)
-                return super if Instrumentation.nested_multiread?
-
-                Instrumentation.trace(Ext::RESOURCE_CACHE_MGET, dd_store_name, multi_key: keys) { super }
-              end
-            end
-
-            # Defines instrumentation for ActiveSupport cache fetching
+            # Defines the last legacy monkey-patching instrumentation for ActiveSupport cache fetching
             module Fetch
               include InstanceMethods
 
@@ -140,42 +143,13 @@ def fetch(*args, &block)
               end
             end
 
-            # Defines instrumentation for ActiveSupport cache fetching of multiple keys
-            module FetchMulti
-              include InstanceMethods
-
-              def fetch_multi(*args, &block)
-                return super if Instrumentation.nested_multiread?
-
-                keys = args[-1].instance_of?(Hash) ? args[0..-2] : args
-                Instrumentation.trace(Ext::RESOURCE_CACHE_MGET, dd_store_name, multi_key: keys) { super }
-              end
-            end
-
-            # Defines instrumentation for ActiveSupport cache writing
-            module Write
-              include InstanceMethods
-
-              def write(*args, &block)
-                Instrumentation.trace(Ext::RESOURCE_CACHE_SET, dd_store_name, key: args[0]) { super }
-              end
-            end
-
-            # Defines instrumentation for ActiveSupport cache writing of multiple keys
-            module WriteMulti
-              include InstanceMethods
-
-              def write_multi(hash, options = nil)
-                Instrumentation.trace(Ext::RESOURCE_CACHE_MSET, dd_store_name, multi_key: hash.keys) { super }
-              end
-            end
-
-            # Defines instrumentation for ActiveSupport cache deleting
-            module Delete
-              include InstanceMethods
-
-              def delete(*args, &block)
-                Instrumentation.trace(Ext::RESOURCE_CACHE_DELETE, dd_store_name, key: args[0]) { super }
+            # Backports the payload[:store] key present since Rails 6.1:
+            # https://github.com/rails/rails/commit/6fa747f2946ee244b2aab0cd8c3c064f05d950a5
+            module Store
+              def instrument(operation, key, options = nil)
+                polyfill_options = options&.dup || {}
+                polyfill_options[:store] = self.class.name
+                super(operation, key, polyfill_options)
               end
             end
           end