ActiveSupport::Notifications instrumentation of aws sdk client calls (#37)

Author: alextwoods

README.md

@@ -119,3 +119,32 @@ message = MyMailer.send_email(options)
 message['X-SES-FROM-ARN'] = 'arn:aws:ses:us-west-2:012345678910:identity/[email protected]'
 message.deliver
 ```
+
+## Active Support Notification Instrumentation for AWS SDK calls
+To add `ActiveSupport::Notifications` Instrumentation to all AWS SDK client 
+operations call `Aws::Rails.instrument_sdk_operations` before you construct any
+SDK clients.
+
+Example usage in `config/initializers/instrument_aws_sdk.rb`
+```ruby
+Aws::Rails.instrument_sdk_operations
+```
+
+Events are published for each client operation call with the following event
+name: <operation>.<serviceId>.aws.  For example, S3's put_object has an event
+name of: `put_object.S3.aws`.  The payload of the event is the 
+[request context](https://docs.aws.amazon.com/sdk-for-ruby/v3/api/Seahorse/Client/RequestContext.html).
+
+You can subscribe to these events as you would other
+ `ActiveSupport::Notifications`:
+ 
+ ```ruby
+ActiveSupport::Notifications.subscribe('put_object.s3.aws') do |name, start, finish, id, payload|
+  # process event
+end
+
+# Or use a regex to subscribe to all service notifications
+ActiveSupport::Notifications.subscribe(/s3[.]aws/) do |name, start, finish, id, payload|
+  # process event
+end
+```

lib/aws-sdk-rails.rb

@@ -1,10 +1,13 @@
 # frozen_string_literal: true
 
 require_relative 'aws/rails/mailer'
+require_relative 'aws/rails/notifications_instrument_plugin'
 
 module Aws
+
   # Use the Rails namespace.
   module Rails
+
     # @api private
     class Railtie < ::Rails::Railtie
       initializer 'aws-sdk-rails.initialize',
@@ -47,5 +50,20 @@ def self.use_rails_encrypted_credentials
           .to_h.slice(*aws_credential_keys)
       )
     end
+
+    # Adds ActiveSupport Notifications instrumentation to AWS SDK
+    # client operations.  Each operation will produce an event with a name:
+    # <operation>.<service>.aws.  For example, S3's put_object has an event
+    # name of: put_object.S3.aws
+    def self.instrument_sdk_operations
+      Aws.constants.each do|c|
+        m = Aws.const_get(c)
+        if m.is_a?(Module) &&
+           m.const_defined?(:Client) &&
+           m.const_get(:Client).superclass == Seahorse::Client::Base
+          m.const_get(:Client).add_plugin(Aws::Rails::Notifications)
+        end
+      end
+    end
   end
 end

lib/aws/rails/notifications_instrument_plugin.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require 'aws-sdk-core'
+require 'active_support/notifications'
+
+module Aws
+  module Rails
+
+    # Instruments client operation calls for ActiveSupport::Notifications
+    # Each client operation will produce an event with name:
+    # <operation>.<service>.aws
+    # @api private
+    class Notifications < Seahorse::Client::Plugin
+
+      def add_handlers(handlers, config)
+        # This plugin needs to be first
+        # which means it is called first in the stack, to start recording time,
+        # and returns last
+        handlers.add(Handler, step: :initialize, priority: 99)
+      end
+
+      class Handler < Seahorse::Client::Handler
+
+        def call(context)
+          event_name = "#{context.operation_name}.#{context.config.api.metadata['serviceId']}.aws"
+          ActiveSupport::Notifications.instrument(event_name, context: context) do
+            @handler.call(context)
+          end
+        end
+      end
+    end
+  end
+end

spec/aws/rails/notifications_spec.rb

@@ -0,0 +1,49 @@
+require_relative '../../spec_helper'
+
+module Aws
+
+  # Test namespaces
+  module Service1; end
+  module Service2; end
+
+  module Rails
+    describe 'NotificationsInstrument' do
+
+      let(:base_client) { Aws::SES::Client }
+
+      describe '.instrument_sdk_operations' do
+        it 'adds the plugin to each AWS Client' do
+          Aws::Service1::Client = base_client.dup
+          Aws::Service2::Client = base_client.dup
+
+          expect(Aws::Service1::Client).to receive(:add_plugin).with(Aws::Rails::Notifications)
+          expect(Aws::Service2::Client).to receive(:add_plugin).with(Aws::Rails::Notifications)
+
+          # Ensure other Clients don't get plugin added
+          allow_any_instance_of(Class).to receive(:add_plugin)
+
+          Aws::Rails.instrument_sdk_operations
+        end
+      end
+
+      describe 'NotificationsInstrument Plugin' do
+        let(:client) do
+          Client = base_client.dup
+          Client.add_plugin(Aws::Rails::Notifications)
+          Client.new(stub_responses: true, logger: nil)
+        end
+
+        it 'adds instrumentation on each call' do
+          out = {}
+          ActiveSupport::Notifications.subscribe(/aws/) do |name, start, finish, id, payload|
+            out[:name] = name
+            out[:payload] = payload
+          end
+          client.get_send_quota
+          expect(out[:name]).to eq('get_send_quota.SES.aws')
+          expect(out[:payload][:context]).to be_a(Seahorse::Client::RequestContext)
+        end
+      end
+    end
+  end
+end