From 28181842c035e055b960cc587225ef36b30d6d14 Mon Sep 17 00:00:00 2001
From: Sam Protas <sprotas@50onred.com>
Date: Tue, 18 Sep 2018 14:20:34 -0400
Subject: [PATCH] Add ECS tags helper

---
 .../coursera/metrics/datadog/AwsHelper.java   | 56 +++++++++++++++++++
 .../metrics/datadog/DatadogReporter.java      | 54 +++++++++++++++++-
 2 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/metrics-datadog/src/main/java/org/coursera/metrics/datadog/AwsHelper.java b/metrics-datadog/src/main/java/org/coursera/metrics/datadog/AwsHelper.java
index c409c3d9..9d0eac60 100644
--- a/metrics-datadog/src/main/java/org/coursera/metrics/datadog/AwsHelper.java
+++ b/metrics-datadog/src/main/java/org/coursera/metrics/datadog/AwsHelper.java
@@ -1,5 +1,8 @@
 package org.coursera.metrics.datadog;
 
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.PropertyNamingStrategy;
 import org.apache.http.client.fluent.Request;
 
 import java.io.IOException;
@@ -8,6 +11,11 @@ public class AwsHelper {
 
   public static final String url = "http://169.254.169.254/latest/meta-data/instance-id";
 
+  // ECS details require the ecs config ECS_ENABLE_CONTAINER_METADATA=true on the underlying instance
+  private static final String ecs_metadata_file_env = "ECS_CONTAINER_METADATA_FILE";
+  private static final ObjectMapper objectMapper = new ObjectMapper()
+          .setPropertyNamingStrategy(PropertyNamingStrategy.UPPER_CAMEL_CASE);
+
   public static String getEc2InstanceId() throws IOException {
     try {
       return Request.Get(url).execute().returnContent().asString();
@@ -15,4 +23,52 @@ public static String getEc2InstanceId() throws IOException {
       throw new IOException(t);
     }
   }
+
+  public static String getEcsMetadataFileName() {
+    String metadataFilename = System.getenv(ecs_metadata_file_env);
+    if (metadataFilename == null || metadataFilename.equals("")) {
+      return null;
+    }
+
+    return metadataFilename;
+  }
+
+  public static boolean isRunningInEcs() {
+    return getEcsMetadataFileName() != null;
+  }
+
+  public static EcsMetadata getEcsMetadata() throws IOException {
+    String metadataFileName = getEcsMetadataFileName();
+    if (metadataFileName == null) {
+      return null;
+    }
+
+    EcsMetadata metadata;
+    try {
+      metadata = objectMapper.readValue(metadataFileName, EcsMetadata.class);
+    } catch (IOException e) {
+      throw e;
+    } catch (Throwable e) {
+      throw new IOException(e);
+    }
+    return metadata.getTaskDefinitionFamily() != null && metadata.getTaskDefinitionRevision() != null ? metadata : null;
+  }
+
+  public class EcsMetadata {
+
+    @JsonProperty
+    private String taskDefinitionFamily = null;
+
+    @JsonProperty
+    private String taskDefinitionRevision = null;
+
+    public String getTaskDefinitionFamily() {
+      return taskDefinitionFamily;
+    }
+
+    public String getTaskDefinitionRevision() {
+      return taskDefinitionRevision;
+    }
+  }
+
 }
\ No newline at end of file
diff --git a/metrics-datadog/src/main/java/org/coursera/metrics/datadog/DatadogReporter.java b/metrics-datadog/src/main/java/org/coursera/metrics/datadog/DatadogReporter.java
index 79ea2f67..61491161 100644
--- a/metrics-datadog/src/main/java/org/coursera/metrics/datadog/DatadogReporter.java
+++ b/metrics-datadog/src/main/java/org/coursera/metrics/datadog/DatadogReporter.java
@@ -308,6 +308,58 @@ public Builder withEC2Host() throws IOException {
       return this;
     }
 
+    /**
+     * Adds the task_name and task_version tags to the list sent to datadog with each and every metric. Gracefully
+     * ignored if not running in AWS ECS.
+     *
+     * As this information isn't guaranteed on container startup, allows for a configurable retry. See
+     * withECSTaskTagsNoWait for a method  that never retries.
+     *
+     * @param attempts
+     * @param delayMillis
+     * @return
+     * @throws IOException
+     */
+    public Builder withECSTaskTags(int attempts, long delayMillis) throws IOException {
+      if (!AwsHelper.isRunningInEcs() || attempts <= 0) {
+          return this;
+      }
+
+      long delay = Math.max(0, delayMillis);
+      int remainingAttempts = attempts;
+
+      while (true) {
+        AwsHelper.EcsMetadata metadata = AwsHelper.getEcsMetadata();
+
+        if (metadata != null) {
+          List<String> taskTags = new ArrayList<>(2);
+          taskTags.add("task_name:" + metadata.getTaskDefinitionFamily());
+          taskTags.add("task_version:" + metadata.getTaskDefinitionRevision());
+          return withTags(taskTags);
+        }
+
+        try {
+          remainingAttempts--;
+          if (remainingAttempts <= 0) {
+              return this;
+          }
+          Thread.sleep(delay);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          return this;
+        }
+      }
+    }
+
+    /**
+     * Convenience method to run withECSTaskTags only a single attempt.
+     * @return
+     * @throws IOException
+     */
+    public Builder withECSTaskTagsNoWait() throws IOException {
+        return withECSTaskTags(1, 0);
+    }
+
     public Builder withExpansions(EnumSet<Expansion> expansions) {
       this.expansions = expansions;
       return this;
@@ -329,7 +381,7 @@ public Builder convertRatesTo(TimeUnit rateUnit) {
      * @param tags List of tags eg: [env:prod, version:1.0.1, name:kafka_client] etc
      */
     public Builder withTags(List<String> tags) {
-      this.tags = tags;
+      this.tags = TagUtils.mergeTags(this.tags, tags);
       return this;
     }