Add Datadog::Registry

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']

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

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

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

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

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

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

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

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'
@@ -15,45 +16,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,
-      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,
-                  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
@@ -66,16 +52,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
@@ -89,5 +70,9 @@ def without_warnings
         $VERBOSE = v
       end
     end
+
+    class << self
+      attr_accessor :registry
+    end
   end
 end

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

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

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