From 09245ac9297c8ef1563a27450c1a4ed7910a441d Mon Sep 17 00:00:00 2001
From: Juan Carrey <juankarrey@gmail.com>
Date: Fri, 24 May 2024 12:48:25 +0200
Subject: [PATCH] ISSUE-1337 Add EntityID to DatadogAgentReporter (#1338)

* ISSUE-1337 Add EntityID to DatadogAgentReporter

* fix: EntityID could be provided from config, but for now only env supported
---
 .../kamon/datadog/DatadogAgentReporter.scala   | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/reporters/kamon-datadog/src/main/scala/kamon/datadog/DatadogAgentReporter.scala b/reporters/kamon-datadog/src/main/scala/kamon/datadog/DatadogAgentReporter.scala
index b6a804ec3..f7cb23e51 100644
--- a/reporters/kamon-datadog/src/main/scala/kamon/datadog/DatadogAgentReporter.scala
+++ b/reporters/kamon-datadog/src/main/scala/kamon/datadog/DatadogAgentReporter.scala
@@ -50,6 +50,9 @@ class DatadogAgentReporter private[datadog] (@volatile private var config: Datad
     new DecimalFormat("#.################################################################", symbols)
   private val valueFormat = new DecimalFormat("#0.#########", symbols)
 
+  private val DD_ENTITY_ID_ENV_VAR = "DD_ENTITY_ID"
+  private val ENTITY_ID_TAG_NAME = "dd.internal.entity_id"
+
   logger.info("Started the Kamon Datadog reporter")
 
   override def stop(): Unit = {}
@@ -68,7 +71,7 @@ class DatadogAgentReporter private[datadog] (@volatile private var config: Datad
         counter.name,
         config.measurementFormatter.formatMeasurement(
           encodeDatadogCounter(instrument.value, counter.settings.unit),
-          instrument.tags
+          updateTagsWithEntityID(instrument.tags)
         )
       )
     }
@@ -81,7 +84,7 @@ class DatadogAgentReporter private[datadog] (@volatile private var config: Datad
         gauge.name,
         config.measurementFormatter.formatMeasurement(
           encodeDatadogGauge(instrument.value, gauge.settings.unit),
-          instrument.tags
+          updateTagsWithEntityID(instrument.tags)
         )
       )
     }
@@ -94,7 +97,7 @@ class DatadogAgentReporter private[datadog] (@volatile private var config: Datad
 
       val bucketData = config.measurementFormatter.formatMeasurement(
         encodeDatadogHistogramBucket(bucket.value, bucket.frequency, metric.settings.unit),
-        instruments.tags
+        updateTagsWithEntityID(instruments.tags)
       )
       config.packetBuffer.appendMeasurement(metric.name, bucketData)
     }
@@ -103,6 +106,15 @@ class DatadogAgentReporter private[datadog] (@volatile private var config: Datad
 
   }
 
+  private def updateTagsWithEntityID(tags: TagSet): TagSet = {
+    val entityId = System.getenv("DD_ENTITY_ID")
+    if (entityId != null && !entityId.trim().isEmpty()) {
+      return tags.withTag(ENTITY_ID_TAG_NAME, entityId);
+    }
+
+    return tags;
+  }
+
   private def encodeDatadogHistogramBucket(value: Long, frequency: Long, unit: MeasurementUnit): String = {
     val metricType = if (unit.dimension == Dimension.Time) "ms" else "h"
     val samplingRate: Double = 1d / frequency.toDouble