Lambda instrumentation implementation

README.md

@@ -84,6 +84,17 @@ user_config = {
 XRay.recorder.configure(user_config)
 ```
 
+**Working with Lambda**
+
+```ruby
+# Bundle aws-xray-sdk with your code.
+# Require any aws-sdks and http libraries to be used, then...
+require 'aws-xray-sdk/lambda'
+# aws-sdk and http calls from here on will be instrumented
+```
+See also [lib/aws-xray-sdk/lambda.rb](lib/aws-xray-sdk/lambda.rb)
+
+
 **Working with Rails**
 
 ```ruby

lib/aws-xray-sdk/exceptions.rb

@@ -28,4 +28,7 @@ class InvalidConfigurationError < AwsXRaySdkError
 
   class UnsupportedPatchingTargetError < AwsXRaySdkError
   end
+
+  class UnsupportedOperationError < AwsXRaySdkError
+  end
 end

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

@@ -20,8 +20,21 @@ def initialize(*options)
         super(*options)
       end
 
+      # HTTP requests to AWS Lambda Ruby Runtime will begin with the
+      # value set in ENV['AWS_LAMBDA_RUNTIME_API']
+      def lambda_runtime_request?(req)
+        ENV['AWS_LAMBDA_RUNTIME_API']  &&
+          req.uri &&
+          req.uri.to_s.start_with?('http://'+ENV['AWS_LAMBDA_RUNTIME_API']+'/')
+      end
+
+      def xray_sampling_request?(req)
+        req.path && (req.path == ('/GetSamplingRules') || req.path == ('/SamplingTargets'))
+      end
+
       def request(req, body = nil, &block)
-        if req.path && (req.path == ('/GetSamplingRules') || req.path == ('/SamplingTargets'))
+        # Do not trace requests to xray or aws lambda runtime
+        if xray_sampling_request?(req) || lambda_runtime_request?(req)
           return super
         end
 

lib/aws-xray-sdk/lambda.rb

@@ -0,0 +1,35 @@
+# AWS Lambda functions should require this file to configure the XRay.recorder
+# for use within the function.
+
+require_relative 'lambda/lambda_recorder'
+require_relative 'lambda/facade_segment'
+require_relative 'lambda/lambda_context'
+require_relative 'lambda/lambda_streamer'
+
+module XRay
+  @recorder = LambdaRecorder.new
+
+  # provide an instance of LambdaRecorder as the global XRay.recorder
+  def self.recorder
+    @recorder
+  end
+end
+
+
+# Set `XRAY_LAMBDA_PATCH_CONFIG` before requiring `aws-xray-sdk/lambda`
+# to configure which libraries (if any) to instrument.
+#
+# By default, both `net_http` and `aws_sdk` will be instrumented
+unless defined? XRAY_LAMBDA_PATCH_CONFIG
+  XRAY_LAMBDA_PATCH_CONFIG = %I[net_http aws_sdk]
+end
+
+# Configure the XRay.recorder with Lambda specific config.
+#
+# From here, a lambda may create subsegments manually, or via
+# the instrumented libraries setup by XRAY_LAMBDA_PATCH_CONFIG
+XRay.recorder.configure(
+  patch: XRAY_LAMBDA_PATCH_CONFIG,
+  context: XRay::LambdaContext.new,
+  streamer: XRay::LambdaStreamer.new
+)

lib/aws-xray-sdk/lambda/facade_segment.rb

@@ -0,0 +1,109 @@
+module XRay
+  class FacadeSegment < XRay::Segment
+
+    class ImmutableEmptyCollection
+      def [](key)
+        nil
+      end
+
+      def []=(k, v)
+        raise UnsupportedOperationError
+      end
+
+      def update(h)
+        raise UnsupportedOperationError
+      end
+
+      def to_h
+        {}
+      end
+    end
+
+
+    def initialize(trace_id: nil, name: nil, parent_id: nil, id: nil, sampled: true)
+      super(trace_id: trace_id, name: name, parent_id: parent_id)
+      @id = id
+      @sampled = sampled
+      @empty_collection = ImmutableEmptyCollection.new
+    end
+
+    def ready_to_send?
+      false #never send this facade. AWS Lambda has already created a Segment with these ids
+    end
+
+    #
+    #Methods from Entity that are not supported
+    #
+    def close(end_time: nil)
+      raise UnsupportedOperationError
+    end
+    def apply_status_code(status:)
+      raise UnsupportedOperationError
+    end
+    def merge_http_request(request:)
+      raise UnsupportedOperationError
+    end
+    def merge_http_response(response:)
+      raise UnsupportedOperationError
+    end
+    def add_exception(exception:, remote: false)
+      raise UnsupportedOperationError
+    end
+
+    #
+    # Mutation accessors from Entity that are not supported
+    #
+    def parent=(value)
+      raise UnsupportedOperationError
+    end
+    def throttle=(value)
+      raise UnsupportedOperationError
+    end
+    def error=(value)
+      raise UnsupportedOperationError
+    end
+    def fault=(value)
+      raise UnsupportedOperationError
+    end
+    def sampled=(value)
+      raise UnsupportedOperationError
+    end
+    def aws=(value)
+      raise UnsupportedOperationError
+    end
+    def start_time=(value)
+      raise UnsupportedOperationError
+    end
+    def end_time=(value)
+      raise UnsupportedOperationError
+    end
+
+    #
+    # Mutation accessors from Segment that are not supported
+    #
+    def origin=(value)
+      raise UnsupportedOperationError
+    end
+    def user=(value)
+      raise UnsupportedOperationError
+    end
+    def service=(value)
+      raise UnsupportedOperationError
+    end
+
+    #
+    # Annotations are read only
+    #
+    def annotations
+      @empty_collection
+    end
+
+    #
+    # Metadata is read only
+    #
+    def metadata(namespace: :default)
+      @empty_collection
+    end
+
+  end
+end

lib/aws-xray-sdk/lambda/lambda_context.rb

@@ -0,0 +1,37 @@
+require_relative '../model/trace_header'
+
+module XRay
+  # LambdaContext extends the default context so that
+  # we can provide an appropriate FacadeSegment as the
+  # root context for each function invocation.
+  class LambdaContext < XRay::DefaultContext
+
+    TRACE_ID_ENV_VAR = '_X_AMZN_TRACE_ID'.freeze
+
+    def lambda_trace_id
+      ENV[TRACE_ID_ENV_VAR]
+    end
+
+    # If the environment trace id changes, create a new facade for that
+    # segment and make it the context's current entity
+    def check_context
+      #Create a new FacadeSegment if the _X_AMZN_TRACE_ID changes.
+      return if lambda_trace_id == @current_trace_id
+
+      @current_trace_id = lambda_trace_id
+      trace_header = XRay::TraceHeader.from_header_string(header_str: @current_trace_id)
+      segment = FacadeSegment.new(trace_id: trace_header.root,
+        parent_id: trace_header.parent_id,
+        id: trace_header.parent_id,
+        name: 'lambda_context',
+        sampled: trace_header.sampled == 1
+      )
+      store_entity(entity: segment)
+    end
+
+    def current_entity
+      check_context #ensure the FacadeSegment is current whenever the current_entity is retrieved
+      super
+    end
+  end
+end

lib/aws-xray-sdk/lambda/lambda_recorder.rb

@@ -0,0 +1,19 @@
+require 'aws-xray-sdk/recorder'
+
+module XRay
+  class LambdaRecorder < Recorder
+    include Logging
+
+    def begin_segment(name, trace_id: nil, parent_id: nil, sampled: nil)
+      # no-op
+      logger.warn('Cannot create segments inside Lambda function. Returning current segment.')
+      return current_segment
+    end
+
+    def end_segment(end_time: nil)
+      # no-op
+      logger.warn('Cannot end segment inside Lambda function. Ignored.')
+      nil
+    end
+  end
+end

lib/aws-xray-sdk/lambda/lambda_streamer.rb

@@ -0,0 +1,9 @@
+module XRay
+  # LambdaStreamer extends DefaultStreamer so that subsegments 
+  # are sent to the XRay endpoint as they are available.
+  class LambdaStreamer < XRay::DefaultStreamer
+    def initialize
+      @stream_threshold = 1 #Stream every subsegment as it is available
+    end
+  end
+end

test/aws-xray-sdk/tc_facet_rack.rb

@@ -34,27 +34,20 @@ class TestFacetRack < Minitest::Test
     "REQUEST_PATH" => "/index.html"
   }.freeze
 
-  @@recorder = XRay::Recorder.new
-  config = {
-    name: "rack_test",
-    emitter: XRay::TestHelper::StubbedEmitter.new,
-    sampler: XRay::TestHelper::StubbedDefaultSampler.new
-  }
-  @@recorder.configure(config)
 
   @@app = Minitest::Mock.new
   def @@app.call(env)
     return 200, {}, ""
   end
 
   def setup
-    @@recorder.context.clear!
-    @@recorder.emitter.clear
-  end
-
-  def teardown
-    @@recorder.context.clear!
-    @@recorder.emitter.clear
+    @@recorder = XRay::Recorder.new
+    config = {
+      name: "rack_test",
+      emitter: XRay::TestHelper::StubbedEmitter.new,
+      sampler: XRay::TestHelper::StubbedDefaultSampler.new
+    }
+    @@recorder.configure(config)
   end
 
   def test_rack_http_url_excludes_query_string