Refactor Pin to automatically set service info

p-lambert · p-lambert · commit e2f2a5455ea8 · 2018-01-12T10:43:19.000-05:00
diff --git a/lib/ddtrace/configuration/pin_setup.rb b/lib/ddtrace/configuration/pin_setup.rb
@@ -22,7 +22,7 @@ def call
 
       attr_reader :pin, :opts
 
-      ATTRS = [:service_name, :app, :tags, :app_type, :name, :tracer].freeze
+      ATTRS = [:app, :tags, :app_type, :name, :tracer, :service_name].freeze
 
       private_constant :ATTRS
     end
diff --git a/lib/ddtrace/contrib/aws/patcher.rb b/lib/ddtrace/contrib/aws/patcher.rb
@@ -23,7 +23,6 @@ def patch
 
             add_pin
             add_plugin(Seahorse::Client::Base, *loaded_constants)
-            Datadog.tracer.set_service_info(get_option(:service_name), 'aws', Ext::AppTypes::WEB)
 
             @patched = true
           rescue => e
@@ -38,9 +37,9 @@ def patched?
           private
 
           def add_pin
-            Pin.new(get_option(:service_name), app_type: Ext::AppTypes::WEB).tap do |pin|
-              pin.onto(::Aws)
-            end
+            Pin
+              .new(get_option(:service_name), app: 'aws', app_type: Ext::AppTypes::WEB)
+              .onto(::Aws)
           end
 
           def add_plugin(*targets)
diff --git a/lib/ddtrace/contrib/dalli/patcher.rb b/lib/ddtrace/contrib/dalli/patcher.rb
@@ -22,7 +22,6 @@ def patch
 
             add_pin!
             Instrumentation.patch!
-            Datadog.tracer.set_service_info(get_option(:service_name), 'dalli', Ext::AppTypes::CACHE)
 
             @patched = true
           rescue => e
@@ -43,9 +42,9 @@ def compatible?
           end
 
           def add_pin!
-            Pin.new(get_option(:service_name), app_type: Ext::AppTypes::CACHE).tap do |pin|
-              pin.onto(::Dalli)
-            end
+            Pin
+              .new(get_option(:service_name), app: 'dalli', app_type: Ext::AppTypes::CACHE)
+              .onto(::Dalli)
           end
         end
       end
diff --git a/lib/ddtrace/contrib/elasticsearch/patcher.rb b/lib/ddtrace/contrib/elasticsearch/patcher.rb
@@ -34,7 +34,6 @@ def patch
               require 'ddtrace/contrib/elasticsearch/quantize'
 
               patch_elasticsearch_transport_client()
-              Datadog.tracer.set_service_info(get_option(:service_name), 'elasticsearch', Ext::AppTypes::DB)
 
               @patched = true
             rescue StandardError => e
diff --git a/lib/ddtrace/contrib/faraday/patcher.rb b/lib/ddtrace/contrib/faraday/patcher.rb
@@ -31,7 +31,6 @@ def patch
 
             add_pin
             add_middleware
-            register_service(get_option(:service_name))
 
             @patched = true
           rescue => e
@@ -56,10 +55,12 @@ def compatible?
           end
 
           def add_pin
-            Pin.new(get_option(:service_name), app_type: Ext::AppTypes::WEB).tap do |pin|
-              pin.onto(::Faraday)
-              pin.tracer = get_option(:tracer)
-            end
+            Pin.new(
+              get_option(:service_name),
+              app: 'faraday',
+              app_type: Ext::AppTypes::WEB,
+              tracer: get_option(:tracer)
+            ).onto(::Faraday)
           end
 
           def add_middleware
diff --git a/lib/ddtrace/contrib/grape/patcher.rb b/lib/ddtrace/contrib/grape/patcher.rb
@@ -35,9 +35,6 @@ def patch
               service = get_option(:service_name)
               pin = Datadog::Pin.new(service, app: 'grape', app_type: Datadog::Ext::AppTypes::WEB)
               pin.onto(::Grape)
