Add ECS metadata allowing cloudwatch-logs to be linked with traces

lib/aws-xray-sdk/plugins/ecs.rb

@@ -3,19 +3,82 @@
 
 module XRay
   module Plugins
-    # Due to lack of ECS container metadata service, the only host information
-    # available is the host name.
     module ECS
       include Logging
 
       ORIGIN = 'AWS::ECS::Container'.freeze
 
+      # Only compatible with v4!
+      # The v3 metadata url does not contain cloudwatch informations
+      METADATA_ENV_KEY = 'ECS_CONTAINER_METADATA_URI_V4'
+
       def self.aws
-        @@aws ||= begin
-          { ecs: { container: Socket.gethostname } }
+        metadata = get_metadata()
+
+        begin
+          metadata[:ecs][:container] = Socket.gethostname
         rescue StandardError => e
-          @@aws = {}
           Logging.logger.warn %(cannot get the ecs container hostname due to: #{e.message}.)
+          metadata[:ecs][:container] = nil
+        end
+
+        @@aws = {
+          ecs: metadata[:ecs],
+          cloudwatch_logs: metadata[:cloudwatch_logs]
+        }
+      end
+
+      private
+
+      def self.get_metadata()
+        begin
+          metadata_uri = URI(ENV[METADATA_ENV_KEY])
+          req = Net::HTTP::Get.new(metadata_uri)
+          metadata_json = do_request(req)
+          return parse_metadata(metadata_json)
+        rescue StandardError => e
+          Logging.logger.warn %(cannot get the ecs instance metadata due to: #{e.message}. Make sure you are using Fargate platform version >=1.4.0)
+          { ecs: {}, cloudwatch_logs: {} }
+        end
+      end
+
+      def self.parse_metadata(json_str)
+        data = JSON(json_str)
+
+        metadata = {
+          ecs: {
+            container_arn: data['ContainerARN'],
+          },
+          cloudwatch_logs: {
+            log_group: data["LogOptions"]['awslogs-group'],
+            log_region: data["LogOptions"]['awslogs-region'],
+            arn: data['ContainerARN']
+          }
+        }
+        metadata
+      end
+
+      def self.do_request(request)
+        begin
+          response = Net::HTTP.start(request.uri.hostname, read_timeout: 1) { |http|
+            http.request(request)
+          }
+
+          if response.code == '200'
+            return response.body
+          else
+            raise(StandardError.new('Unsuccessful response::' + response.code + '::' + response.message))
+          end
+        rescue StandardError => e
+          # Two attempts in total to complete the request successfully
+          @retries ||= 0
+          if @retries < 1
+            @retries += 1
+            retry
+          else
+            Logging.logger.warn %(Failed to complete request due to: #{e.message}.)
+            raise e
+          end
         end
       end
     end

test/aws-xray-sdk/tc_plugin.rb

@@ -25,6 +25,7 @@ def test_get_runtime_context
     WebMock.reset!
   end
 
+  # EC2 Plugin
   def test_ec2_metadata_v2_successful
     dummy_json = '{\"availabilityZone\" : \"us-east-2a\", \"imageId\" : \"ami-03cca83dd001d4666\",
                   \"instanceId\" : \"i-07a181803de94c666\", \"instanceType\" : \"t3.xlarge\"}'
@@ -43,6 +44,7 @@ def test_ec2_metadata_v2_successful
         ami_id: 'ami-03cca83dd001d4666'
       }
     }
+    # We should probably use `assert_equal` here ? Always true otherwise...
     assert expected, XRay::Plugins::EC2.aws
     WebMock.reset!
   end
@@ -80,4 +82,53 @@ def test_ec2_metadata_fail
     assert expected, XRay::Plugins::EC2.aws
     WebMock.reset!
   end
+
+  # ECS Plugin
+  def test_ecs_metadata_successful
+    dummy_metadata_uri = 'http://169.254.170.2/v4/a_random_id'
+    dummy_json = {
+      "ContainerARN"=>"arn:aws:ecs:eu-central-1:an_id:container/a_cluster/a_cluster_id/a_task_id",
+      "LogOptions"=>{"awslogs-group"=>"/ecs/a_service_name", "awslogs-region"=>"eu-central-1", "awslogs-stream"=>"ecs/a_service_name/a_task_id"},
+    }
+
+    ENV[XRay::Plugins::ECS::METADATA_ENV_KEY] = dummy_metadata_uri
+    stub_request(:get, dummy_metadata_uri)
+      .to_return(status: 200, body: dummy_json.to_json, headers: {})
+
+    expected = {
+      ecs: {
+        container: Socket.gethostname,
+        container_arn: 'arn:aws:ecs:eu-central-1:an_id:container/a_cluster/a_cluster_id/a_task_id',
+      },
+      cloudwatch_logs: {:log_group=>"/ecs/a_service_name", :log_region=>"eu-central-1", :arn=>"arn:aws:ecs:eu-central-1:an_id:container/a_cluster/a_cluster_id/a_task_id"}
+    }
+    assert_equal expected, XRay::Plugins::ECS.aws
+    WebMock.reset!
+    ENV.delete(XRay::Plugins::ECS::METADATA_ENV_KEY)
+  end
+
+  def test_ecs_metadata_fail
+    dummy_metadata_uri = 'http://169.254.170.2/v4/a_random_id'
+    ENV['ECS_CONTAINER_METADATA_URI_V4'] = dummy_metadata_uri
+
+    stub_request(:get, dummy_metadata_uri)
+      .to_raise(StandardError)
+
+    expected = {
+      ecs: {container: Socket.gethostname},
+      cloudwatch_logs: {}
+    }
+    assert_equal expected, XRay::Plugins::ECS.aws
+    WebMock.reset!
+    ENV.delete(XRay::Plugins::ECS::METADATA_ENV_KEY)
+  end
+
+  def test_ecs_metadata_not_defined
+    expected = {
+      ecs: {container: Socket.gethostname},
+      cloudwatch_logs: {}
+    }
+    assert_equal expected, XRay::Plugins::ECS.aws
+    WebMock.reset!
+  end
 end