Fix possible memory leak in ObjectNameAttributeFilter

collector/src/main/java/io/prometheus/jmx/JmxScraper.java

@@ -22,23 +22,8 @@
 import io.prometheus.jmx.logger.LoggerFactory;
 import java.io.IOException;
 import java.lang.management.ManagementFactory;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedHashMap;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.TreeSet;
-import javax.management.Attribute;
-import javax.management.AttributeList;
-import javax.management.JMException;
-import javax.management.MBeanAttributeInfo;
-import javax.management.MBeanInfo;
-import javax.management.MBeanServerConnection;
-import javax.management.ObjectInstance;
-import javax.management.ObjectName;
+import java.util.*;
+import javax.management.*;
 import javax.management.openmbean.CompositeData;
 import javax.management.openmbean.CompositeType;
 import javax.management.openmbean.TabularData;
@@ -144,8 +129,10 @@ public void doScrape() throws Exception {
                 }
             }
 
-            // Now that we have *only* the whitelisted mBeans, remove any old ones from the cache:
+            // Now that we have *only* the whitelisted mBeans, remove any old ones from the cache
+            // and dynamic attribute filter:
             jmxMBeanPropertyCache.onlyKeepMBeans(mBeanNames);
+            objectNameAttributeFilter.onlyKeepMBeans(mBeanNames);
 
             for (ObjectName objectName : mBeanNames) {
                 long start = System.nanoTime();

collector/src/main/java/io/prometheus/jmx/ObjectNameAttributeFilter.java

@@ -41,13 +41,15 @@ public class ObjectNameAttributeFilter {
     public static final String AUTO_EXCLUDE_OBJECT_NAME_ATTRIBUTES =
             "autoExcludeObjectNameAttributes";
 
-    private final Map<ObjectName, Set<String>> excludeObjectNameAttributesMap;
+    private final Map<ObjectName, Set<String>> configExcludeObjectNameAttributesMap;
+    private final Map<ObjectName, Set<String>> dynamicExcludeObjectNameAttributesMap;
 
     private boolean autoExcludeObjectNameAttributes;
 
     /** Constructor */
     private ObjectNameAttributeFilter() {
-        excludeObjectNameAttributesMap = new ConcurrentHashMap<>();
+        configExcludeObjectNameAttributesMap = new ConcurrentHashMap<>();
+        dynamicExcludeObjectNameAttributesMap = new ConcurrentHashMap<>();
     }
 
     /**
@@ -69,15 +71,10 @@ private ObjectNameAttributeFilter initialize(Map<String, Object> yamlConfig)
                 List<String> attributeNames = (List<String>) entry.getValue();
 
                 Set<String> attributeNameSet =
-                        excludeObjectNameAttributesMap.computeIfAbsent(
+                        configExcludeObjectNameAttributesMap.computeIfAbsent(
                                 objectName, o -> Collections.synchronizedSet(new HashSet<>()));
 
                 attributeNameSet.addAll(attributeNames);
-                for (String attribueName : attributeNames) {
-                    attributeNameSet.add(attribueName);
-                }
-
-                excludeObjectNameAttributesMap.put(objectName, attributeNameSet);
             }
         }
 
@@ -102,7 +99,7 @@ private ObjectNameAttributeFilter initialize(Map<String, Object> yamlConfig)
     public void add(ObjectName objectName, String attributeName) {
         if (autoExcludeObjectNameAttributes) {
             Set<String> attribteNameSet =
-                    excludeObjectNameAttributesMap.computeIfAbsent(
+                    dynamicExcludeObjectNameAttributesMap.computeIfAbsent(
                             objectName, o -> Collections.synchronizedSet(new HashSet<>()));
 
             LOGGER.log(
@@ -115,6 +112,16 @@ public void add(ObjectName objectName, String attributeName) {
         }
     }
 
+    public void onlyKeepMBeans(Set<ObjectName> aliveMBeans) {
+        if (autoExcludeObjectNameAttributes) {
+            for (ObjectName prevName : dynamicExcludeObjectNameAttributesMap.keySet()) {
+                if (!aliveMBeans.contains(prevName)) {
+                    dynamicExcludeObjectNameAttributesMap.remove(prevName);
+                }
+            }
+        }
+    }
+
     /**
      * Method to check if an attribute should be excluded
      *
@@ -123,15 +130,21 @@ public void add(ObjectName objectName, String attributeName) {
      * @return true if it should be excluded, false otherwise
      */
     public boolean exclude(ObjectName objectName, String attributeName) {
-        boolean result = false;
+        return exclude(configExcludeObjectNameAttributesMap, objectName, attributeName)
+                || exclude(dynamicExcludeObjectNameAttributesMap, objectName, attributeName);
+    }
 
-        if (excludeObjectNameAttributesMap.size() > 0) {
-            Set<String> attributeNameSet = excludeObjectNameAttributesMap.get(objectName);
+    private boolean exclude(
+            Map<ObjectName, Set<String>> exclusionMap,
+            ObjectName objectName,
+            String attributeName) {
+        boolean result = false;
+        if (!exclusionMap.isEmpty()) {
+            Set<String> attributeNameSet = exclusionMap.get(objectName);
             if (attributeNameSet != null) {
                 result = attributeNameSet.contains(attributeName);
             }
         }
-
         return result;
     }
 

collector/src/test/java/io/prometheus/jmx/ObjectNameAttributeFilterTest.java

@@ -0,0 +1,165 @@
+package io.prometheus.jmx;
+
+import static org.junit.Assert.*;
+
+import java.lang.reflect.Field;
+import java.util.*;
+import javax.management.MalformedObjectNameException;
+import javax.management.ObjectName;
+import org.junit.Test;
+import org.yaml.snakeyaml.Yaml;
+
+public class ObjectNameAttributeFilterTest {
+
+    static final String CONFIG_EXCLUDE_MAP_FIELD = "configExcludeObjectNameAttributesMap";
+    static final String DYNAMIC_EXCLUDE_MAP_FIELD = "dynamicExcludeObjectNameAttributesMap";
+
+    @Test
+    public void emptyConfig() throws Exception {
+        ObjectNameAttributeFilter filter = initEmptyConfigFilter();
+
+        assertEquals(0, excludeMapSize(filter, CONFIG_EXCLUDE_MAP_FIELD));
+        assertEquals(0, excludeMapSize(filter, DYNAMIC_EXCLUDE_MAP_FIELD));
+    }
+
+    @Test
+    public void nonEmptyConfigShouldInitializeConfigExcludeMap() throws Exception {
+        ObjectNameAttributeFilter filter = initNonEmptyConfigFilter();
+
+        Map<ObjectName, Set<String>> configMap =
+                getInternalMapValue(filter, CONFIG_EXCLUDE_MAP_FIELD);
+        Map<ObjectName, Set<String>> dynamicMap =
+                getInternalMapValue(filter, DYNAMIC_EXCLUDE_MAP_FIELD);
+
+        assertEquals(2, configMap.size());
+        assertEquals(1, configMap.get(new ObjectName("java.lang:type=OperatingSystem")).size());
+        assertEquals(2, configMap.get(new ObjectName("java.lang:type=Runtime")).size());
+        assertEquals(0, dynamicMap.size());
+    }
+
+    @Test
+    public void excludeShouldCheckConfigExclusionMap() throws Exception {
+        ObjectNameAttributeFilter filter = initNonEmptyConfigFilter();
+
+        boolean result = filter.exclude(new ObjectName("java.lang:type=Runtime"), "ClassPath");
+
+        assertTrue("java.lang:type=Runtime<>ClassPath should be excluded by config", result);
+    }
+
+    @Test
+    public void excludeShouldCheckDynamicExclusionMap() throws Exception {
+        ObjectNameAttributeFilter filter = initEmptyConfigFilter();
+        filter.add(new ObjectName("boolean:Type=Test"), "Value");
+
+        boolean result = filter.exclude(new ObjectName("boolean:Type=Test"), "Value");
+
+        assertTrue("boolean:Type=Test<>Value should be excluded dynamically", result);
+    }
+
+    @Test
+    public void onlyKeepMBeansShouldRemoveDeadDynamicRecords() throws Exception {
+        ObjectNameAttributeFilter filter = initEmptyConfigFilter();
+        Set<ObjectName> aliveMBeans = getAliveMBeans();
+        int size = aliveMBeans.size();
+        for (ObjectName objectName : aliveMBeans) {
+            filter.add(objectName, "Value");
+        }
+
+        Map<ObjectName, Set<String>> dynamicMap =
+                getInternalMapValue(filter, DYNAMIC_EXCLUDE_MAP_FIELD);
+
+        Iterator<ObjectName> iterator = aliveMBeans.iterator();
+        ObjectName unregisteredObjectName = iterator.next();
+        iterator.remove();
+        filter.onlyKeepMBeans(aliveMBeans);
+        assertEquals("onlyKeepMBeans should shrink the map", size - 1, dynamicMap.size());
+        for (ObjectName objectName : aliveMBeans) {
+            assertTrue(
+                    objectName
+                            + "<>Value should still be excluded dynamically after onlyKeepMBeans",
+                    filter.exclude(objectName, "Value"));
+        }
+        assertFalse(
+                unregisteredObjectName + "<>Value should not be excluded dynamically before add",
+                filter.exclude(unregisteredObjectName, "Value"));
+    }
+
+    @Test
+    public void onlyKeepMBeansShouldNotRemoveConfigRecords() throws Exception {
+        ObjectNameAttributeFilter objectNameAttributeFilter = initNonEmptyConfigFilter();
+        Set<ObjectName> aliveMBeans = getAliveMBeans();
+
+        for (ObjectName objectName : aliveMBeans) {
+            objectNameAttributeFilter.add(objectName, "Value");
+        }
+        objectNameAttributeFilter.onlyKeepMBeans(Collections.emptySet());
+        Map<ObjectName, Set<String>> configExcludeObjectNameAttributesMap =
+                getInternalMapValue(objectNameAttributeFilter, CONFIG_EXCLUDE_MAP_FIELD);
+        Map<ObjectName, Set<String>> dynamicExcludeObjectNameAttributesMap =
+                getInternalMapValue(objectNameAttributeFilter, DYNAMIC_EXCLUDE_MAP_FIELD);
+        assertEquals(
+                "configExcludeObjectNameAttributesMap should be left untouched after"
+                        + " onlyKeepMBeans(emptyList())",
+                2,
+                configExcludeObjectNameAttributesMap.size());
+        assertEquals(
+                "dynamicExcludeObjectNameAttributesMap should be empty after"
+                        + " onlyKeepMBeans(emptyList())",
+                0,
+                dynamicExcludeObjectNameAttributesMap.size());
+        assertTrue(
+                "java.lang:type=Runtime<>ClassPath should be excluded by config and not removed by"
+                        + " onlyKeepMBeans",
+                objectNameAttributeFilter.exclude(
+                        new ObjectName("java.lang:type=Runtime"), "ClassPath"));
+        for (ObjectName objectName : aliveMBeans) {
+            assertFalse(
+                    objectName + "<>Value should not be excluded dynamically before call to add",
+                    objectNameAttributeFilter.exclude(objectName, "Value"));
+        }
+    }
+
+    private static ObjectNameAttributeFilter initEmptyConfigFilter() {
+        return ObjectNameAttributeFilter.create(
+                new Yaml().load("---\nexcludeObjectNameAttributes: {}\n"));
+    }
+
+    private static ObjectNameAttributeFilter initNonEmptyConfigFilter() {
+        return ObjectNameAttributeFilter.create(
+                new Yaml()
+                        .load(
+                                "---\n"
+                                        + "excludeObjectNameAttributes:\n"
+                                        + "  \"java.lang:type=OperatingSystem\":\n"
+                                        + "    - \"ObjectName\"\n"
+                                        + "  \"java.lang:type=Runtime\":\n"
+                                        + "    - \"ClassPath\"\n"
+                                        + "    - \"SystemProperties\"\n"));
+    }
+
+    private static Set<ObjectName> getAliveMBeans() throws MalformedObjectNameException {
+        return new HashSet<>(
+                Arrays.asList(
+                        new ObjectName("boolean:Type=Test1"),
+                        new ObjectName("boolean:Type=Test2"),
+                        new ObjectName("boolean:Type=Test3"),
+                        new ObjectName("boolean:Type=Test4")));
+    }
+
+    private int excludeMapSize(ObjectNameAttributeFilter filter, String mapName) throws Exception {
+        return getInternalMapValue(filter, mapName).size();
+    }
+
+    private Map getInternalMapValue(ObjectNameAttributeFilter filter, String mapName)
+            throws Exception {
+        return getInternalFieldValue(filter, mapName, Map.class);
+    }
+
+    @SuppressWarnings("java:S1172")
+    private static <T> T getInternalFieldValue(
+            Object object, String fieldName, Class<T> ignoredFieldType) throws Exception {
+        Field field = object.getClass().getDeclaredField(fieldName);
+        field.setAccessible(true);
+        return (T) field.get(object);
+    }
+}