handler.rb in aspecto-opentelemetry-instrumentation-aws_sdk-0.1.8

- old
+ new

@@ -15,11 +15,11 @@
         SNS_PUBLISH = 'SNS.Publish'
 
         def call(context) # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
           return super unless context
 
-          service_name = context.client.class.api.metadata['serviceId'] || context.client.class.to_s.split('::')[1]
+          service_name = context.client.class.api.metadata['serviceId'] || context.client.class.to_s.split('::')[1] rescue context.client.class.to_s.split('::')[1] # rubocop:disable Style/RescueModifier
           operation = context.operation&.name
           client_method = "#{service_name}.#{operation}"
           attributes = {
             'aws.region' => context.config.region,
             OpenTelemetry::SemanticConventions::Trace::RPC_SYSTEM => 'aws-api',
@@ -28,21 +28,25 @@
           }
           attributes[SemanticConventions::Trace::DB_SYSTEM] = 'dynamodb' if service_name == 'DynamoDB'
           MessagingHelper.apply_sqs_attributes(attributes, context, client_method) if service_name == 'SQS'
           MessagingHelper.apply_sns_attributes(attributes, context, client_method) if service_name == 'SNS'
 
+          prepare_extract_context(context, client_method)
           tracer.in_span(span_name(context, client_method), attributes: attributes, kind: span_kind(client_method)) do |span|
             inject_context(context, client_method)
+
             if instrumentation_config[:suppress_internal_instrumentation]
               OpenTelemetry::Common::Utilities.untraced { super }
             else
               super
             end.tap do |response|
               if (err = response.error)
                 span.record_exception(err)
                 span.status = Trace::Status.error(err.to_s)
               end
+
+              MessagingHelper.create_sqs_processing_spans(context, tracer, response.messages) if client_method == SQS_RECEIVE_MESSAGE && response.respond_to?(:messages)
             end
           end
         end
 
         private
@@ -68,10 +72,17 @@
             context.params[:message_attributes] ||= {}
             OpenTelemetry.propagation.inject(context.params[:message_attributes], setter: MessageAttributeSetter)
           end
         end
 
+        def prepare_extract_context(context, client_method)
+          return unless client_method == SQS_RECEIVE_MESSAGE
+          return unless instrumentation_config[:extract_messaging_context]
+
+          context.params[:message_attribute_names] = ['All']
+        end
+
         def span_kind(client_method)
           case client_method
           when SQS_SEND_MESSAGE, SQS_SEND_MESSAGE_BATCH, SNS_PUBLISH
             OpenTelemetry::Trace::SpanKind::PRODUCER
           when SQS_RECEIVE_MESSAGE
@@ -82,12 +93,12 @@
         end
 
         def span_name(context, client_method)
           case client_method
           when SQS_SEND_MESSAGE, SQS_SEND_MESSAGE_BATCH, SNS_PUBLISH
-            "#{MessagingHelper.queue_name(context)} send"
+            "#{MessagingHelper.destination_name(context)} send"
           when SQS_RECEIVE_MESSAGE
-            "#{MessagingHelper.queue_name(context)} receive"
+            "#{MessagingHelper.destination_name(context)} receive"
           else
             client_method
           end
         end
       end