experimental(contrib): add aws_lambda_ric integration

.rubocop_todo.yml

@@ -64,6 +64,7 @@ Naming/FileName:
     - 'lib/datadog/profiling/preload.rb'
     - 'lib/datadog/transport/ext.rb'
     - 'lib/datadog/version.rb'
+    - 'lib/datadog/lambda.rb'
 
 # Offense count: 1
 RSpec/IteratedExpectation:

lib/datadog/lambda.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module Datadog
+  # Lambda module streamlines instrumenting AWS Lambda functions with Datadog.
+  module Lambda
+    def self.configure_apm
+      # Load tracing
+      require_relative 'tracing'
+      require_relative 'tracing/contrib'
+
+      # Load other products (must follow tracing)
+      require_relative 'profiling'
+      require_relative 'appsec'
+      require_relative 'di'
+      require_relative 'error_tracking'
+      require_relative 'kit'
+
+      require_relative 'tracing/transport/io'
+
+      # Needed to keep trace flushes on a single line
+      $stdout.sync = true
+
+      Datadog.configure do |c|
+        # unless Datadog::Utils.extension_running?
+        #   c.tracing.writer = Datadog::Tracing::SyncWriter.new(
+        #     transport: Datadog::Tracing::Transport::IO.default
+        #   )
+        # end
+        c.tags = { "_dd.origin": 'lambda' }
+        # Enable AWS SDK instrumentation
+        # c.tracing.instrument :aws if trace_managed_services?
+
+        yield(c) if block_given?
+      end
+    end
+
+    def self.trace_managed_services?
+      dd_trace_managed_services = ENV['DD_TRACE_MANAGED_SERVICES']
+      return true if dd_trace_managed_services.nil?
+
+      dd_trace_managed_services.downcase == 'true'
+    end
+  end
+end

lib/datadog/tracing/contrib.rb

@@ -39,6 +39,7 @@ module Contrib
 require_relative 'contrib/active_record/integration'
 require_relative 'contrib/active_support/integration'
 require_relative 'contrib/aws/integration'
+require_relative 'contrib/aws_lambda_ric/integration'
 require_relative 'contrib/concurrent_ruby/integration'
 require_relative 'contrib/dalli/integration'
 require_relative 'contrib/delayed_job/integration'

lib/datadog/tracing/contrib/aws_lambda_ric/configuration/settings.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require_relative '../../configuration/settings'
+require_relative '../ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module AwsLambdaRic
+        module Configuration
+          # Custom settings for the AWS Lambda RIC integration
+          class Settings < Contrib::Configuration::Settings
+            option :enabled, default: true
+            option :analytics_enabled, default: false
+            option :analytics_sample_rate, default: 1.0
+            option :service_name, default: Ext::TAG_DEFAULT_AGENT
+            option :peer_service, default: nil
+            option :on_error, default: nil
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws_lambda_ric/ext.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module AwsLambdaRic
+        # AWS integration constants
+        # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
+        module Ext
+          ENV_ENABLED = 'DD_TRACE_AWS_LAMBDA_RIC_ENABLED'
+          ENV_SERVICE_NAME = 'DD_TRACE_AWS_LAMBDA_RIC_SERVICE_NAME'
+          ENV_PEER_SERVICE = 'DD_TRACE_AWS_LAMBDA_RIC_PEER_SERVICE'
+          # @!visibility private
+          ENV_ANALYTICS_ENABLED = 'DD_TRACE_AWS_LAMBDA_RIC_ANALYTICS_ENABLED'
+          ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_AWS_LAMBDA_RIC_ANALYTICS_SAMPLE_RATE'
+          # DEFAULT_PEER_SERVICE_NAME = 'aws'
+          SPAN_COMMAND = 'aws.lambda'
+          # TAG_AGENT = 'aws.agent'
+          # TAG_COMPONENT = 'aws'
+          TAG_DEFAULT_AGENT = 'aws-sdk-ruby'
+          # TAG_HOST = 'host'
+          TAG_OPERATION = 'aws.lambda'
+          # TAG_OPERATION_COMMAND = 'command'
+          # TAG_PATH = 'path'
+          TAG_AWS_REGION = 'aws.region'
+          TAG_REGION = 'region'
+          TAG_AWS_ACCOUNT = 'aws_account'
+          TAG_FUNCTION_NAME = 'functionname'
+          PEER_SERVICE_SOURCES = Array[TAG_FUNCTION_NAME,
+            Tracing::Metadata::Ext::TAG_PEER_HOSTNAME,
+            Tracing::Metadata::Ext::NET::TAG_DESTINATION_NAME,
+            Tracing::Metadata::Ext::NET::TAG_TARGET_HOST,].freeze
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws_lambda_ric/instrumentation.rb

