From dee040318ae28b230072eb9e5e0602dba67eb0a8 Mon Sep 17 00:00:00 2001
From: Ehren Kret <ehren@signal.org>
Date: Thu, 30 Jul 2020 11:59:16 -0500
Subject: [PATCH] Add the host tag to datadog metric reporting

---
 .../textsecuregcm/WhisperServerService.java   |  7 +++
 .../metrics/SignalDatadogReporterFactory.java | 63 +++++++++++++++++++
 .../io.dropwizard.metrics.ReporterFactory     |  3 +-
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/metrics/SignalDatadogReporterFactory.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index d6fbb6d7e..de341dc96 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -22,6 +22,7 @@ import com.amazonaws.auth.AWSStaticCredentialsProvider;
 import com.amazonaws.auth.BasicAWSCredentials;
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.AmazonS3Client;
+import com.amazonaws.util.EC2MetadataUtils;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.codahale.metrics.jdbi3.strategies.DefaultNameStrategy;
 import com.fasterxml.jackson.annotation.JsonAutoDetect;
@@ -252,6 +253,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     {
       final MicrometerConfiguration micrometerDatadogConfig = micrometerConfigurationByName.get("datadog");
 
+      final String instanceId = EC2MetadataUtils.getInstanceId();
       Metrics.addRegistry(new DatadogMeterRegistry(new DatadogConfig() {
         @Override
         public String get(final String key) {
@@ -262,6 +264,11 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         public String apiKey() {
           return micrometerDatadogConfig.getApiKey();
         }
+
+        @Override
+        public String hostTag() {
+          return instanceId;
+        }
       }, Clock.SYSTEM));
     }
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/SignalDatadogReporterFactory.java b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/SignalDatadogReporterFactory.java
new file mode 100644
index 000000000..e73c4a847
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/SignalDatadogReporterFactory.java
@@ -0,0 +1,63 @@
+package org.whispersystems.textsecuregcm.metrics;
+
+import com.amazonaws.util.EC2MetadataUtils;
+import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.ScheduledReporter;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonTypeName;
+import io.dropwizard.metrics.BaseReporterFactory;
+import org.coursera.metrics.datadog.DatadogReporter;
+import org.coursera.metrics.datadog.DefaultMetricNameFormatterFactory;
+import org.coursera.metrics.datadog.DynamicTagsCallbackFactory;
+import org.coursera.metrics.datadog.MetricNameFormatterFactory;
+import org.coursera.metrics.datadog.transport.AbstractTransportFactory;
+
+import javax.validation.Valid;
+import javax.validation.constraints.NotNull;
+import java.util.EnumSet;
+import java.util.List;
+
+@SuppressWarnings("FieldMayBeFinal")
+@JsonTypeName("signaldatadog")
+public class SignalDatadogReporterFactory extends BaseReporterFactory {
+
+  @JsonProperty
+  private List<String> tags = null;
+
+  @Valid
+  @JsonProperty
+  private DynamicTagsCallbackFactory dynamicTagsCallback = null;
+
+  @JsonProperty
+  private String prefix = null;
+
+  @Valid
+  @NotNull
+  @JsonProperty
+  private EnumSet<DatadogReporter.Expansion> expansions = DatadogReporter.Expansion.ALL;
+
+  @Valid
+  @NotNull
+  @JsonProperty
+  private MetricNameFormatterFactory metricNameFormatter = new DefaultMetricNameFormatterFactory();
+
+  @Valid
+  @NotNull
+  @JsonProperty
+  private AbstractTransportFactory transport = null;
+
+  public ScheduledReporter build(MetricRegistry registry) {
+    return DatadogReporter.forRegistry(registry)
+                          .withTransport(transport.build())
+                          .withHost(EC2MetadataUtils.getInstanceId())
+                          .withTags(tags)
+                          .withPrefix(prefix)
+                          .withExpansions(expansions)
+                          .withMetricNameFormatter(metricNameFormatter.build())
+                          .withDynamicTagCallback(dynamicTagsCallback != null ? dynamicTagsCallback.build() : null)
+                          .filter(getFilter())
+                          .convertDurationsTo(getDurationUnit())
+                          .convertRatesTo(getRateUnit())
+                          .build();
+  }
+}
diff --git a/service/src/main/resources/META-INF/services/io.dropwizard.metrics.ReporterFactory b/service/src/main/resources/META-INF/services/io.dropwizard.metrics.ReporterFactory
index bdc070f96..be127b6fd 100644
--- a/service/src/main/resources/META-INF/services/io.dropwizard.metrics.ReporterFactory
+++ b/service/src/main/resources/META-INF/services/io.dropwizard.metrics.ReporterFactory
@@ -1 +1,2 @@
-org.whispersystems.textsecuregcm.metrics.JsonMetricsReporterFactory
\ No newline at end of file
+org.whispersystems.textsecuregcm.metrics.JsonMetricsReporterFactory
+org.whispersystems.textsecuregcm.metrics.SignalDatadogReporterFactory