Added log correlation to ECS plugin (#263)

* added log correlation to ECS plugin

* comments and ci

Author: willarmiros

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/EC2MetadataFetcher.java

@@ -18,22 +18,14 @@
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
-import java.io.UnsupportedEncodingException;
-import java.net.HttpURLConnection;
 import java.net.MalformedURLException;
-import java.net.ProtocolException;
-import java.net.SocketTimeoutException;
 import java.net.URL;
-import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-import org.checkerframework.checker.nullness.qual.Nullable;
 
 class EC2MetadataFetcher {
     private static final Log logger = LogFactory.getLog(EC2MetadataFetcher.class);
@@ -47,8 +39,7 @@ enum EC2Metadata {
         AMI_ID,
     }
 
-    private static final int CONNECT_TIMEOUT_MILLIS = 100;
-    private static final int READ_TIMEOUT_MILLIS = 1000;
+    private static final String METADATA_SERVICE_NAME = "IMDS";
     private static final String DEFAULT_IMDS_ENDPOINT = "169.254.169.254";
 
     private final URL identityDocumentUrl;
@@ -120,86 +111,11 @@ Map<EC2Metadata, String> fetch() {
     }
 
     private String fetchToken() {
-        return fetchString("PUT", tokenUrl, "", true);
+        return MetadataUtils.fetchString("PUT", tokenUrl, "", true, METADATA_SERVICE_NAME);
     }
 
     private String fetchIdentity(String token) {
-        return fetchString("GET", identityDocumentUrl, token, false);
-    }
-
-    // Generic HTTP fetch function for IMDS.
-    private static String fetchString(String httpMethod, URL url, String token, boolean includeTtl) {
-        final HttpURLConnection connection;
-        try {
-            connection = (HttpURLConnection) url.openConnection();
-        } catch (Exception e) {
-            logger.debug("Error connecting to IMDS.", e);
-            return "";
-        }
-
-        try {
-            connection.setRequestMethod(httpMethod);
-        } catch (ProtocolException e) {
-            logger.warn("Unknown HTTP method, this is a programming bug.", e);
-            return "";
-        }
-
-        connection.setConnectTimeout(CONNECT_TIMEOUT_MILLIS);
-        connection.setReadTimeout(READ_TIMEOUT_MILLIS);
-
-        if (includeTtl) {
-            connection.setRequestProperty("X-aws-ec2-metadata-token-ttl-seconds", "60");
-        }
-        if (!token.isEmpty()) {
-            connection.setRequestProperty("X-aws-ec2-metadata-token", token);
-        }
-
-        final int responseCode;
-        try {
-            responseCode = connection.getResponseCode();
-        } catch (Exception e) {
-            if (e instanceof SocketTimeoutException) {
-                logger.debug("Timed out trying to connect to IMDS, likely not operating in EC2 environment");
-            } else {
-                logger.debug("Error connecting to IMDS.", e);
-            }
-            return "";
-        }
-
-        if (responseCode != 200) {
-            logger.warn("Error reponse from IMDS: code (" + responseCode + ") text " + readResponseString(connection));
-        }
-
-        return readResponseString(connection).trim();
-    }
-
-    private static String readResponseString(HttpURLConnection connection) {
-        ByteArrayOutputStream os = new ByteArrayOutputStream();
-        try (InputStream is = connection.getInputStream()) {
-            readTo(is, os);
-        } catch (IOException e) {
-            // Only best effort read if we can.
-        }
-        try (InputStream is = connection.getErrorStream()) {
-            readTo(is, os);
-        } catch (IOException e) {
-            // Only best effort read if we can.
-        }
-        try {
-            return os.toString(StandardCharsets.UTF_8.name());
-        } catch (UnsupportedEncodingException e) {
-            throw new IllegalStateException("UTF-8 not supported can't happen.");
-        }
-    }
-
-    private static void readTo(@Nullable InputStream is, ByteArrayOutputStream os) throws IOException {
-        if (is == null) {
-            return;
-        }
-        int b;
-        while ((b = is.read()) != -1) {
-            os.write(b);
-        }
+        return MetadataUtils.fetchString("GET", identityDocumentUrl, token, false, METADATA_SERVICE_NAME);
     }
 
     private static String getEndpoint() {

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/ECSMetadataFetcher.java

@@ -0,0 +1,119 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amazonaws.xray.plugins;
+
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.JsonToken;
+import java.io.IOException;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+class ECSMetadataFetcher {
+    private static final Log logger = LogFactory.getLog(ECSMetadataFetcher.class);
+
+    private static final String METADATA_SERVICE_NAME = "TMDE";
+    private static final JsonFactory JSON_FACTORY = new JsonFactory();
+
+    @Nullable
+    private final URL containerUrl;
+
+    // TODO: Record additional attributes in runtime context from Task Metadata Endpoint
+    enum ECSContainerMetadata {
+        LOG_DRIVER,
+        LOG_GROUP_REGION,
+        LOG_GROUP_NAME,
+        CONTAINER_ARN,
+    }
+
+    ECSMetadataFetcher(@Nullable String endpoint) {
+        if (endpoint == null) {
+            this.containerUrl = null;
+            return;
+        }
+
+        try {
+            this.containerUrl = new URL(endpoint);
+        } catch (MalformedURLException e) {
+            throw new IllegalArgumentException("Illegal endpoint: " + endpoint);
+        }
+    }
+
+    Map<ECSContainerMetadata, String> fetchContainer() {
+        if (this.containerUrl == null) {
+            return Collections.emptyMap();
+        }
+
+        String metadata = MetadataUtils.fetchString("GET", this.containerUrl, "", false, METADATA_SERVICE_NAME);
+
+        Map<ECSContainerMetadata, String> result = new HashMap<>();
+        try (JsonParser parser = JSON_FACTORY.createParser(metadata)) {
+            parser.nextToken();
+            parseContainerJson(parser, result);
+        } catch (IOException e) {
+            logger.warn("Could not parse container metadata.", e);
+            return Collections.emptyMap();
+        }
+
+        // This means the document didn't have all the metadata fields we wanted.
+        if (result.size() != ECSContainerMetadata.values().length) {
+            logger.debug("Container metadata response missing metadata: " + metadata);
+        }
+
+        return Collections.unmodifiableMap(result);
+    }
+
+    // Helper method to shallow-parse a JSON object, assuming the parser is located at the start of an object,
+    // and record the desired fields in the result map in-place
+    private void parseContainerJson(JsonParser parser, Map<ECSContainerMetadata, String> result) throws IOException {
+        if (!parser.isExpectedStartObjectToken()) {
+            logger.warn("Container metadata endpoint returned invalid JSON");
+            return;
+        }
+
+        while (parser.nextToken() != JsonToken.END_OBJECT) {
+            String value = parser.nextTextValue();
+            switch (parser.getCurrentName()) {
+                case "LogDriver":
+                    result.put(ECSContainerMetadata.LOG_DRIVER, value);
+                    break;
+                case "ContainerARN":
+                    result.put(ECSContainerMetadata.CONTAINER_ARN, value);
+                    break;
+                case "awslogs-group":
+                    result.put(ECSContainerMetadata.LOG_GROUP_NAME, value);
+                    break;
+                case "awslogs-region":
+                    result.put(ECSContainerMetadata.LOG_GROUP_REGION, value);
+                    break;
+                case "LogOptions":
+                    parseContainerJson(parser, result);  // Parse the LogOptions object for log fields
+                    break;
+                default:
+                    parser.skipChildren();
+            }
+            if (result.size() == ECSContainerMetadata.values().length) {
+                return;
+            }
+        }
+    }
+}

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/plugins/ECSPlugin.java

@@ -15,12 +15,15 @@
 
 package com.amazonaws.xray.plugins;
 
+import com.amazonaws.xray.entities.AWSLogReference;
 import com.amazonaws.xray.utils.DockerUtils;
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -38,29 +41,43 @@ public class ECSPlugin implements Plugin {
     private static final Log logger = LogFactory.getLog(ECSPlugin.class);
 
     private static final String SERVICE_NAME = "ecs";
-    private static final String ECS_METADATA_KEY = "ECS_CONTAINER_METADATA_URI";
-    private static final String HTTP_PREFIX = "http://";
+    private static final String ECS_METADATA_V3_KEY = "ECS_CONTAINER_METADATA_URI";
+    private static final String ECS_METADATA_V4_KEY = "ECS_CONTAINER_METADATA_URI_V4";
     private static final String CONTAINER_ID_KEY = "container_id";
+    private static final String CONTAINER_NAME_KEY = "container";
+    private static final String CONTAINER_ARN_KEY = "container_arn";
 
+    private final ECSMetadataFetcher fetcher;
     private final HashMap<String, @Nullable Object> runtimeContext;
     private final DockerUtils dockerUtils;
+    private final Set<AWSLogReference> logReferences;
+    private final Map<ECSMetadataFetcher.ECSContainerMetadata, String> containerMetadata;
 
+    @SuppressWarnings("nullness:method.invocation.invalid")
     public ECSPlugin() {
         runtimeContext = new HashMap<>();
         dockerUtils = new DockerUtils();
+        logReferences = new HashSet<>();
+        fetcher = new ECSMetadataFetcher(getTmdeFromEnv());
+        containerMetadata = this.fetcher.fetchContainer();
+    }
+
+    // Exposed for testing
+    ECSPlugin(ECSMetadataFetcher fetcher) {
+        runtimeContext = new HashMap<>();
+        dockerUtils = new DockerUtils();
+        logReferences = new HashSet<>();
+        this.fetcher = fetcher;
+        containerMetadata = this.fetcher.fetchContainer();
     }
 
     /**
      * Returns true if the environment variable added by ECS is present and contains a valid URI
      */
     @Override
     public boolean isEnabled() {
-        String ecsMetadataUri = System.getenv(ECS_METADATA_KEY);
-        if (ecsMetadataUri == null) {
-            return false;
-        }
-
-        return ecsMetadataUri.startsWith(HTTP_PREFIX);
+        String ecsMetadataUri = getTmdeFromEnv();
+        return ecsMetadataUri != null && ecsMetadataUri.startsWith("http://");
     }
 
     @Override
@@ -70,7 +87,7 @@ public String getServiceName() {
 
     public void populateRuntimeContext() {
         try {
-            runtimeContext.put("container", InetAddress.getLocalHost().getHostName());
+            runtimeContext.put(CONTAINER_NAME_KEY, InetAddress.getLocalHost().getHostName());
         } catch (UnknownHostException uhe) {
             logger.error("Could not get docker container ID from hostname.", uhe);
         }
@@ -80,6 +97,10 @@ public void populateRuntimeContext() {
         } catch (IOException e) {
             logger.error("Failed to read full container ID from container instance.", e);
         }
+
+        if (containerMetadata.containsKey(ECSMetadataFetcher.ECSContainerMetadata.CONTAINER_ARN)) {
+            runtimeContext.put(CONTAINER_ARN_KEY, containerMetadata.get(ECSMetadataFetcher.ECSContainerMetadata.CONTAINER_ARN));
+        }
     }
 
     @Override
@@ -88,6 +109,34 @@ public void populateRuntimeContext() {
         return runtimeContext;
     }
 
+    @Override
+    public Set<AWSLogReference> getLogReferences() {
+        if (logReferences.isEmpty()) {
+            populateLogReferences();
+        }
+
+        return logReferences;
+    }
+
+    // See: https://docs.aws.amazon.com/AmazonECS/latest/developerguide/ecs-account-settings.html#ecs-resource-ids
+    private void populateLogReferences() {
+        String logGroup = containerMetadata.get(ECSMetadataFetcher.ECSContainerMetadata.LOG_GROUP_NAME);
+        if (logGroup == null) {
+            return;
+        }
+        AWSLogReference logReference = new AWSLogReference();
+        logReference.setLogGroup(logGroup);
+
+        String logRegion = containerMetadata.get(ECSMetadataFetcher.ECSContainerMetadata.LOG_GROUP_REGION);
+        String containerArn = containerMetadata.get(ECSMetadataFetcher.ECSContainerMetadata.CONTAINER_ARN);
+        String logAccount = containerArn != null ? containerArn.split(":")[4] : null;
+
+        if (logRegion != null && logAccount != null) {
+            logReference.setArn("arn:aws:logs:" + logRegion + ":" + logAccount + ":log-group:" + logGroup);
+        }
+        logReferences.add(logReference);
+    }
+
     @Override
     public String getOrigin() {
         return ORIGIN;
@@ -109,4 +158,17 @@ public boolean equals(@Nullable Object o) {
     public int hashCode() {
         return this.getOrigin().hashCode();
     }
+
+    /**
+     * @return V4 Metadata endpoint if present, otherwise V3 endpoint if present, otherwise null
+     */
+    @Nullable
+    private String getTmdeFromEnv() {
+        String ecsMetadataUri = System.getenv(ECS_METADATA_V4_KEY);
+        if (ecsMetadataUri == null) {
+            ecsMetadataUri = System.getenv(ECS_METADATA_V3_KEY);
+        }
+
+        return ecsMetadataUri;
+    }
 }