-              if pin.tracer && pin.service
-                pin.tracer.set_service_info(pin.service, 'grape', pin.app_type)
-              end
 
               # subscribe to ActiveSupport events
               Datadog::Contrib::Grape::Endpoint.subscribe()
diff --git a/lib/ddtrace/contrib/mongodb/patcher.rb b/lib/ddtrace/contrib/mongodb/patcher.rb
@@ -67,9 +67,6 @@ def initialize(*args, &blk)
               service = Datadog.configuration[:mongo][:service_name]
               pin = Datadog::Pin.new(service, app: APP, app_type: Datadog::Ext::AppTypes::DB)
               pin.onto(self)
-              if pin.tracer && pin.service
-                pin.tracer.set_service_info(pin.service, 'mongodb', pin.app_type)
-              end
             end
 
             def datadog_pin
diff --git a/lib/ddtrace/contrib/redis/patcher.rb b/lib/ddtrace/contrib/redis/patcher.rb
@@ -54,7 +54,6 @@ def datadog_pin
         end
 
         # rubocop:disable Metrics/MethodLength
-        # rubocop:disable Metrics/BlockLength
         def patch_redis_client
           ::Redis::Client.class_eval do
             alias_method :initialize_without_datadog, :initialize
@@ -66,9 +65,6 @@ def initialize(*args)
               service = Datadog.configuration[:redis][:service_name]
               pin = Datadog::Pin.new(service, app: 'redis', app_type: Datadog::Ext::AppTypes::DB)
               pin.onto(self)
-              if pin.tracer && pin.service
-                pin.tracer.set_service_info(pin.service, pin.app, pin.app_type)
-              end
               initialize_without_datadog(*args)
             end
 
diff --git a/lib/ddtrace/contrib/resque/patcher.rb b/lib/ddtrace/contrib/resque/patcher.rb
@@ -40,9 +40,9 @@ def patched?
           private
 
           def add_pin
-            Pin.new(get_option(:service_name), app_type: Ext::AppTypes::WORKER).tap do |pin|
-              pin.onto(::Resque)
-            end
+            Pin
+              .new(get_option(:service_name), app: 'resque', app_type: Ext::AppTypes::WORKER)
+              .onto(::Resque)
           end
         end
       end
diff --git a/lib/ddtrace/contrib/resque/resque_job.rb b/lib/ddtrace/contrib/resque/resque_job.rb
@@ -25,7 +25,6 @@ def around_perform(*args)
 Resque.before_first_fork do
   pin = Datadog::Pin.get_from(Resque)
   next unless pin && pin.tracer
-  pin.tracer.set_service_info(pin.service, 'resque', Datadog::Ext::AppTypes::WORKER)
 
   # Create SyncWriter instance before forking
   sync_writer = Datadog::SyncWriter.new(transport: pin.tracer.writer.transport)
diff --git a/lib/ddtrace/contrib/sucker_punch/patcher.rb b/lib/ddtrace/contrib/sucker_punch/patcher.rb
@@ -24,7 +24,6 @@ def patch
           add_pin!
           ExceptionHandler.patch!
           Instrumentation.patch!
-          Datadog.tracer.set_service_info(get_option(:service_name), 'sucker_punch', Ext::AppTypes::WORKER)
 
           @patched = true
         rescue => e
@@ -43,7 +42,7 @@ def compatible?
         end
 
         def add_pin!
-          Pin.new(get_option(:service_name), app_type: Ext::AppTypes::WORKER).tap do |pin|
+          Pin.new(get_option(:service_name), app: 'sucker_punch', app_type: Ext::AppTypes::WORKER).tap do |pin|
             pin.onto(::SuckerPunch)
           end
         end
diff --git a/lib/ddtrace/pin.rb b/lib/ddtrace/pin.rb
@@ -18,17 +18,14 @@ def self.get_from(obj)
     attr_accessor :tracer
     attr_accessor :config
 
