! all: introduced support for metric tags.

1 files changed, 11 insertions, 6 deletions

diff --git a/kamon-newrelic/src/main/scala/kamon/newrelic/CustomMetricExtractor.scala b/kamon-newrelic/src/main/scala/kamon/newrelic/CustomMetricExtractor.scala
index 41ed9661..6919a967 100644
--- a/kamon-newrelic/src/main/scala/kamon/newrelic/CustomMetricExtractor.scala
+++ b/kamon-newrelic/src/main/scala/kamon/newrelic/CustomMetricExtractor.scala
@@ -16,17 +16,22 @@
 
 package kamon.newrelic
 
-import kamon.metric.{ SimpleMetricsImpl, EntitySnapshot, Entity }
+import kamon.metric.{ EntitySnapshot, Entity }
 import kamon.metric.instrument.CollectionContext
 
 object CustomMetricExtractor extends MetricExtractor {
 
   def extract(settings: AgentSettings, collectionContext: CollectionContext, metrics: Map[Entity, EntitySnapshot]): Map[MetricID, MetricData] = {
-    metrics.get(SimpleMetricsImpl.SimpleMetricsEntity).map { allSimpleMetrics ⇒
-      allSimpleMetrics.metrics.map {
-        case (key, snapshot) ⇒ Metric(snapshot, key.unitOfMeasurement, s"Custom/${key.name}", None)
-      }
+    def onlySimpleMetrics(kv: (Entity, EntitySnapshot)): Boolean =
+      kamon.metric.SingleInstrumentEntityRecorder.AllCategories.contains(kv._1.category)
 
-    } getOrElse (Map.empty)
+    def toNewRelicMetric(kv: (Entity, EntitySnapshot)): (MetricID, MetricData) = {
+      val (entity, entitySnapshot) = kv
+      val (metricKey, instrumentSnapshot) = entitySnapshot.metrics.head
+
+      Metric(instrumentSnapshot, metricKey.unitOfMeasurement, s"Custom/${entity.name}", None)
+    }
+
+    metrics.filter(onlySimpleMetrics).map(toNewRelicMetric)
   }
 }