Add support for AWS

Appraisals

@@ -120,6 +120,7 @@ if RUBY_VERSION >= '2.2.2' && RUBY_PLATFORM != 'java'
     gem 'sqlite3'
     gem 'activerecord'
     gem 'sidekiq'
+    gem 'aws-sdk'
   end
 else
   appraise 'contrib-old' do
@@ -132,5 +133,6 @@ else
     gem 'sqlite3'
     gem 'activerecord', '3.2.22.5'
     gem 'sidekiq', '4.0.0'
+    gem 'aws-sdk'
   end
 end

Rakefile

@@ -51,7 +51,7 @@ namespace :test do
     t.test_files = FileList['test/contrib/rails/**/*disable_env*_test.rb']
   end
 
-  [:elasticsearch, :http, :redis, :sinatra, :sidekiq, :rack, :grape].each do |contrib|
+  [:elasticsearch, :http, :redis, :sinatra, :sidekiq, :rack, :grape, :aws].each do |contrib|
     Rake::TestTask.new(contrib) do |t|
       t.libs << %w[test lib]
       t.test_files = FileList["test/contrib/#{contrib}/*_test.rb"]
@@ -140,12 +140,14 @@ task :ci do
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake test:sidekiq'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake test:rack'
     sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake test:grape'
+    sh 'rvm $MRI_VERSIONS --verbose do appraisal contrib rake test:aws'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:monkey'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:elasticsearch'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:http'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:redis'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:sinatra'
     sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:rack'
+    sh 'rvm $MRI_OLD_VERSIONS --verbose do appraisal contrib-old rake test:aws'
     sh 'rvm $SIDEKIQ_OLD_VERSIONS --verbose do appraisal contrib-old rake test:sidekiq'
   when 2
     sh 'rvm $RAILS3_VERSIONS --verbose do appraisal rails30-postgres rake test:rails'

gemfiles/contrib.gemfile

@@ -12,5 +12,6 @@ gem "sinatra"
 gem "sqlite3"
 gem "activerecord"
 gem "sidekiq"
+gem "aws-sdk"
 
 gemspec path: "../"

lib/ddtrace/contrib/aws/instrumentation.rb

@@ -0,0 +1,43 @@
+module Datadog
+  module Contrib
+    module Aws
+      # A Seahorse::Client::Plugin that enables instrumentation for all AWS services
+      class Instrumentation < Seahorse::Client::Plugin
+        def add_handlers(handlers, _)
+          handlers.add(Handler, step: :validate)
+        end
+      end
+
+      # Generates Spans for all interactions with AWS
+      class Handler < Seahorse::Client::Handler
+        def call(context)
+          pin = Datadog::Pin.get_from(::Aws)
+
+          return @handler.call(context) unless pin && pin.tracer
+
+          pin.tracer.trace(RESOURCE) do |span|
+            result = @handler.call(context)
+            annotate!(span, pin, ParsedContext.new(context))
+            result
+          end
+        end
+
+        private
+
+        def annotate!(span, pin, context)
+          span.service = pin.service
+          span.span_type = pin.app_type
+          span.name = context.safely(:resource)
+          span.resource = RESOURCE
+          span.set_tag('aws.agent', AGENT)
+          span.set_tag('aws.operation', context.safely(:operation))
+          span.set_tag('aws.region', context.safely(:region))
+          span.set_tag('path', context.safely(:path))
+          span.set_tag('host', context.safely(:host))
+          span.set_tag(Ext::HTTP::METHOD, context.safely(:http_method))
+          span.set_tag(Ext::HTTP::STATUS_CODE, context.safely(:status_code))
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/aws/parsed_context.rb

@@ -0,0 +1,56 @@
+module Datadog
+  module Contrib
+    module Aws
+      # A wrapper around Seahorse::Client::RequestContext
+      class ParsedContext
+        def initialize(context)
+          @context = context
+        end
+
+        def safely(attr, fallback = nil)
+          public_send(attr) rescue fallback
+        end
+
+        def resource
+          "#{service}.#{operation}"
+        end
+
+        def operation
+          context.operation_name
+        end
+
+        def status_code
+          context.http_response.status_code
+        end
+
+        def http_method
+          context.http_request.http_method
+        end
+
+        def region
+          context.client.config.region
+        end
+
+        def retry_attempts
+          context.retries
+        end
+
+        def path
+          context.http_request.endpoint.path
+        end
+
+        def host
+          context.http_request.endpoint.host
+        end
+
+        private
+
+        attr_reader :context
+
+        def service
+          context.client.class.to_s.split('::')[1].downcase
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/aws/patcher.rb

