add whitelist for prometheus metrics (#280)

* add whitelist for prometheus metrics

Author: BrandonArp

src/main/java/com/arpnetworking/metrics/common/sources/PrometheusHttpSource.java

@@ -16,8 +16,11 @@
 package com.arpnetworking.metrics.common.sources;
 
 import com.arpnetworking.metrics.mad.parsers.PrometheusToRecordParser;
+import com.google.common.collect.Sets;
 import net.sf.oval.constraint.NotNull;
 
+import java.util.HashSet;
+
 /**
  * Processes Prometheus messages, extracts data and emits metrics.
  *
@@ -67,6 +70,18 @@ public Builder setOutputDebugFiles(final Boolean value) {
             return this;
         }
 
+        /**
+         * Whitelist for reserved metrics names (names with aggregation keys). By default, we ignore metrics with
+         * aggregation keys unless specified here. Cannot be null.
+         *
+         * @param value the value
+         * @return this {@link Builder}
+         */
+        public Builder setReservedNameWhitelist(final HashSet<String> value) {
+            _reservedNameWhitelist = value;
+            return this;
+        }
+
         @Override
         protected Builder self() {
             return this;
@@ -78,9 +93,12 @@ protected Builder self() {
         @NotNull
         private Boolean _outputDebugFiles = false;
 
+        @NotNull
+        private HashSet<String> _reservedNameWhitelist = Sets.newHashSet();
+
         @Override
         public PrometheusHttpSource build() {
-            setParser(new PrometheusToRecordParser(_interpretUnits, _outputDebugFiles));
+            setParser(new PrometheusToRecordParser(_interpretUnits, _outputDebugFiles, _reservedNameWhitelist));
             return super.build();
         }
     }

src/main/java/com/arpnetworking/metrics/mad/parsers/PrometheusToRecordParser.java

@@ -45,6 +45,7 @@
 import java.time.Instant;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -63,10 +64,15 @@ public final class PrometheusToRecordParser implements Parser<List<Record>, Http
      *
      * @param interpretUnits specifies whether to interpret units.
      * @param outputDebugInfo specifies whether to output debug files.
+     * @param reservedNameWhitelist specifies the whitelist set for reserved names (names with aggregation keys).
      */
-    public PrometheusToRecordParser(final boolean interpretUnits, final boolean outputDebugInfo) {
+    public PrometheusToRecordParser(
+            final boolean interpretUnits,
+            final boolean outputDebugInfo,
+            final HashSet<String> reservedNameWhitelist) {
         _interpretUnits = interpretUnits;
         _outputDebugInfo = outputDebugInfo;
+        _reservedNameWhitelist = reservedNameWhitelist;
     }
 
     /*
@@ -128,7 +134,7 @@ public List<Record> parse(final HttpRequest data) throws ParsingException {
                 }
                 final ParseResult result = parseNameAndUnit(nameOpt.orElse("").trim());
                 // We don't currently support aggregate metrics from prometheus
-                if (result.getAggregationKey().isPresent()) {
+                if (result.getAggregationKey().isPresent() && !_reservedNameWhitelist.contains(result.getName())) {
                     continue;
                 }
                 final String metricName = result.getName();
@@ -204,6 +210,7 @@ private Quantity createQuantity(final Types.Sample sample, final Optional<Unit>
     private final boolean _interpretUnits;
     private final AtomicInteger _outputFileNumber = new AtomicInteger(0);
     private final boolean _outputDebugInfo;
+    private final HashSet<String> _reservedNameWhitelist;
 
     private static final ImmutableMap<String, Unit> UNIT_MAP = ImmutableMap.of(
             "seconds", Unit.SECOND,

src/test/java/com/arpnetworking/metrics/mad/parsers/PrometheusToRecordParserTest.java

@@ -26,6 +26,7 @@
 import com.arpnetworking.metrics.mad.model.Unit;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
+import com.google.common.collect.Sets;
 import com.google.common.io.Resources;
 import org.junit.Assert;
 import org.junit.Ignore;
@@ -35,6 +36,7 @@
 import java.time.Instant;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import javax.annotation.Nullable;
@@ -166,6 +168,15 @@ public void testLive1() throws ParsingException, IOException {
         Assert.assertEquals(294, records.size());
     }
 
+    @Test
+    public void testLive1Whitelist() throws ParsingException, IOException {
+        final HashSet<String> whitelist = Sets.newHashSet("container_network_transmit_packets_total");
+        final PrometheusToRecordParser parser = new PrometheusToRecordParser(true, false, whitelist);
+        final List<Record> records = parseRecords("PrometheusParserTest/testLivePrometheus1", parser);
+
+        Assert.assertEquals(500, records.size());
+    }
+
     private static void testUnitParsing(final String prometheusUnit, final Unit expected) {
         assertUnitNewName(prometheusUnit, expected, null);
         assertUnitNewName("foo_" + prometheusUnit, expected, null);
@@ -204,9 +215,9 @@ private static List<Record> parseRecords(
     }
 
     private static PrometheusToRecordParser createParser() {
-        return new PrometheusToRecordParser(true, false);
+        return new PrometheusToRecordParser(true, false, Sets.newHashSet());
     }
     private static PrometheusToRecordParser createParserWithoutInterpreter() {
-        return new PrometheusToRecordParser(false, false);
+        return new PrometheusToRecordParser(false, false, Sets.newHashSet());
     }
 }