melix / jmh-gradle-plugin

Integrates the JMH benchmarking framework with Gradle
Apache License 2.0
666 stars 88 forks source link

Ability to set environment of the forked JVM #122

Open pkolaczk opened 6 years ago

pkolaczk commented 6 years ago

The plugin allows to set jvm args, but not environment variables. Controlling the environment would be very useful, particularly, when some third-party library referenced by the code under test uses System.getenv and you can't change it to System.getProperty.

No3x commented 5 years ago

I had to populate the environment variables and did the following: Map the env vars from the system to jmh.jvmArgsAppend with -D flag. Then map these arguments to the env vars inside the forked JVM. After this the System.getenv calls work inside the forked JVM.

build.gradle

jmh {
    List<String> args = System.getenv().entrySet().stream().map{es -> String.format("-D%s=%s", es.key, es.value)}.collect(Collectors.toList())
    jvmArgs = ['-Djmh.separateClasspathJAR=true']
    jvmArgsAppend = args
}

MyJMHTest.java

@Setup(Level.Trial)
public synchronized void initialize() {
    EnvironmentHacks.printEnvVars(); // Somehow this is set already, but it's the only one: SystemRoot=C:\WINDOWS
    EnvironmentHacks.mapJVMArgsToEnvVars();
    EnvironmentHacks.printEnvVars(); //Now populated: SystemRoot=C:\WINDOWS TEMP=C:\Users\x\AppData\Local\Temp ....
}

EnvironmentHacks.java

/*
 * Based on: https://blog.sebastian-daschner.com/entries/changing_env_java
 */
public class EnvironmentHacks {

    public static void injectEnvironmentVariable(String key, String value)
            throws Exception {

        Class<?> processEnvironment = Class.forName("java.lang.ProcessEnvironment");

        Field unmodifiableMapField = getAccessibleField(processEnvironment, "theUnmodifiableEnvironment");
        Object unmodifiableMap = unmodifiableMapField.get(null);
        injectIntoUnmodifiableMap(key, value, unmodifiableMap);

        Field mapField = getAccessibleField(processEnvironment, "theEnvironment");
        Map<String, String> map = (Map<String, String>) mapField.get(null);
        map.put(key, value);
    }

    private static Field getAccessibleField(Class<?> clazz, String fieldName)
            throws NoSuchFieldException {

        Field field = clazz.getDeclaredField(fieldName);
        field.setAccessible(true);
        return field;
    }

    private static void injectIntoUnmodifiableMap(String key, String value, Object map)
            throws ReflectiveOperationException {

        Class unmodifiableMap = Class.forName("java.util.Collections$UnmodifiableMap");
        Field field = getAccessibleField(unmodifiableMap, "m");
        Object obj = field.get(map);
        ((Map<String, String>) obj).put(key, value);
    }

    public static void mapJVMArgsToEnvVars() {
        // get the jvm's input arguments as a list of strings
        final List<String> inputArguments = ManagementFactory.getRuntimeMXBean().getInputArguments();

        for (String arg : inputArguments) {
            final String key = arg.split("=")[0];
            final String value = arg.split("=")[1];
            final String envKey = key.substring(2);
            if(key.startsWith("-D") && !envKey.isEmpty() ) {
                try {
                    injectEnvironmentVariable(envKey, value);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }

    public static void printEnvVars() {
        System.out.println("ENV vars:");
        Map<String, String> env = System.getenv();
        for (String envName : env.keySet()) {
            System.out.format("%s=%s%n", envName, env.get(envName));
        }
    }
}

Please consider this as workaround. I'm not familiar with gradle and groovy to integrate this properly.