@@ -0,0 +1,135 @@
+# frozen_string_literal: true
+
+require_relative '../../metadata/ext'
+require_relative '../analytics'
+require_relative '../ext'
+require_relative 'ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module AwsLambdaRic
+        # AWS Lambda RIC instrumentation module
+        module Instrumentation
+          def self.included(base)
+            base.prepend(InstanceMethods)
+          end
+
+          # Instance methods for AWS Lambda RIC classes
+          module InstanceMethods
+            def run_user_code
+              return super unless enabled?
+
+              # Store the arguments for later use in the trace
+              @lambda_user_code_args = method(__method__).super_method.parameters.map do |type, name|
+                [type, name, binding.local_variable_get(name)] if binding.local_variable_defined?(name)
+              end.compact
+
+              super
+            end
+
+            def call_handler(request:, context:)
+              return super unless enabled?
+
+              # Store the arguments for later use in the trace
+              @lambda_call_handler_args = { request: request, context: context }
+
+              # Start the trace when the handler is called
+              start_lambda_trace
+              result = super
+              # Store the result for later use in the trace
+              @lambda_handler_result = result
+              # Finish the trace when the handler completes
+              finish_lambda_trace
+              result
+            end
+
+            private
+
+            def start_lambda_trace
+              service = datadog_configuration[:service_name]
+              on_error = datadog_configuration[:on_error]
+
+              # Start a new span manually
+              @lambda_span = Tracing.start_span(
+                Ext::SPAN_COMMAND,
+                service: service,
+                resource: 'lambda_execution',
+                type: Tracing::Metadata::Ext::HTTP::TYPE_INBOUND
+              )
+
+              annotate_span!(@lambda_span, 'lambda_execution')
+
+              # Set analytics sample rate
+              Contrib::Analytics.set_sample_rate(@lambda_span, analytics_sample_rate) if analytics_enabled?
+
+              # Set error handler
+              @lambda_span.on_error = on_error if on_error
+            end
+
+            def finish_lambda_trace
+              return unless @lambda_span
+
+              # Add any additional tags or metadata before finishing
+              @lambda_span.set_tag('lambda.user_code_args', @lambda_user_code_args.to_json) if @lambda_user_code_args
+
+              if @lambda_call_handler_args
+                @lambda_span.set_tag('lambda.call_handler.request', @lambda_call_handler_args[:request].to_json)
+                @lambda_span.set_tag('lambda.call_handler.context', @lambda_call_handler_args[:context].to_json)
+              end
+
+              if @lambda_handler_result
+                handler_response, content_type = @lambda_handler_result
+                @lambda_span.set_tag('lambda.handler_response', handler_response.to_json)
+                @lambda_span.set_tag('lambda.content_type', content_type.to_json) if content_type
+              end
+
+              # Finish the span
+              @lambda_span.finish
+
+              # Clean up instance variables
+              @lambda_span = nil
+              @lambda_user_code_args = nil
+              @lambda_call_handler_args = nil
+              @lambda_handler_result = nil
+            end
+
+            def annotate_span!(span, operation)
+              span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_DEFAULT_AGENT)
+              span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION)
+              span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_SERVER)
+
+              # Add AWS Lambda specific tags
+              span.set_tag(Ext::TAG_FUNCTION_NAME, function_name) if respond_to?(:function_name)
+              span.set_tag(Ext::TAG_AWS_REGION, aws_region) if respond_to?(:aws_region)
+              span.set_tag(Ext::TAG_AWS_ACCOUNT, aws_account) if respond_to?(:aws_account)
+
+              # Tag original global service name if not used
+              if span.service != Datadog.configuration.service
+                span.set_tag(Tracing::Contrib::Ext::Metadata::TAG_BASE_SERVICE, Datadog.configuration.service)
+              end
+
+              Contrib::SpanAttributeSchema.set_peer_service!(span, Ext::PEER_SERVICE_SOURCES)
+            end
+
+            def datadog_configuration
+              Datadog.configuration.tracing[:aws_lambda_ric]
+            end
+
+            def analytics_enabled?
+              datadog_configuration[:analytics_enabled]
+            end
+
+            def analytics_sample_rate
+              datadog_configuration[:analytics_sample_rate]
+            end
+
+            def enabled?
+              datadog_configuration[:enabled]
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws_lambda_ric/integration.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require_relative '../integration'
+require_relative 'configuration/settings'
+require_relative 'patcher'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module AwsLambdaRic
+        # Description of AWS Lambda RIC integration
+        class Integration
+          include Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new('3.0')
+
+          # @public_api Changing the integration name or integration options can cause breaking changes
+          register_as :aws_lambda_ric, auto_patch: true
+          def self.gem_name
+            'aws_lambda_ric'
+          end
+
+          def self.version
+            Gem.loaded_specs['aws_lambda_ric'].version if Gem.loaded_specs['aws_lambda_ric']
+          end
+
+          def self.loaded?
+            !defined?(::AwsLambdaRIC).nil?
+          end
+
+          def self.compatible?
+            super && version >= MINIMUM_VERSION
+          end
+
+          def new_configuration
+            Configuration::Settings.new
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws_lambda_ric/patcher.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require_relative '../patcher'
+require_relative 'instrumentation'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module AwsLambdaRic
+        # Patcher enables patching of 'aws_lambda_ric' module.
+        module Patcher
+          include Contrib::Patcher
+
+          module_function
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            patch_lambda_runner
+            patch_lambda_handler
+          end
+
+          def patch_lambda_runner
+            return unless defined?(::AwsLambdaRIC::LambdaRunner)
+
+            ::AwsLambdaRIC::LambdaRunner.include(Instrumentation)
+          rescue StandardError => e
+            Datadog.logger.debug("Unable to patch AwsLambdaRIC::LambdaRunner: #{e}")
+          end
+
+          def patch_lambda_handler
+            return unless defined?(::AwsLambdaRIC::LambdaHandler)
+
+            ::AwsLambdaRIC::LambdaHandler.include(Instrumentation)
+          rescue StandardError => e
+            Datadog.logger.debug("Unable to patch AwsLambdaRIC::LambdaHandler: #{e}")
+          end
+        end
+      end
+    end
+  end
+end

