Support use of snakeyaml-engine on java 8+ and legacy snakeyaml on Java 7

pom.xml

@@ -61,6 +61,9 @@
         <lombok.version>1.18.38</lombok.version>
         <mockito.version>2.28.2</mockito.version>
         <slf4j.version>1.7.32</slf4j.version>
+        <!-- used on Java 8+ when available on runtime classpath -->
+        <snakeyaml-engine.version>2.9</snakeyaml-engine.version>
+        <!-- used on Java 7 or if snakeyaml-engine is not available -->
         <snakeyaml.version>2.0</snakeyaml.version>
         <testcontainers.version>1.19.8</testcontainers.version>
 
@@ -215,6 +218,13 @@
             <artifactId>snakeyaml</artifactId>
             <version>${snakeyaml.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.snakeyaml</groupId>
+            <artifactId>snakeyaml-engine</artifactId>
+            <version>${snakeyaml-engine.version}</version>
+            <scope>provided</scope>
+            <optional>true</optional>
+        </dependency>
         <dependency>
             <groupId>javax.management.j2ee</groupId>
             <artifactId>javax.management.j2ee-api</artifactId>

src/main/java/org/datadog/jmxfetch/App.java

@@ -84,8 +84,8 @@ public class App {
     private static int loopCounter;
     private int lastJsonConfigTs;
     private Map<String, Object> adJsonConfigs;
-    private Map<String, YamlParser> configs;
-    private Map<String, YamlParser> adPipeConfigs = new ConcurrentHashMap<>();
+    private Map<String, ConfigYaml> configs;
+    private Map<String, ConfigYaml> adPipeConfigs = new ConcurrentHashMap<>();
     private List<Instance> instances = new ArrayList<>();
     private Map<String, Instance> brokenInstanceMap = new ConcurrentHashMap<>();
     private AtomicBoolean reinit = new AtomicBoolean(false);
@@ -403,7 +403,7 @@ public boolean processAutoDiscovery(byte[] buffer) {
             final String name = getAutoDiscoveryName(config);
             log.debug("Attempting to apply config. Name: " + name);
             final InputStream stream = new ByteArrayInputStream(config.getBytes(UTF_8));
-            final YamlParser yaml = new YamlParser(stream);
+            final ConfigYaml yaml = new ConfigYaml(stream);
 
             if (this.addConfig(name, yaml)) {
                 reinit = true;
@@ -700,7 +700,7 @@ public TaskStatusHandler invoke(
      * Adds a configuration to the auto-discovery pipe-collected configuration list. This method is
      * deprecated.
      */
-    public boolean addConfig(final String name, final YamlParser config) {
+    public boolean addConfig(final String name, final ConfigYaml config) {
         // named groups not supported with Java6:
         //
         // "AUTO_DISCOVERY_PREFIX(?<check>.{1,80})_(?<version>\\d{0,AD_MAX_MAG_INSTANCES})"
@@ -746,8 +746,8 @@ public boolean addJsonConfig(final String name, final String json) {
         return false;
     }
 
-    private Map<String, YamlParser> getConfigs(final AppConfig config) {
-        final Map<String, YamlParser> configs = new ConcurrentHashMap<>();
+    private Map<String, ConfigYaml> getConfigs(final AppConfig config) {
+        final Map<String, ConfigYaml> configs = new ConcurrentHashMap<>();
 
         this.loadFileConfigs(config, configs);
         this.loadResourceConfigs(config, configs);
@@ -756,7 +756,7 @@ private Map<String, YamlParser> getConfigs(final AppConfig config) {
         return configs;
     }
 
-    private void loadFileConfigs(final AppConfig config, final Map<String, YamlParser> configs) {
+    private void loadFileConfigs(final AppConfig config, final Map<String, ConfigYaml> configs) {
         final List<String> fileList = config.getYamlFileList();
         if (fileList != null) {
             for (final String fileName : fileList) {
@@ -765,7 +765,7 @@ private void loadFileConfigs(final AppConfig config, final Map<String, YamlParse
                 final String yamlPath = file.getAbsolutePath();
                 log.info("Reading {}", yamlPath);
                 try (FileInputStream yamlInputStream = new FileInputStream(yamlPath)) {
-                    configs.put(name, new YamlParser(yamlInputStream));
+                    configs.put(name, new ConfigYaml(yamlInputStream));
                 } catch (FileNotFoundException e) {
                     log.warn("Cannot find " + yamlPath);
                 } catch (Exception e) {
@@ -776,7 +776,7 @@ private void loadFileConfigs(final AppConfig config, final Map<String, YamlParse
     }
 
     private void loadResourceConfigs(
-            final AppConfig config, final Map<String, YamlParser> configs) {
+            final AppConfig config, final Map<String, ConfigYaml> configs) {
         final List<String> resourceConfigList = config.getInstanceConfigResources();
         if (resourceConfigList != null) {
             final ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
@@ -788,7 +788,7 @@ private void loadResourceConfigs(
                     log.warn("Cannot find " + resourceName);
                 } else {
                     try {
-                        configs.put(name, new YamlParser(inputStream));
+                        configs.put(name, new ConfigYaml(inputStream));
                     } catch (Exception e) {
                         log.warn("Cannot parse yaml file " + resourceName, e);
                     } finally {
@@ -933,11 +933,11 @@ public void init(final boolean forceNewConnection) {
         final Set<String> instanceNamesSeen = new HashSet<>();
 
         log.info("Dealing with YAML config instances...");
-        final Iterator<Entry<String, YamlParser>> it = this.configs.entrySet().iterator();
-        final Iterator<Entry<String, YamlParser>> itPipeConfigs = this.adPipeConfigs
+        final Iterator<Entry<String, ConfigYaml>> it = this.configs.entrySet().iterator();
+        final Iterator<Entry<String, ConfigYaml>> itPipeConfigs = this.adPipeConfigs
                         .entrySet().iterator();
         while (it.hasNext() || itPipeConfigs.hasNext()) {
-            Map.Entry<String, YamlParser> entry;
+            Map.Entry<String, ConfigYaml> entry;
             boolean fromPipeIterator = false;
             if (it.hasNext()) {
                 entry = it.next();
@@ -947,14 +947,14 @@ public void init(final boolean forceNewConnection) {
             }
 
             final String name = entry.getKey();
-            final YamlParser yamlConfig = entry.getValue();
+            final ConfigYaml yamlConfig = entry.getValue();
             // AD config cache doesn't remove configs - it just overwrites.
             if (!fromPipeIterator) {
                 it.remove();
             }
 
             final List<Map<String, Object>> configInstances =
-                    ((List<Map<String, Object>>) yamlConfig.getYamlInstances());
+                    ((List<Map<String, Object>>) yamlConfig.getInstances());
             if (configInstances == null || configInstances.size() == 0) {
                 final String warning = "No instance found in :" + name;
                 log.warn(warning);

src/main/java/org/datadog/jmxfetch/ConfigYaml.java

@@ -0,0 +1,96 @@
+package org.datadog.jmxfetch;
+
+import org.datadog.jmxfetch.util.JavaVersion;
+import org.snakeyaml.engine.v2.api.Dump;
+import org.snakeyaml.engine.v2.api.DumpSettings;
+import org.snakeyaml.engine.v2.api.Load;
+import org.snakeyaml.engine.v2.api.LoadSettings;
+import org.yaml.snakeyaml.Yaml;
+
+import java.io.InputStream;
+import java.util.Map;
+
+class ConfigYaml {
+    private static final boolean USE_YAML_ENGINE;
+
+    // used on Java 8+
+    private static final ThreadLocal<Load> YAML_LOAD;
+    private static final ThreadLocal<Dump> YAML_DUMP;
+
+    // used on Java 7
+    private static final ThreadLocal<Yaml> YAML_LEGACY;
+
+    static {
+        ThreadLocal<Load> yamlLoad = null;
+        ThreadLocal<Dump> yamlDump = null;
+
+        if (JavaVersion.atLeastJava(8)) {
+            try {
+                final LoadSettings loadSettings = LoadSettings.builder().build();
+                yamlLoad = new ThreadLocal<Load>() {
+                    @Override
+                    public Load initialValue() {
+                        return new Load(loadSettings);
+                    }
+                };
+                final DumpSettings dumpSettings = DumpSettings.builder().build();
+                yamlDump = new ThreadLocal<Dump>() {
+                    @Override
+                    public Dump initialValue() {
+                        return new Dump(dumpSettings);
+                    }
+                };
+            } catch (Throwable ignore) {
+                // snakeyaml-engine not available, fallback to legacy snakeyaml
+            }
+        }
+
+        USE_YAML_ENGINE = yamlLoad != null && yamlDump != null;
+
+        if (USE_YAML_ENGINE) {
+            YAML_LOAD = yamlLoad;
+            YAML_DUMP = yamlDump;
+            YAML_LEGACY = null;
+        } else {
+            YAML_LOAD = null;
+            YAML_DUMP = null;
+            YAML_LEGACY = new ThreadLocal<Yaml>() {
+                @Override
+                public Yaml initialValue() {
+                    return new Yaml();
+                }
+            };
+        }
+    }
+
+    private final Map<Object, Object> parsedYaml;
+
+    public ConfigYaml(InputStream yamlInputStream) {
+        parsedYaml = parse(yamlInputStream);
+    }
+
+    public Object getInstances() {
+        return parsedYaml.get("instances");
+    }
+
+    public Object getInitConfig() {
+        return parsedYaml.get("init_config");
+    }
+
+    @SuppressWarnings("unchecked")
+    public static <T> T parse(InputStream yamlInputStream) {
+        if (USE_YAML_ENGINE) {
+            return (T) YAML_LOAD.get().loadFromInputStream(yamlInputStream);
+        } else {
+            return YAML_LEGACY.get().load(yamlInputStream);
+        }
+    }
+
+    public static String dump(Object parsedYaml) {
+        if (USE_YAML_ENGINE) {
+            return YAML_DUMP.get().dumpToString(parsedYaml);
+        } else {
+            return YAML_LEGACY.get().dump(parsedYaml);
+        }
+    }
+}

src/main/java/org/datadog/jmxfetch/Instance.java

@@ -5,9 +5,6 @@
 import org.datadog.jmxfetch.service.ConfigServiceNameProvider;
 import org.datadog.jmxfetch.service.ServiceNameProvider;
 import org.datadog.jmxfetch.util.InstanceTelemetry;
-import org.yaml.snakeyaml.Yaml;
-
-
 
 import java.io.File;
 import java.io.FileInputStream;
@@ -43,14 +40,6 @@ public class Instance {
     public static final String JVM_DIRECT = "jvm_direct";
     public static final String ATTRIBUTE = "Attribute: ";
 
-    private static final ThreadLocal<Yaml> YAML =
-        new ThreadLocal<Yaml>() {
-            @Override
-            public Yaml initialValue() {
-                return new Yaml();
-            }
-        };
-
     private Set<ObjectName> beans;
     private List<String> beanScopes;
     private List<Configuration> configurationList = new ArrayList<Configuration>();
@@ -297,8 +286,7 @@ public static boolean isDirectInstance(Map<String, Object> configInstance) {
 
     private void loadDefaultConfig(String configResourcePath) {
         List<Map<String, Object>> defaultConf =
-                (List<Map<String, Object>>)
-                        YAML.get().load(this.getClass().getResourceAsStream(configResourcePath));
+                ConfigYaml.parse(this.getClass().getResourceAsStream(configResourcePath));
         for (Map<String, Object> conf : defaultConf) {
             configurationList.add(new Configuration(conf));
         }
@@ -321,8 +309,7 @@ static void loadMetricConfigFiles(
                 try {
                     yamlInputStream = new FileInputStream(yamlPath);
                     List<Map<String, Object>> confs =
-                            (List<Map<String, Object>>)
-                                    YAML.get().load(yamlInputStream);
+                            ConfigYaml.parse(yamlInputStream);
                     for (Map<String, Object> conf : confs) {
                         configurationList.add(new Configuration(conf));
                     }
@@ -360,8 +347,7 @@ static void loadMetricConfigResources(
                 } else {
                     try {
                         Map<String, List<Map<String, Object>>> topYaml =
-                                (Map<String, List<Map<String, Object>>>)
-                                        YAML.get().load(inputStream);
+                                ConfigYaml.parse(inputStream);
                         List<Map<String, Object>> jmxConf =
                                 topYaml.get("jmx_metrics");
                         if (jmxConf != null) {

src/main/java/org/datadog/jmxfetch/Status.java

@@ -4,7 +4,6 @@
 
 import org.datadog.jmxfetch.util.InstanceTelemetry;
 import org.datadog.jmxfetch.util.MetadataHelper;
-import org.yaml.snakeyaml.Yaml;
 
 import java.io.BufferedOutputStream;
 import java.io.File;
@@ -151,7 +150,7 @@ private String generateYaml() {
         status.put("timestamp", System.currentTimeMillis());
         status.put("checks", this.instanceStats);
         status.put("errors", this.errors);
-        return new Yaml().dump(status);
+        return ConfigYaml.dump(status);
     }
 
     private String generateJson() throws IOException {

src/main/java/org/datadog/jmxfetch/util/JavaVersion.java

@@ -0,0 +1,44 @@
+package org.datadog.jmxfetch.util;
+
+public final class JavaVersion {
+    private static final int JAVA_VERSION = getMajorJavaVersion();
+
+    private JavaVersion() {}
+
+    /**
+     * Tests whether the Java version is at least the expected version.
+     *
+     * @param expectedVersion the expected version
+     * @return {@code true} if the Java version is at least the expected version
+     */
+    public static boolean atLeastJava(int expectedVersion) {
+        return expectedVersion <= JAVA_VERSION;
+    }
+
+    private static int getMajorJavaVersion() {
+        try {
+            return parseMajorJavaVersion(System.getProperty("java.version"));
+        } catch (Throwable e) {
+            return 7; // assume Java 7, i.e. the lowest version JMXFetch supports
+        }
+    }
+
+    private static int parseMajorJavaVersion(String str) {
+        int value = 0;
+        for (int i = 0; i < str.length(); i++) {
+            char ch = str.charAt(i);
+            if (ch >= '0' && ch <= '9') {
+                value = value * 10 + (ch - '0');
+            } else if (ch == '.') {
+                if (value == 1) {
+                    value = 0; // skip leading 1.
+                } else {
+                    break;
+                }
+            } else {
+                throw new NumberFormatException();
+            }
+        }
+        return value;
+    }
+}