diff --git a/lib/ddtrace.rb b/lib/ddtrace.rb index 038d6fc3380..449ebf8cb01 100644 --- a/lib/ddtrace.rb +++ b/lib/ddtrace.rb @@ -1,4 +1,4 @@ -require 'ddtrace/monkey' +require 'ddtrace/registry' require 'ddtrace/pin' require 'ddtrace/tracer' require 'ddtrace/error' @@ -6,7 +6,8 @@ # \Datadog global namespace that includes all tracing functionality for Tracer and Span classes. module Datadog - @tracer = Datadog::Tracer.new() + @tracer = Tracer.new + @registry = Registry.new # Default tracer that can be used as soon as +ddtrace+ is required: # @@ -25,8 +26,14 @@ module Datadog def self.tracer @tracer end + + def self.registry + @registry + end end +require 'ddtrace/monkey' + # Datadog auto instrumentation for frameworks if defined?(Rails::VERSION) if !ENV['DISABLE_DATADOG_RAILS'] diff --git a/lib/ddtrace/contrib/active_record/patcher.rb b/lib/ddtrace/contrib/active_record/patcher.rb index d088533a39c..efad7816de0 100644 --- a/lib/ddtrace/contrib/active_record/patcher.rb +++ b/lib/ddtrace/contrib/active_record/patcher.rb @@ -4,6 +4,9 @@ module ActiveRecord # Patcher enables patching of 'active_record' module. # This is used in monkey.rb to manually apply patches module Patcher + include Base + register_as :active_record, auto_patch: false + @patched = false module_function diff --git a/lib/ddtrace/contrib/aws/patcher.rb b/lib/ddtrace/contrib/aws/patcher.rb index 58c233ae4cd..16563b95a2e 100644 --- a/lib/ddtrace/contrib/aws/patcher.rb +++ b/lib/ddtrace/contrib/aws/patcher.rb @@ -7,6 +7,9 @@ module Aws # Responsible for hooking the instrumentation into aws-sdk module Patcher + include Base + register_as :aws, auto_patch: true + @patched = false class << self diff --git a/lib/ddtrace/contrib/base.rb b/lib/ddtrace/contrib/base.rb new file mode 100644 index 00000000000..6084221eb3b --- /dev/null +++ b/lib/ddtrace/contrib/base.rb @@ -0,0 +1,12 @@ +require 'ddtrace/registry' + +module Datadog + module Contrib + # Base provides features that are shared across all integrations + module Base + def self.included(base) + base.send(:include, Registry::Registerable) + end + end + end +end diff --git a/lib/ddtrace/contrib/dalli/patcher.rb b/lib/ddtrace/contrib/dalli/patcher.rb index 83f7dac44da..827ed8277bb 100644 --- a/lib/ddtrace/contrib/dalli/patcher.rb +++ b/lib/ddtrace/contrib/dalli/patcher.rb @@ -8,6 +8,9 @@ module Dalli # Responsible for hooking the instrumentation into `dalli` module Patcher + include Base + register_as :dalli, auto_patch: true + @patched = false class << self diff --git a/lib/ddtrace/contrib/elasticsearch/patcher.rb b/lib/ddtrace/contrib/elasticsearch/patcher.rb index ddd6fdc43e8..609159861ea 100644 --- a/lib/ddtrace/contrib/elasticsearch/patcher.rb +++ b/lib/ddtrace/contrib/elasticsearch/patcher.rb @@ -14,6 +14,9 @@ module Elasticsearch # Patcher enables patching of 'elasticsearch/transport' module. # This is used in monkey.rb to automatically apply patches module Patcher + include Base + register_as :elasticsearch, auto_patch: true + @patched = false module_function diff --git a/lib/ddtrace/contrib/faraday/patcher.rb b/lib/ddtrace/contrib/faraday/patcher.rb index 437719cee77..e09f8a7e3de 100644 --- a/lib/ddtrace/contrib/faraday/patcher.rb +++ b/lib/ddtrace/contrib/faraday/patcher.rb @@ -6,6 +6,9 @@ module Faraday # Responsible for hooking the instrumentation into faraday module Patcher + include Base + register_as :faraday, auto_patch: true + @patched = false class << self diff --git a/lib/ddtrace/contrib/grape/patcher.rb b/lib/ddtrace/contrib/grape/patcher.rb index f24e7aa82ac..259d3af646b 100644 --- a/lib/ddtrace/contrib/grape/patcher.rb +++ b/lib/ddtrace/contrib/grape/patcher.rb @@ -6,6 +6,9 @@ module Grape # Patcher that introduces more instrumentation for Grape endpoints, so that # new signals are executed at the beginning of each step (filters, render and run) module Patcher + include Base + register_as :grape, auto_patch: true + @patched = false module_function diff --git a/lib/ddtrace/contrib/http/patcher.rb b/lib/ddtrace/contrib/http/patcher.rb index d7f7b8ebe77..fda7999e12d 100644 --- a/lib/ddtrace/contrib/http/patcher.rb +++ b/lib/ddtrace/contrib/http/patcher.rb @@ -51,6 +51,9 @@ def should_skip_distributed_tracing?(pin) # Patcher enables patching of 'net/http' module. # This is used in monkey.rb to automatically apply patches module Patcher + include Base + register_as :http, auto_patch: true + @patched = false module_function diff --git a/lib/ddtrace/contrib/mongodb/patcher.rb b/lib/ddtrace/contrib/mongodb/patcher.rb index 1b4fafe21a9..97c41770d2e 100644 --- a/lib/ddtrace/contrib/mongodb/patcher.rb +++ b/lib/ddtrace/contrib/mongodb/patcher.rb @@ -11,6 +11,9 @@ module MongoDB # Use the `Datadog::Monkey.patch_module(:mongodb)` to activate tracing for # this module. module Patcher + include Base + register_as :mongo, auto_patch: true + @patched = false module_function diff --git a/lib/ddtrace/contrib/redis/patcher.rb b/lib/ddtrace/contrib/redis/patcher.rb index cb744b96a9a..a14faa9ca9a 100644 --- a/lib/ddtrace/contrib/redis/patcher.rb +++ b/lib/ddtrace/contrib/redis/patcher.rb @@ -9,6 +9,9 @@ module Redis # Patcher enables patching of 'redis' module. # This is used in monkey.rb to automatically apply patches module Patcher + include Base + register_as :redis, auto_patch: true + @patched = false module_function diff --git a/lib/ddtrace/contrib/resque/patcher.rb b/lib/ddtrace/contrib/resque/patcher.rb index f90eaffca2c..80a423f105e 100644 --- a/lib/ddtrace/contrib/resque/patcher.rb +++ b/lib/ddtrace/contrib/resque/patcher.rb @@ -5,6 +5,9 @@ module Resque # Patcher for Resque integration - sets up the pin for the integration module Patcher + include Base + register_as :resque, auto_patch: true + @patched = false class << self diff --git a/lib/ddtrace/contrib/sucker_punch/patcher.rb b/lib/ddtrace/contrib/sucker_punch/patcher.rb index 99ffa0dc335..00b1ccbe1d4 100644 --- a/lib/ddtrace/contrib/sucker_punch/patcher.rb +++ b/lib/ddtrace/contrib/sucker_punch/patcher.rb @@ -6,6 +6,9 @@ module SuckerPunch # Responsible for hooking the instrumentation into `sucker_punch` module Patcher + include Base + register_as :sucker_punch, auto_patch: true + @patched = false module_function diff --git a/lib/ddtrace/monkey.rb b/lib/ddtrace/monkey.rb index 53bf6d7d1f2..44996cc95dc 100644 --- a/lib/ddtrace/monkey.rb +++ b/lib/ddtrace/monkey.rb @@ -3,6 +3,7 @@ # We import all patchers for every module we support, but this is fine # because patchers do not include any 3rd party module nor even our # patching code, which is required on demand, when patching. +require 'ddtrace/contrib/base' require 'ddtrace/contrib/active_record/patcher' require 'ddtrace/contrib/elasticsearch/patcher' require 'ddtrace/contrib/faraday/patcher' @@ -18,51 +19,30 @@ module Datadog # Monkey is used for monkey-patching 3rd party libs. module Monkey - @patched = [] - @autopatch_modules = { - elasticsearch: true, - http: true, - redis: true, - grape: true, - faraday: true, - aws: true, - sucker_punch: true, - mongo: true, - dalli: true, - resque: true, - active_record: false - } # Patchers should expose 2 methods: # - patch, which applies our patch if needed. Should be idempotent, # can be call twice but should just do nothing the second time. # - patched?, which returns true if the module has been succesfully # patched (patching might have failed if requirements were not here) - @patchers = { elasticsearch: Datadog::Contrib::Elasticsearch::Patcher, - http: Datadog::Contrib::HTTP::Patcher, - redis: Datadog::Contrib::Redis::Patcher, - grape: Datadog::Contrib::Grape::Patcher, - faraday: Datadog::Contrib::Faraday::Patcher, - aws: Datadog::Contrib::Aws::Patcher, - sucker_punch: Datadog::Contrib::SuckerPunch::Patcher, - mongo: Datadog::Contrib::MongoDB::Patcher, - dalli: Datadog::Contrib::Dalli::Patcher, - resque: Datadog::Contrib::Resque::Patcher, - active_record: Datadog::Contrib::ActiveRecord::Patcher } + @mutex = Mutex.new + @registry = Datadog.registry module_function + attr_accessor :registry + def autopatch_modules - @autopatch_modules.clone + registry.to_h end def patch_all - patch @autopatch_modules + patch(autopatch_modules) end def patch_module(m) @mutex.synchronize do - patcher = @patchers[m] + patcher = registry[m] raise "Unsupported module #{m}" unless patcher patcher.patch end @@ -75,16 +55,11 @@ def patch(modules) end def get_patched_modules - patched = autopatch_modules - @patchers.each do |k, v| + registry.each_with_object({}) do |entry, patched| @mutex.synchronize do - if v - patcher = @patchers[k] - patched[k] = patcher.patched? if patcher - end + patched[entry.name] = entry.klass.patched? end end - patched end def without_warnings @@ -98,5 +73,9 @@ def without_warnings $VERBOSE = v end end + + class << self + attr_accessor :registry + end end end diff --git a/lib/ddtrace/registry.rb b/lib/ddtrace/registry.rb new file mode 100644 index 00000000000..5c8c26a536e --- /dev/null +++ b/lib/ddtrace/registry.rb @@ -0,0 +1,42 @@ +require_relative 'registry/registerable' + +module Datadog + # Registry provides insertion/retrieval capabilities for integrations + class Registry + include Enumerable + + Entry = Struct.new(:name, :klass, :auto_patch) + + def initialize + @data = {} + @mutex = Mutex.new + end + + def add(name, klass, auto_patch = false) + @mutex.synchronize do + @data[name] = Entry.new(name, klass, auto_patch).freeze + end + end + + def each + @mutex.synchronize do + @data.each { |_, entry| yield(entry) } + end + end + + def [](name) + @mutex.synchronize do + entry = @data[name] + entry.klass if entry + end + end + + def to_h + @mutex.synchronize do + @data.each_with_object({}) do |(_, entry), hash| + hash[entry.name] = entry.auto_patch + end + end + end + end +end diff --git a/lib/ddtrace/registry/registerable.rb b/lib/ddtrace/registry/registerable.rb new file mode 100644 index 00000000000..4fd63703cf0 --- /dev/null +++ b/lib/ddtrace/registry/registerable.rb @@ -0,0 +1,20 @@ +module Datadog + class Registry + # Registerable provides a convenience method for self-registering + module Registerable + def self.included(base) + base.singleton_class.send(:include, ClassMethods) + end + + # ClassMethods + module ClassMethods + def register_as(name, options = {}) + registry = options.fetch(:registry, Datadog.registry) + auto_patch = options.fetch(:auto_patch, false) + + registry.add(name, self, auto_patch) + end + end + end + end +end diff --git a/test/registry_test.rb b/test/registry_test.rb new file mode 100644 index 00000000000..97649c9f2cd --- /dev/null +++ b/test/registry_test.rb @@ -0,0 +1,60 @@ +require 'minitest/autorun' +require 'ddtrace' +require 'ddtrace/registry' + +module Datadog + class RegistryTest < Minitest::Test + def test_object_retrieval + registry = Registry.new + + object1 = Object.new + object2 = Object.new + + registry.add(:object1, object1) + registry.add(:object2, object2) + + assert_same(object1, registry[:object1]) + assert_same(object2, registry[:object2]) + end + + def test_hash_coercion + registry = Registry.new + + object1 = Object.new + object2 = Object.new + + registry.add(:object1, object1, true) + registry.add(:object2, object2, false) + + assert_equal({ object1: true, object2: false }, registry.to_h) + end + + def test_enumeration + registry = Registry.new + + object1 = Object.new + object2 = Object.new + + registry.add(:object1, object1, true) + registry.add(:object2, object2, false) + + assert(registry.respond_to?(:each)) + assert_kind_of(Enumerable, registry) + + # Enumerable#map + objects = registry.map(&:klass) + assert_kind_of(Array, objects) + assert_equal(2, objects.size) + assert_includes(objects, object1) + assert_includes(objects, object2) + end + + def test_registry_entry + entry = Registry::Entry.new(:array, Array, true) + + assert_equal(:array, entry.name) + assert_equal(Array, entry.klass) + assert_equal(true, entry.auto_patch) + end + end +end