aws-xray-sdk 0.9.0

data/lib/aws-xray-sdk/facets/helper.rb

@@ -0,0 +1,61 @@
+require 'aws-xray-sdk/model/trace_header'
+
+module XRay
+  module Facets
+    # Hepler functions shared for all external frameworks/libraries
+    # like make sampling decisions from incoming http requests etc.
+    module Helper
+      TRACE_HEADER = 'X-Amzn-Trace-Id'.freeze
+      TRACE_HEADER_PROXY = 'HTTP_X_AMZN_TRACE_ID'.freeze
+
+      # Construct a `TraceHeader` object from headers
+      # of the incoming request. This method should always return
+      # a `TraceHeader` object regardless of tracing header's presence
+      # in the incoming request.
+      # @param [Hash] headers Hash that contains X-Ray trace header key.
+      # @return [TraceHeader] The new constructed trace header object.
+      def construct_header(headers:)
+        if v = headers[TRACE_HEADER_PROXY] || headers[TRACE_HEADER]
+          TraceHeader.from_header_string header_str: v
+        else
+          TraceHeader.empty_header
+        end
+      end
+
+      # The sampling decision coming from `trace_header` always has
+      # the highest precedence. If the `trace_header` doesn't contain
+      # sampling decision then it checks if sampling is enabled or not
+      # in the recorder. If not enbaled it returns 'true'. Otherwise it uses
+      # sampling rule to decide.
+      def should_sample?(header_obj:, recorder:,
+                         host: nil, method: nil, path: nil,
+                         **args)
+        # check outside decision
+        if i = header_obj.sampled
+          if i.zero?
+            false
+          else
+            true
+          end
+        # check sampling rules
+        elsif recorder.sampling_enabled?
+          recorder.sampler.sample_request?(service_name: host, url_path: path,
+                                           http_method: method)
+        # sample if no reason not to
+        else
+          true
+        end
+      end
+
+      # Prepares a X-Ray header string based on the provided Segment/Subsegment.
+      def prep_header_str(entity:)
+        return '' if entity.nil?
+        root = entity.segment.trace_id
+        parent_id = entity.id
+        sampled = entity.sampled ? 1 : 0
+        header = TraceHeader.new root: root, parent_id: parent_id, sampled: sampled
+        header.header_string
+      end
+    end
+  end
+end

data/lib/aws-xray-sdk/facets/net_http.rb

