Merge pull request #2730 from DataDog/ar/ConsistentApmSpanTagsForAwsR…

…equests

lib/datadog/core/configuration.rb

@@ -208,7 +208,11 @@ def components(allow_initialization: true)
       # Used internally to ensure a clean environment between test runs.
       def reset!
         safely_synchronize do |write_components|
-          @components.shutdown! if components?
+          if components?
+            @components.shutdown!
+            @temp_logger = nil # Reset to ensure instance and log level are reset for next run
+          end
+
           write_components.call(nil)
           configuration.reset!
         end

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

@@ -18,7 +18,17 @@ module Ext
           TAG_OPERATION = 'aws.operation'.freeze
           TAG_OPERATION_COMMAND = 'command'.freeze
           TAG_PATH = 'path'.freeze
-          TAG_REGION = 'aws.region'.freeze
+          TAG_AWS_REGION = 'aws.region'.freeze
+          TAG_REGION = 'region'.freeze
+          TAG_AWS_SERVICE = 'aws_service'.freeze
+          TAG_AWS_ACCOUNT = 'aws_account'.freeze
+          TAG_QUEUE_NAME = 'queuename'.freeze
+          TAG_TOPIC_NAME = 'topicname'.freeze
+          TAG_TABLE_NAME = 'tablename'.freeze
+          TAG_STREAM_NAME = 'streamname'.freeze
+          TAG_RULE_NAME = 'rulename'.freeze
+          TAG_STATE_MACHINE_NAME = 'statemachinename'.freeze
+          TAG_BUCKET_NAME = 'bucketname'.freeze
         end
       end
     end

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

@@ -33,6 +33,12 @@ def annotate!(span, context)
             span.span_type = Tracing::Metadata::Ext::HTTP::TYPE_OUTBOUND
             span.name = Ext::SPAN_COMMAND
             span.resource = context.safely(:resource)
+            aws_service = span.resource.split('.')[0]
+            span.set_tag(Ext::TAG_AWS_SERVICE, aws_service)
+            params = context.safely(:params)
+            if (handler = Datadog::Tracing::Contrib::Aws::SERVICE_HANDLERS[aws_service])
+              handler.add_tags(span, params)
+            end
 
             span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_CLIENT)
 
@@ -55,6 +61,7 @@ def annotate!(span, context)
             span.set_tag(Ext::TAG_AGENT, Ext::TAG_DEFAULT_AGENT)
             span.set_tag(Ext::TAG_OPERATION, context.safely(:operation))
             span.set_tag(Ext::TAG_REGION, context.safely(:region))
+            span.set_tag(Ext::TAG_AWS_REGION, context.safely(:region))
             span.set_tag(Ext::TAG_PATH, context.safely(:path))
             span.set_tag(Ext::TAG_HOST, context.safely(:host))
             span.set_tag(Tracing::Metadata::Ext::HTTP::TAG_METHOD, context.safely(:http_method))

lib/datadog/tracing/contrib/aws/parsed_context.rb

@@ -20,6 +20,10 @@ def operation
             context.operation_name
           end
 
+          def params
+            context.params
+          end
+
           def status_code
             context.http_response.status_code
           end

lib/datadog/tracing/contrib/aws/service/base.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # Base class for all AWS service-specific tag handlers.
+          class Base
+            def add_tags(span, params); end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/service/dynamodb.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # DynamoDB tag handlers.
+          class DynamoDB < Base
+            def add_tags(span, params)
+              table_name = params[:table_name]
+              span.set_tag(Aws::Ext::TAG_TABLE_NAME, table_name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/service/eventbridge.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # EventBridge tag handlers.
+          class EventBridge < Base
+            def add_tags(span, params)
+              rule_name = params[:name] || params[:rule]
+              span.set_tag(Aws::Ext::TAG_RULE_NAME, rule_name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/service/kinesis.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # Kinesis tag handlers.
+          class Kinesis < Base
+            def add_tags(span, params)
+              stream_arn = params[:stream_arn]
+              stream_name = params[:stream_name]
+
+              if stream_arn
+                # example stream_arn: arn:aws:kinesis:us-east-1:123456789012:stream/my-stream
+                parts = stream_arn.split(':')
+                stream_name = parts[-1].split('/')[-1]
+                aws_account = parts[-2]
+                span.set_tag(Aws::Ext::TAG_AWS_ACCOUNT, aws_account)
+              end
+
+              span.set_tag(Aws::Ext::TAG_STREAM_NAME, stream_name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/service/s3.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # S3 tag handlers.
+          class S3 < Base
+            def add_tags(span, params)
+              bucket_name = params[:bucket]
+              span.set_tag(Aws::Ext::TAG_BUCKET_NAME, bucket_name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/service/sns.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # SNS tag handlers.
+          class SNS < Base
+            def add_tags(span, params)
+              topic_arn = params[:topic_arn]
+              topic_name = params[:name]
+              if topic_arn
+                # example topic_arn: arn:aws:sns:us-west-2:123456789012:my-topic-name
+                parts = topic_arn.split(':')
+                topic_name = parts[-1]
+                aws_account = parts[-2]
+                span.set_tag(Aws::Ext::TAG_AWS_ACCOUNT, aws_account)
+              end
+              span.set_tag(Aws::Ext::TAG_TOPIC_NAME, topic_name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/service/sqs.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # SQS tag handlers.
+          class SQS < Base
+            def add_tags(span, params)
+              queue_url = params[:queue_url]
+              queue_name = params[:queue_name]
+              if queue_url
+                _, _, _, aws_account, queue_name = queue_url.split('/')
+                span.set_tag(Aws::Ext::TAG_AWS_ACCOUNT, aws_account)
+              end
+              span.set_tag(Aws::Ext::TAG_QUEUE_NAME, queue_name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/service/stepfunctions.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Aws
+        module Service
+          # States tag handlers.
+          class States < Base
+            def add_tags(span, params)
+              state_machine_name = params[:name]
+              state_machine_arn = params[:state_machine_arn]
+              execution_arn = params[:execution_arn]
+              state_machine_account_id = ''
+
+              if execution_arn
+                # 'arn:aws:states:us-east-1:123456789012:execution:example-state-machine:example-execution'
+                parts = execution_arn.split(':')
+                state_machine_name = parts[-2]
+                state_machine_account_id = parts[4]
+              end
+
+              if state_machine_arn
+                # example statemachinearn: arn:aws:states:us-east-1:123456789012:stateMachine:MyStateMachine
+                parts = state_machine_arn.split(':')
+                state_machine_name ||= parts[-1]
+                state_machine_account_id = parts[-3]
+              end
+              span.set_tag(Aws::Ext::TAG_AWS_ACCOUNT, state_machine_account_id)
+              span.set_tag(Aws::Ext::TAG_STATE_MACHINE_NAME, state_machine_name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/aws/services.rb

@@ -113,6 +113,16 @@ module Aws
           WorkSpaces
           XRay
         ].freeze
+
+        SERVICE_HANDLERS = {
+          'sqs' => Service::SQS.new,
+          'sns' => Service::SNS.new,
+          'dynamodb' => Service::DynamoDB.new,
+          'kinesis' => Service::Kinesis.new,
+          'eventbridge' => Service::EventBridge.new,
+          'states' => Service::States.new,
+          's3' => Service::S3.new
+        }.freeze
       end
     end
   end