Skip to content

Commit

Permalink
Fix possible memory leak in ObjectNameAttributeFilter
Browse files Browse the repository at this point in the history
Signed-off-by: Pawel Kwasny <pwkwasny@gmail.com>
  • Loading branch information
kw4s committed Jun 21, 2024
1 parent 77bd2a4 commit 8af30cf
Show file tree
Hide file tree
Showing 3 changed files with 196 additions and 31 deletions.
23 changes: 5 additions & 18 deletions collector/src/main/java/io/prometheus/jmx/JmxScraper.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<>();
}

/**
Expand All @@ -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);
}
}

Expand All @@ -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(
Expand All @@ -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
*
Expand All @@ -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;
}

Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}

0 comments on commit 8af30cf

Please sign in to comment.