@@ -0,0 +1,61 @@
+module Datadog
+  module Contrib
+    module Aws
+      COMPATIBLE_WITH = Gem::Version.new('2.0.0')
+      SERVICE = 'aws'.freeze
+      AGENT = 'aws-sdk-ruby'.freeze
+      RESOURCE = 'aws.command'.freeze
+
+      # Responsible for hooking the instrumentation into aws-sdk
+      module Patcher
+        class << self
+          @patched = false
+
+          def patch
+            return @patched if patched? || !compatible?
+
+            require 'ddtrace/ext/app_types'
+            require 'ddtrace/contrib/aws/parsed_context'
+            require 'ddtrace/contrib/aws/instrumentation'
+
+            add_pin
+            add_plugin(Seahorse::Client::Base, *loaded_constants)
+
+            @patched = true
+          rescue => e
+            Datadog::Tracer.log.error("Unable to apply AWS integration: #{e}")
+            @patched
+          end
+
+          def patched?
+            @patched
+          end
+
+          private
+
+          def compatible?
+            return unless defined?(::Aws::VERSION)
+
+            Gem::Version.new(::Aws::VERSION) >= COMPATIBLE_WITH
+          end
+
+          def add_pin
+            Pin.new(SERVICE, app_type: Ext::AppTypes::WEB).tap do |pin|
+              pin.onto(::Aws)
+            end
+          end
+
+          def add_plugin(*targets)
+            targets.each { |klass| klass.add_plugin(Instrumentation) }
+          end
+
+          def loaded_constants
+            ::Aws::SERVICE_MODULE_NAMES
+              .reject { |klass| ::Aws.autoload?(klass) }
+              .map { |klass| ::Aws.const_get("#{klass}::Client") }
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/monkey.rb

@@ -8,6 +8,7 @@
 require 'ddtrace/contrib/grape/patcher'
 require 'ddtrace/contrib/redis/patcher'
 require 'ddtrace/contrib/http/patcher'
+require 'ddtrace/contrib/aws/patcher'
 
 module Datadog
   # Monkey is used for monkey-patching 3rd party libs.
@@ -18,6 +19,7 @@ module Monkey
       http: true,
       redis: true,
       grape: true,
+      aws: true,
       active_record: false
     }
     # Patchers should expose 2 methods:
@@ -29,6 +31,7 @@ module Monkey
                   http: Datadog::Contrib::HTTP::Patcher,
                   redis: Datadog::Contrib::Redis::Patcher,
                   grape: Datadog::Contrib::Grape::Patcher,
+                  aws: Datadog::Contrib::Aws::Patcher,
                   active_record: Datadog::Contrib::ActiveRecord::Patcher }
     @mutex = Mutex.new
 

test/contrib/aws/instrumentation_test.rb

@@ -0,0 +1,63 @@
+require 'helper'
+require 'aws-sdk'
+require 'ddtrace'
+require 'ddtrace/contrib/aws/patcher'
+require 'ddtrace/ext/http'
+
+module Datadog
+  module Contrib
+    module Aws
+      class InstrumentationTest < Minitest::Test
+        def setup
+          @tracer = enable_tracer
+          client = ::Aws::S3::Client.new(stub_responses: true)
+          client.list_buckets
+
+          @span = @tracer.writer.spans.find { |span| span.resource == RESOURCE }
+        end
+
+        def test_span_service
+          assert_equal(pin.service, @span.service)
+        end
+
+        def test_span_type
+          assert_equal(pin.app_type, @span.span_type)
+        end
+
+        def test_span_resource
+          assert_equal('aws.command', @span.resource)
+        end
+
+        def test_span_tags
+          assert_equal('aws-sdk-ruby', @span.get_tag('aws.agent'))
+          assert_equal('list_buckets', @span.get_tag('aws.operation'))
+          assert_equal('us-stubbed-1', @span.get_tag('aws.region'))
+          assert_equal('/', @span.get_tag('path'))
+          assert_equal('s3.us-stubbed-1.amazonaws.com', @span.get_tag('host'))
+          assert_equal('GET', @span.get_tag(Datadog::Ext::HTTP::METHOD))
+          assert_equal('200', @span.get_tag(Datadog::Ext::HTTP::STATUS_CODE))
+        end
+
+        def test_client_response
+          client = ::Aws::S3::Client.new(
+            stub_responses: { list_buckets: { buckets: [{ name: 'bucket1' }] } }
+          )
+
+          buckets = client.list_buckets.buckets.map(&:name)
+          assert_equal(['bucket1'], buckets)
+        end
+
+        private
+
+        def enable_tracer
+          Patcher.patch
+          get_test_tracer.tap { |tracer| pin.tracer = tracer }
+        end
+
+        def pin
+          ::Aws.datadog_pin
+        end
+      end
+    end
+  end
+end

test/contrib/aws/parsed_context_test.rb

@@ -0,0 +1,50 @@
+require 'helper'
+require 'aws-sdk'
+require 'ddtrace/contrib/aws/parsed_context'
+
+class ParsedContextTest < Minitest::Test
+  def setup
+    @client = Aws::S3::Client.new(region: 'us-west-2', stub_responses: true)
+    response = @client.list_buckets
+    @context = Datadog::Contrib::Aws::ParsedContext.new(response.context)
+  end
+
+  def test_resource
+    assert_equal('s3.list_buckets', @context.resource)
+  end
+
+  def test_operation
+    assert_equal(:list_buckets, @context.operation)
+  end
+
+  def test_status_code
+    assert_equal(200, @context.status_code)
+  end
+
+  def test_http_method
+    assert_equal('GET', @context.http_method)
+  end
+
+  def test_region
+    assert_equal('us-west-2', @context.region)
+  end
+
+  def test_retry_attempts
+    assert_equal(0, @context.retry_attempts)
+  end
+
+  def test_path
+    assert_equal('/', @context.path)
+  end
+
+  def test_host
+    assert_equal('s3-us-west-2.amazonaws.com', @context.host)
+  end
+
+  def test_param_safety
+    @context.stub :resource, -> { raise } do
+      actual = @context.safely(:resource, 'fallback_name')
+      assert_equal('fallback_name', actual)
+    end
+  end
+end