-    alias service_name= service=
-
-    # [ruby19] named parameters would be more idiomatic here, but would break backward compatibility
-    def initialize(service, options = { app: nil, tags: nil, app_type: nil, tracer: nil, config: nil })
-      @service = service
-      @app = options.fetch(:app, nil)
-      @tags = options.fetch(:tags, nil)
-      @app_type = options.fetch(:app_type, nil)
+    def initialize(service, options = {})
+      @app = options[:app]
+      @tags = options[:tags]
+      @app_type = options[:app_type]
       @name = nil # this would rarely be overriden as it's really span-specific
       @tracer = options[:tracer] || Datadog.tracer
-      @config = options.fetch(:config, nil)
+      @config = options[:config]
+      self.service = service
     end
 
     def enabled?
@@ -57,6 +54,16 @@ def datadog_pin
       obj.datadog_pin = self
     end
 
+    def service=(name)
+      if name && app && app_type
+        tracer.set_service_info(name, app, app_type)
+      end
+
+      @service = name
+    end
+
+    alias service_name= service=
+
     def to_s
       "Pin(service:#{@service},app:#{@app},app_type:#{@app_type},name:#{@name})"
     end
diff --git a/test/configuration/pin_setup_test.rb b/test/configuration/pin_setup_test.rb
@@ -12,7 +12,7 @@ def setup
       end
 
       def test_setting_options
-        custom_tracer = Object.new
+        custom_tracer = get_test_tracer
 
         custom_options = {
           service_name: 'my-service',
diff --git a/test/contrib/aws/instrumentation_test.rb b/test/contrib/aws/instrumentation_test.rb
@@ -17,7 +17,7 @@ def setup
         def test_pin_defaults
           assert_equal('aws', pin.service)
           assert_equal('web', pin.app_type)
-          assert_nil(pin.app)
+          assert_equal('aws', pin.app)
           assert_nil(pin.name)
           assert_nil(pin.config)
         end
diff --git a/test/contrib/redis/redis_test.rb b/test/contrib/redis/redis_test.rb
@@ -127,13 +127,14 @@ def test_quantize
   end
 
   def test_service_name
-    service_name = 'test'
     driver = Redis.new(host: REDIS_HOST, port: REDIS_PORT, driver: :ruby)
-    pin = Datadog::Pin.get_from(driver)
-    pin.tracer = @tracer
-    pin.service = service_name
-    @tracer.writer.services() # empty queue
-    @tracer.set_service_info("redis-#{service_name}", 'redis', Datadog::Ext::AppTypes::CACHE)
+    @tracer.writer.services # empty queue
+    Datadog.configure(
+      driver,
+      service_name: 'redis-test',
+      tracer: @tracer,
+      app_type: Datadog::Ext::AppTypes::CACHE
+    )
     driver.set 'FOO', 'bar'
     services = @tracer.writer.services
     assert_equal(1, services.length)
diff --git a/test/pin_test.rb b/test/pin_test.rb
@@ -67,4 +67,23 @@ def test_custom_getter_setter
     pin.onto(obj)
     assert_equal('The PIN is set!', Datadog::Pin.get_from(obj))
   end
+
+  def test_service_info_update
+    tracer = get_test_tracer
+    Datadog::Pin.new('test-service', app: 'test-app', app_type: 'test-type', tracer: tracer)
+
+    assert(tracer.services.key?('test-service'))
+    assert_equal(
+      { 'app' => 'test-app', 'app_type' => 'test-type' },
+      tracer.services['test-service']
+    )
+  end
+
+  def test_service_info_update_with_missing_params
+    tracer = get_test_tracer
+    # instantiating `Pin` without an `app` param (which is necessary for service info)
+    Datadog::Pin.new('test-service', app_type: 'test-type', tracer: tracer)
+
+    assert(tracer.services.empty?)
+  end
 end