sig/datadog/tracing/contrib/aws_lambda_ric/configuration/settings.rbs

@@ -0,0 +1,12 @@
+module Datadog
+  module Tracing
+    module Contrib
+      module AwsLambdaRic
+        module Configuration
+          class Settings < Contrib::Configuration::Settings
+          end
+        end
+      end
+    end
+  end
+end

sig/datadog/tracing/contrib/aws_lambda_ric/ext.rbs

@@ -0,0 +1,30 @@
+module Datadog
+  module Tracing
+    module Contrib
+      module AwsLambdaRic
+        module Ext
+          ENV_ENABLED: "DD_TRACE_AWS_LAMBDA_RIC_ENABLED"
+
+          ENV_SERVICE_NAME: "DD_TRACE_AWS_LAMBDA_RIC_SERVICE_NAME"
+
+          ENV_PEER_SERVICE: "DD_TRACE_AWS_LAMBDA_RIC_PEER_SERVICE"
+          ENV_ANALYTICS_ENABLED: "DD_TRACE_AWS_LAMBDA_RIC_ANALYTICS_ENABLED"
+
+          ENV_ANALYTICS_SAMPLE_RATE: "DD_TRACE_AWS_LAMBDA_RIC_ANALYTICS_SAMPLE_RATE"
+          SPAN_COMMAND: "aws.lambda"
+          TAG_DEFAULT_AGENT: "aws-sdk-ruby"
+          TAG_OPERATION: "aws.lambda"
+          TAG_AWS_REGION: "aws.region"
+
+          TAG_REGION: "region"
+
+          TAG_AWS_ACCOUNT: "aws_account"
+
+          TAG_FUNCTION_NAME: "functionname"
+
+          PEER_SERVICE_SOURCES: untyped
+        end
+      end
+    end
+  end
+end