diff --git a/lib/ddtrace.rb b/lib/ddtrace.rb index 303b35fd3b1..f4be69e4ae4 100644 --- a/lib/ddtrace.rb +++ b/lib/ddtrace.rb @@ -1,11 +1,12 @@ -require 'ddtrace/monkey' +require 'ddtrace/registry' require 'ddtrace/pin' require 'ddtrace/tracer' require 'ddtrace/error' # \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: # @@ -24,8 +25,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/base.rb b/lib/ddtrace/contrib/base.rb new file mode 100644 index 00000000000..17407ba5e89 --- /dev/null +++ b/lib/ddtrace/contrib/base.rb @@ -0,0 +1,14 @@ +require 'ddtrace/registry' + +module Datadog + module Contrib + module Base + def self.included(base) + base.class_eval do + include Registry::Registerable + module_function :register_as + end + end + end + end +end 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/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/redis/patcher.rb b/lib/ddtrace/contrib/redis/patcher.rb index 87c0175cfbf..5ca2564f7d5 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/monkey.rb b/lib/ddtrace/monkey.rb index adcfcdf1039..16f7887d531 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/grape/patcher' @@ -12,39 +13,29 @@ module Datadog # Monkey is used for monkey-patching 3rd party libs. module Monkey - @patched = [] - @autopatch_modules = { - elasticsearch: true, - http: true, - redis: true, - grape: 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, - 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 @@ -57,16 +48,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 @@ -80,5 +66,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..3324e8ceed5 --- /dev/null +++ b/lib/ddtrace/registry.rb @@ -0,0 +1,41 @@ +require_relative 'registry/registerable' + +module Datadog + 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..a8b730f73f8 --- /dev/null +++ b/lib/ddtrace/registry/registerable.rb @@ -0,0 +1,12 @@ +module Datadog + class Registry + module Registerable + 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 diff --git a/test/registry_test.rb b/test/registry_test.rb new file mode 100644 index 00000000000..6330df43c46 --- /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