@@ -0,0 +1,61 @@
+require 'net/http'
+require 'aws-xray-sdk/facets/helper'
+
+module XRay
+  # Patch net/http to be traced by X-Ray
+  module NetHttp
+    # Class level interceptor to capture http requests as subsegments
+    module HTTPClassInterceptor
+      def new(*options)
+        o = super(*options)
+        o
+      end
+    end
+
+    # Instance level interceptor to capture http requests as subsegments
+    module HTTPInstanceInterceptor
+      include XRay::Facets::Helper
+
+      def initialize(*options)
+        super(*options)
+      end
+
+      def request(req, body = nil, &block)
+        entity = XRay.recorder.current_entity
+        capture = !(entity && entity.namespace && entity.namespace == 'aws'.freeze)
+        if started? && capture && entity
+          XRay.recorder.capture(address, namespace: 'remote') do |subsegment|
+            protocol = use_ssl? ? 'https'.freeze : 'http'.freeze
+            # avoid modifying original variable
+            iport = port.nil? ? nil : %(:#{port})
+            # do not capture query string
+            path = req.path.split('?')[0] if req.path
+            uri = %(#{protocol}://#{address}#{iport}#{path})
+            req_meta = {
+              url: uri,
+              method: req.method
+            }
+            subsegment.merge_http_request request: req_meta
+            req[TRACE_HEADER] = prep_header_str entity: subsegment
+            begin
+              res = super
+              res_meta = {
+                status: res.code.to_i,
+                content_length: res.content_length
+              }
+              subsegment.merge_http_response response: res_meta
+            rescue Exception => e
+              subsegment.add_exception exception: e
+              raise e
+            end
+          end
+        else
+          super
+        end
+      end
+    end
+
+    ::Net::HTTP.singleton_class.prepend HTTPClassInterceptor
+    ::Net::HTTP.prepend HTTPInstanceInterceptor
+  end
+end

data/lib/aws-xray-sdk/facets/rack.rb

@@ -0,0 +1,87 @@
+require 'rack/request'
+require 'aws-xray-sdk'
+require 'aws-xray-sdk/facets/helper'
+
+module XRay
+  module Rack
+    # Rack middleware that generates a segment for each request/response cycle.
+    class Middleware
+      include XRay::Facets::Helper
+      X_FORWARD = 'HTTP_X_FORWARDED_FOR'.freeze
+
+      def initialize(app, recorder: nil)
+        @app = app
+        @recorder = recorder || XRay.recorder
+      end
+
+      def call(env)
+        header = construct_header(headers: env)
+        req = ::Rack::Request.new(env)
+
+        # params required for path based sampling
+        host = req.host
+        url_path = req.path
+        method = req.request_method
+
+        # get sampling decision
+        sampled = should_sample?(
+          header_obj: header, recorder: @recorder,
+          host: host, method: method, path: url_path
+        )
+
+        # get segment name from host header if applicable
+        seg_name = @recorder.segment_naming.provide_name(host: req.host)
+
+        # begin the segment
+        segment = @recorder.begin_segment seg_name, trace_id: header.root, parent_id: header.parent_id,
+                                                    sampled: sampled
+
+        # add neccessary http request metadata to the segment
+        req_meta = extract_request_meta(req)
+        segment.merge_http_request request: req_meta unless req_meta.empty?
+        begin
+          status, headers, body = @app.call env
+          resp_meta = {}
+          resp_meta[:status] = status
+          # Don't set content_length if it is not available on headers.
+          resp_obj = ::Rack::Response.new body: body, status: status, headers: headers
+          if len = resp_obj.content_length
+            resp_meta[:content_length] = len
+          end
+          segment.merge_http_response response: resp_meta
+          [status, headers, body]
+        rescue Exception => e
+          segment.apply_status_code status: 500
+          segment.add_exception exception: e
+          raise e
+        ensure
+          @recorder.end_segment
+        end
+      end
+
+      private
+
+      def extract_request_meta(req)
+        req_meta = {}
+        req_meta[:url] = req.url if req.url
+        req_meta[:user_agent] = req.user_agent if req.user_agent
+        req_meta[:method] = req.request_method if req.request_method
+        if req.has_header?(X_FORWARD)
+          req_meta[:client_ip] = get_ip(req.get_header(X_FORWARD))
+          req_meta[:x_forwarded_for] = true
+        elsif v = req.ip
+          req_meta[:client_ip] = v
+        end
+        req_meta
+      end
+
+      def get_ip(ips)
+        if ips.respond_to?(:length)
+          ips[ips.length - 1]
+        else
+          ips
+        end
+      end
+    end
+  end
+end

data/lib/aws-xray-sdk/facets/rails/active_record.rb

@@ -0,0 +1,66 @@
+require 'active_record'
+
+module XRay
+  module Rails
+    # Recording Rails database transactions as subsegments.
+    module ActiveRecord
+      class << self
+        IGNORE_OPS = ['SCHEMA', 'ActiveRecord::SchemaMigration Load',
+                      'ActiveRecord::InternalMetadata Load'].freeze
+        DB_TYPE_MAPPING = {
+          mysql2: 'MySQL',
+          postgresql: 'PostgreSQL'
+        }.freeze
+
+        def record(transaction)
+          payload = transaction.payload
+          pool, conn = get_pool_n_conn(payload[:connection_id])
+
+          return if IGNORE_OPS.include?(payload[:name]) || pool.nil? || conn.nil?
+          db_config = pool.spec.config
+          name, sql = build_name_sql_meta config: db_config, conn: conn
+          subsegment = XRay.recorder.begin_subsegment name, namespace: 'remote'
+          subsegment.start_time = transaction.time.to_f
+          subsegment.sql = sql
+          XRay.recorder.end_subsegment end_time: transaction.end.to_f
+        end
+
+        private
+
+        def build_name_sql_meta(config:, conn:)
+          # extract all available info
+          adapter = config[:adapter]
+          database = config[:database]
+          host = config[:host].nil? ? nil : %(@#{config[:host]})
+          port = config[:port].nil? ? nil : %(:#{config[:port]})
+          username = config[:username]
+
+          # assemble subsegment name
+          name = %(#{database}#{host})
+          # assemble sql meta
+          sql = {}
+          sql[:user] = username
+          sql[:url] = %(#{adapter}://#{username}#{host}#{port}/#{database})
+          sql[:database_type] = DB_TYPE_MAPPING[adapter.to_sym]
+          [name, sql]
+        end
+
+        def get_pool_n_conn(conn_id)
+          pool, conn = nil, nil
+          ::ActiveRecord::Base.connection_handler.connection_pool_list.each do |p|
+            conn = p.connections.select { |c| c.object_id == conn_id }
+            pool = p unless conn.nil?
+          end
+          [pool, conn]
+        end
+      end
+    end
+  end
+end
+
+# Add a hook on database transactions using Rails instrumentation API
+ActiveSupport::Notifications.subscribe('sql.active_record') do |*args|
+  # We need the full event which has all the timing info
+  transaction = ActiveSupport::Notifications::Event.new(*args)
+  XRay::Rails::ActiveRecord.record(transaction)
+end

data/lib/aws-xray-sdk/facets/rails/ex_middleware.rb

@@ -0,0 +1,24 @@
+require 'aws-xray-sdk'
+
+module XRay
+  module Rails
+    # Middleware for capturing unhandled exceptions from views/controller.
+    # To properly capture exceptions this middleware needs to be placed
+    # after the default exception handling middleware. Otherwise they will
+    # be swallowed.
+    class ExceptionMiddleware
+      def initialize(app, recorder: nil)
+        @app = app
+        @recorder = recorder || XRay.recorder
+      end
+
+      def call(env)
+        @app.call(env)
+      rescue Exception => e
+        segment = @recorder.current_segment
+        segment.add_exception exception: e if segment
+        raise e
+      end
+    end
+  end
+end

data/lib/aws-xray-sdk/facets/rails/railtie.rb

@@ -0,0 +1,23 @@
+require 'aws-xray-sdk/facets/rack'
+require 'aws-xray-sdk/facets/rails/ex_middleware'
+
+module XRay
+  # configure X-Ray instrumentation for rails framework
+  class Railtie < ::Rails::Railtie
+    RAILS_OPTIONS = %I[active_record].freeze
+
+    initializer "aws-xray-sdk.rack_middleware" do |app|
+      app.middleware.insert 0, Rack::Middleware
+      app.middleware.use XRay::Rails::ExceptionMiddleware
+    end
+
+    config.after_initialize do |app|
+      if app.config.respond_to?('xray')
+        options = app.config.xray
+        require 'aws-xray-sdk/facets/rails/active_record' if options[:active_record]
+        general_options = options.reject { |k, v| RAILS_OPTIONS.include?(k) }
+        XRay.recorder.configure(general_options)
+      end
+    end
+  end
+end

data/lib/aws-xray-sdk/facets/resources/aws_params_whitelist.rb

@@ -0,0 +1,340 @@
+module XRay
+  # AWS SDK parameters whitelisted will be recorded
+  # as metadata on AWS subsegments
+  module AwsParams
+    @whitelist = {
+      services: {
+        DynamoDB: {
+          operations: {
+            batch_get_item: {
+              request_descriptors: {
+                request_items: {
+                  map: true,
+                  get_keys: true,
+                  rename_to: :table_names
+                }
+              },
+              response_parameters: %I[
+                consumed_capacity
+              ]
+            },
+            batch_write_item: {
+              request_descriptors: {
+                request_items: {
+                  map: true,
+                  get_keys: true,
+                  rename_to: :table_names
+                }
+              },
+              response_parameters: %I[
+                consumed_capacity
+                item_collection_metrics
+              ]
+            },
+            create_table: {
+              request_parameters: %I[
+                global_secondary_indexes
+                local_secondary_indexes
+                provisioned_throughput
+                table_name
+              ]
+            },
+            delete_item: {
+              request_parameters: %I[
+                table_name
+              ],
+              response_parameters: %I[
+                consumed_capacity
+                item_collection_metrics
+              ]
+            },
+            delete_table: {
+              request_parameters: %I[
+                table_name
+              ]
+            },
+            describe_table: {
+              request_parameters: %I[
+                table_name
+              ]
+            },
+            get_item: {
+              request_parameters: %I[
+                consistent_read
+                projection_expression
+                table_name
+              ],
+              response_parameters: %I[
+                consumed_capacity
+              ]
+            },
+            list_tables: {
+              request_parameters: %I[
+                exclusive_start_table_name
+                limit
+              ],
+              response_descriptors: {
+                table_names: {
+                  list: true,
+                  get_count: true,
+                  rename_to: :table_count
+                }
+              }
+            },
+            put_item: {
+              request_parameters: %I[
+                table_name
+              ],
+              response_parameters: %I[
+                consumed_capacity
+                item_collection_metrics
+              ]
+            },
+            query: {
+              request_parameters: %I[
+                attributes_to_get
+                consistent_read
+                index_name
+                limit
+                projection_expression
+                scan_index_forward
+                select
+                table_name
+              ],
+              response_parameters: %I[
+                consumed_capacity
+              ]
+            },
+            scan: {
+              request_parameters: %I[
+                attributes_to_get
+                consistent_read
+                index_name
+                limit
+                projection_expression
+                segment
+                select
+                table_name
+                total_segments
+              ],
+              response_parameters: %I[
+                consumed_capacity
+                count
+                scanned_count
+              ]
+            },
+            update_item: {
+              request_parameters: %I[
+                table_name
+              ],
+              response_parameters: %I[
+                consumed_capacity
+                item_collection_metrics
+              ]
+            },
+            update_table: {
+              request_parameters: %I[
+                attribute_definitions
+                global_secondary_index_updates
+                provisioned_throughput
+                table_name
+              ]
+            }
+          }
+        },
+        SQS: {
+          operations: {
+            add_permission: {
+              request_parameters: %I[
+                label
+                queue_url
+              ]
+            },
+            change_message_visibility: {
+              request_parameters: %I[
+                queue_url
+                visibility_timeout
+              ]
+            },
+            change_message_visibility_batch: {
+              request_parameters: %I[
+                queue_url
+              ],
+              response_parameters: %I[
+                failed
+              ]
+            },
+            create_queue: {
+              request_parameters: %I[
+                attributes
+                queue_name
+              ]
+            },
+            delete_message: {
+              request_parameters: %I[
+                queue_urls
+              ]
+            },
+            delete_message_batch: {
+              request_parameters: %I[
+                queue_url
+              ],
+              response_parameters: %I[
+                failed
+              ]
+            },
+            delete_queue: {
+              request_parameters: %I[
+                queue_url
+              ]
+            },
+            get_queue_attributes: {
+              request_parameters: %I[
+                queue_url
+              ],
+              response_parameters: %I[
+                attributes
+              ]
+            },
+            get_queue_url: {
+              request_parameters: %I[
+                queue_name
+                queue_owner_aws_account_id
+              ],
+              response_parameters: %I[
+                queue_url
+              ]
+            },
+            list_dead_letter_source_queues: {
+              request_parameters: %I[
+                queue_url
+              ],
+              response_parameters: %I[
+                queue_urls
+              ]
+            },
+            list_queues: {
+              request_parameters: %I[
+                queue_name_prefix
+              ],
+              response_descriptors: {
+                queue_urls: {
+                  list: true,
+                  get_count: true,
+                  rename_to: :queue_count
+                }
+              }
+            },
+            purge_queue: {
+              request_parameters: %I[
+                queue_url
+              ]
+            },
+            receive_message: {
+              request_parameters: %I[
+                attribute_names
+                max_number_of_messages
+                message_attribute_names
+                queue_url
+                visibility_timeout
+                wait_time_seconds
+              ],
+              response_descriptors: {
+                messages: {
+                  list: true,
+                  get_count: true,
+                  rename_to: :message_count
+                }
+              }
+            },
+            remove_permission: {
+              request_parameters: %I[
+                queue_url
+              ]
+            },
+            send_message: {
+              request_parameters: %I[
+                delay_seconds
+                queue_url
+              ],
+              request_descriptors: {
+                message_attributes: {
+                  map: true,
+                  get_keys: true,
+                  rename_to: :message_attribute_names
+                }
+              },
+              response_parameters: %I[
+                message_id
+              ]
+            },
+            send_message_batch: {
+              request_parameters: %I[
+                queue_url
+              ],
+              request_descriptors: {
+                entries: {
+                  list: true,
+                  get_count: true,
+                  rename_to: :message_count
+                }
+              },
+              response_descriptors: {
+                failed: {
+                  list: true,
+                  get_count: true,
+                  rename_to: :failed_count
+                },
+                successful: {
+                  list: true,
+                  get_count: true,
+                  rename_to: :successful_count
+                }
+              }
+            },
+            set_queue_attributes: {
+              request_parameters: %I[
+                queue_url
+              ],
+              request_descriptors: {
+                attributes: {
+                  map: true,
+                  get_keys: true,
+                  rename_to: :attribute_names
+                }
+              }
+            }
+          }
+        },
+        Lambda: {
+          operations: {
+            invoke: {
+              request_parameters: %I[
+                function_name
+                invocation_type
+                log_type
+                qualifier
+              ],
+              response_parameters: %I[
+                function_error
+                status_code
+              ]
+            },
+            invoke_async: {
+              request_parameters: %I[
+                function_name
+              ],
+              response_parameters: %I[
+                status
+              ]
+            }
+          }
+        }
+      }
+    }
+
+    def self.whitelist
+      @whitelist
+    end
+  end
+end