Add type pollution test harness (#1048)

Add a test harness that can be used to check for type pollution issues. Heavily inspired by and similar to https://github.com/RedHatPerf/type-pollution-agent . Key differences are:

- Works with in-VM ByteBuddyAgent. This makes installation somewhat easier in tests
- Fully programmatic API, to make threshold verification work without parsing summary reports. Also means this cannot be used for existing applications, it needs app code that explicitly listens to events
- indy-based implementation instead of invokestatic
- Less focus on performance, this is for testing only
- No type check miss support

gradle/libs.versions.toml

@@ -12,6 +12,7 @@ managed-junit = "5.10.3"
 managed-rest-assured = "5.4.0"
 managed-kotest = "5.9.1"
 managed-spock = "2.3-groovy-4.0"
+managed-bytebuddy = "1.14.17"
 
 kotlin = "1.9.24"
 graal-svm = "23.1.3"
@@ -58,6 +59,9 @@ managed-junit-jupiter-params = { module = "org.junit.jupiter:junit-jupiter-param
 managed-junit-jupiter-api = { module = "org.junit.jupiter:junit-jupiter-api", version.ref = "managed-junit" }
 managed-junit-vintage-engine = { module = "org.junit.vintage:junit-vintage-engine", version.ref = "managed-junit" }
 
+managed-bytebuddy = { module = "net.bytebuddy:byte-buddy", version.ref = "managed-bytebuddy" }
+managed-bytebuddy-agent = { module = "net.bytebuddy:byte-buddy-agent", version.ref = "managed-bytebuddy" }
+
 # BOMs
 boms-junit = { module = "org.junit:junit-bom", version.ref = "managed-junit" }
 boms-kotest = { module = "io.kotest:kotest-bom", version.ref = "managed-kotest" }

settings.gradle

@@ -34,3 +34,4 @@ include "test-kotest5"
 include "test-rest-assured"
 include 'test-suite-sql-r2dbc'
 include 'test-suite-at-sql-jpa'
+include 'test-type-pollution'

test-type-pollution/build.gradle.kts

@@ -0,0 +1,30 @@
+plugins {
+    id("io.micronaut.build.internal.micronaut-test-module")
+}
+
+repositories {
+    mavenCentral()
+}
+
+dependencies {
+    implementation(libs.managed.bytebuddy)
+
+    testImplementation(libs.managed.bytebuddy.agent)
+    testImplementation(libs.managed.junit.jupiter.api)
+    testRuntimeOnly(libs.managed.junit.jupiter.engine)
+}
+
+tasks.withType<Test> {
+    jvmArgs("-XX:+EnableDynamicAgentLoading")
+
+    jacoco {
+        enabled = false
+    }
+}
+
+micronautBuild {
+    // todo: enable after 4.4.0
+    binaryCompatibility {
+        enabled.set(false)
+    }
+}

test-type-pollution/src/main/java/io/micronaut/test/typepollution/ClassSwitch.java

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2017-2024 original authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.micronaut.test.typepollution;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.lang.invoke.MutableCallSite;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * This class dynamically builds and maintains a {@link MethodHandle} that calls any of a number of
+ * downstream {@link MethodHandle}s, depending on a {@link Class} parameter. If a new {@link Class}
+ * is passed to this {@link MethodHandle}, it is mutated to gain a branch for that {@link Class} to
+ * enable fast dispatch in the future.
+ */
+abstract class ClassSwitch {
+    private static final MethodHandle FALLBACK;
+    private static final MethodHandle CLASS_EQUALS;
+
+    static {
+        try {
+            MethodHandles.Lookup lookup = MethodHandles.lookup();
+            FALLBACK = lookup.findVirtual(ClassSwitch.class, "fallback", MethodType.methodType(void.class, Class.class));
+            CLASS_EQUALS = MethodHandles.explicitCastArguments(
+                lookup.findVirtual(Object.class, "equals", MethodType.methodType(boolean.class, Object.class)),
+                MethodType.methodType(boolean.class, Class.class, Class.class)
+            );
+        } catch (NoSuchMethodException | IllegalAccessException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private final MutableCallSite callSite;
+    private final int switchPosition;
+
+    /**
+     * @param expectedType The type of the downstream method returned by {@link #downstream(Class)}
+     * @param switchPosition The index in the parameter list where the {@link Class} parameter
+     *                       should be added that this {@link ClassSwitch} will switch on
+     */
+    ClassSwitch(MethodType expectedType, int switchPosition) {
+        List<Class<?>> combinedArgs = new ArrayList<>(expectedType.parameterList());
+        combinedArgs.add(switchPosition, Class.class);
+        this.callSite = new MutableCallSite(
+            MethodHandles.dropArguments(
+                MethodHandles.throwException(expectedType.returnType(), UnsupportedOperationException.class).bindTo(new UnsupportedOperationException("Type not bound yet")),
+                0,
+                combinedArgs
+            )
+        );
+        this.switchPosition = switchPosition;
+
+        // replace the call site with a handle that first calls fallback and then calls the call
+        // site again. The fallback call will update the call site to include the new type.
+        callSite.setTarget(MethodHandles.foldArguments(
+            dynamicInvoker(),
+            switchPosition,
+            FALLBACK.bindTo(this)
+        ));
+    }
+
+    /**
+     * Build a {@link MethodHandle} that invokes this switch statement.
+     */
+    public final MethodHandle dynamicInvoker() {
+        return callSite.dynamicInvoker();
+    }
+
+    private synchronized void fallback(Class<?> cl) {
+        MethodHandle test = CLASS_EQUALS.bindTo(cl);
+        if (switchPosition != 0) {
+            test = MethodHandles.dropArguments(test, 0, callSite.getTarget().type().parameterList().subList(0, switchPosition));
+        }
+        callSite.setTarget(MethodHandles.guardWithTest(
+            test,
+            MethodHandles.dropArguments(downstream(cl), switchPosition, List.of(Class.class)),
+            callSite.getTarget()
+        ));
+    }
+
+    /**
+     * Construct a new downstream {@link MethodHandle} for the given type.
+     *
+     * @param cl The type that was passed to this switch statement
+     * @return The method handle that will be called every time this type is seen
+     */
+    protected abstract MethodHandle downstream(Class<?> cl);
+}

test-type-pollution/src/main/java/io/micronaut/test/typepollution/ConcreteCounter.java

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2017-2024 original authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.micronaut.test.typepollution;
+
+import io.micronaut.core.annotation.Nullable;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.lang.invoke.MutableCallSite;
+
+/**
+ * This class recognizes the "focus changes" that are then forwarded to the {@link FocusListener}.
+ * <p>
+ * There is one {@link ConcreteCounter} for each concrete type. Each {@link ConcreteCounter} has an
+ * {@link InterfaceHolder} for any interfaces that the concrete type is type checked against. Each
+ * {@link InterfaceHolder} has a {@link MutableCallSite}. For the interface that was last type
+ * checked, this call site does nothing. For any other interface, a type check leads to a "focus
+ * event": The event is forwarded to the {@link FocusListener} for logging. The now-focused
+ * interface changes its call site to do nothing, and the previously focused interface becomes
+ * unfocused.
+ */
+final class ConcreteCounter {
+    static final ClassValue<ConcreteCounter> COUNTERS = new ClassValue<>() {
+        @Override
+        protected ConcreteCounter computeValue(Class<?> type) {
+            return new ConcreteCounter(type);
+        }
+    };
+
+    @Nullable
+    static FocusListener focusListener;
+
+    private static final MethodHandle FOCUS;
+    private static final MethodHandle IGNORE = MethodHandles.empty(MethodType.methodType(void.class));
+
+    static {
+        try {
+            FOCUS = MethodHandles.lookup()
+                .findVirtual(InterfaceHolder.class, "focus", MethodType.methodType(void.class));
+        } catch (NoSuchMethodException | IllegalAccessException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private final Class<?> concreteType;
+
+    private final ClassValue<InterfaceHolder> holdersByInterface = new ClassValue<>() {
+        @Override
+        protected InterfaceHolder computeValue(Class<?> type) {
+            return new InterfaceHolder(type);
+        }
+    };
+    private InterfaceHolder focused = null;
+
+    private ConcreteCounter(Class<?> concreteType) {
+        this.concreteType = concreteType;
+    }
+
+    MethodHandle typeCheckHandle(Class<?> interfaceType) {
+        return holdersByInterface.get(interfaceType).callSite.dynamicInvoker();
+    }
+
+    final class InterfaceHolder {
+        final Class<?> interfaceType;
+        final MethodHandle focus = FOCUS.bindTo(this);
+        final MutableCallSite callSite = new MutableCallSite(focus);
+
+        InterfaceHolder(Class<?> interfaceType) {
+            this.interfaceType = interfaceType;
+        }
+
+        private void unfocus() {
+            assert Thread.holdsLock(ConcreteCounter.this);
+            callSite.setTarget(focus);
+        }
+
+        private void focus() {
+            FocusListener focusListener = ConcreteCounter.focusListener;
+            if (focusListener != null) {
+                focusListener.onFocus(concreteType, interfaceType);
+            }
+            synchronized (ConcreteCounter.this) {
+                if (focused != null) {
+                    focused.unfocus();
+                }
+                focused = this;
+                callSite.setTarget(IGNORE);
+            }
+        }
+    }
+}

test-type-pollution/src/main/java/io/micronaut/test/typepollution/FocusListener.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2017-2024 original authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.micronaut.test.typepollution;
+
+import io.micronaut.core.annotation.Nullable;
+
+/**
+ * Public listener for responding to focus events.<p>
+ * A focus event happens when a certain concrete type is successfully type checked against an
+ * interface that it was not checked against immediately before:
+ * <code><pre>
+ * Impl impl = new Impl();
+ * impl instanceof A // focus event
+ * impl instanceof A // no focus event
+ * impl instanceof A // no focus event
+ * impl instanceof B // focus event
+ * impl instanceof B // no focus event
+ * impl instanceof A // focus event
+ * </pre></code>
+ * Each focus event may invalidate a cache field on the concrete class which can be especially
+ * expensive on machines with many cores (JDK-8180450). Thus, such focus events should be kept off
+ * the hot path when running on JDK versions that still have this bug.
+ */
+public interface FocusListener {
+    /**
+     * Set the global focus listener, or {@code null} to disable listening for focus events.
+     *
+     * @param focusListener The focus listener
+     */
+    static void setFocusListener(@Nullable FocusListener focusListener) {
+        ConcreteCounter.focusListener = focusListener;
+    }
+
+    /**
+     * Called on every focus event, potentially concurrently.
+     *
+     * @param concreteType  The concrete type that was checked
+     * @param interfaceType The interface type that the concrete type was type checked against
+     */
+    void onFocus(Class<?> concreteType, Class<?> interfaceType);
+}