Core Affinity Agent (#8)

* Kotlin coroutines pinned dispatchers functionality

* Core shielding initial commit

* Initial work on shielding agent

* Basic thread shielding functionality seems to be working for bootstrapping as well

* Done with basic agent configuration support, allowing the specification of various affinity configurations

* more core re-organizaition to allow for agent dynamic loading and testing

* Done with agent shielding tests

* gradle test fixes

* Some semantic refactoring
Listener now uses JUL
Javadoc fixes

* More javadoc stuff

* Kotlin Formatting fixes

* Kotlin Formatting fixes

* Needle affinity annotations for needle-agent exclusion

* agent exclusion from codecov.yml

* agent tests enhancements

* Removed unneeded prints

* Done with agent code cleanup

* typo fix

* enable maven publishing

* gradle fixes

* gradle fixes

agent/build.gradle

@@ -0,0 +1,124 @@
+plugins {
+    id 'com.github.johnrengelman.shadow' version '6.0.0'
+}
+
+dependencies {
+
+    implementation project(':needle-core')
+    testImplementation project(':needle-core').sourceSets.test.output
+
+    implementation 'net.bytebuddy:byte-buddy:1.10.14'
+
+    implementation 'com.fasterxml.jackson.core:jackson-core:2.11.2'
+    implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.11.2'
+}
+
+task sourcesJar(type: Jar, dependsOn: classes) {
+    classifier 'sources'
+    from sourceSets.main.allSource
+}
+
+task javadocJar(type: Jar, dependsOn: javadoc) {
+    classifier 'javadoc'
+    from javadoc.destinationDir
+}
+
+jar {
+    enabled = false
+}
+
+afterEvaluate {
+
+    task copyAgent(dependsOn: shadowJar, type: Copy) {
+        from "$buildDir/libs/"
+        into "$buildDir/resources/test"
+        include "*.jar"
+    }
+
+    shadowJar {
+
+        archiveClassifier.set('')
+
+        from sourceSets.main.output
+
+        dependencies {
+            exclude(dependency('org.jetbrains.kotlin:.*:.*'))
+        }
+
+        manifest {
+            attributes(
+                    "Agent-Class": "org.sheinbergon.needle.agent.NeedleAgent",
+                    "Can-Redefine-Classes": true,
+                    "Can-Retransform-Classes": true,
+                    "Boot-Class-Path": archiveFileName.get(),
+                    "Premain-Class": "org.sheinbergon.needle.agent.NeedleAgent")
+        }
+
+        finalizedBy tasks.copyAgent
+    }
+
+    test {
+        forkEvery(1)
+        dependsOn tasks.copyAgent
+        systemProperties = [
+                "jdk.attach.allowAttachSelf": true,
+                "test.agent.jar.path"       : "/${tasks.shadowJar.archiveFileName.get()}"
+        ]
+    }
+
+    publish.dependsOn shadowJar
+    assemble.dependsOn shadowJar
+    build.dependsOn shadowJar
+    uploadArchives.dependsOn shadowJar
+}
+
+publishing {
+    publications {
+        agent(MavenPublication) {
+            project.shadow.component(it)
+            artifactId = "needle-agent"
+            artifact tasks.javadocJar
+            artifact tasks.sourcesJar
+            pom {
+                name = project.name
+                description = 'Feature-rich CPU affinity for the JVM - Affinity Setting Agent'
+                url = 'https://github.com/sheinbergon/needle'
+                inceptionYear = '2020'
+
+                licenses {
+                    license {
+                        name = 'Apache License 2.0'
+                        url = 'https://github.com/sheinbergon/needle/blob/master/LICENSE'
+                        distribution = 'repo'
+                    }
+                }
+
+                developers {
+                    developer {
+                        id = 'sheinbergon'
+                        name = 'Idan Sheinberg'
+                        email = '[email protected]'
+                    }
+                }
+
+                scm {
+                    url = 'https://github.com/sheinbergon/needle'
+                    connection = 'scm:https://github.com/sheinbergon/needle.git'
+                    developerConnection = 'scm:[email protected]:sheinbergon/needle.git'
+                }
+            }
+
+            repositories {
+                mavenLocal()
+                maven {
+                    name "oss-sonatype-nexus"
+                    url nexus.url
+                    credentials {
+                        username = nexus.username
+                        password = nexus.password
+                    }
+                }
+            }
+        }
+    }
+}

agent/src/main/java/org/sheinbergon/needle/agent/AffinityAdvice.java

@@ -0,0 +1,54 @@
+package org.sheinbergon.needle.agent;
+
+import lombok.val;
+import net.bytebuddy.asm.Advice;
+import org.sheinbergon.needle.Needle;
+import org.sheinbergon.needle.Pinned;
+import org.sheinbergon.needle.agent.util.AffinityGroupMatcher;
+import org.sheinbergon.needle.util.NeedleAffinity;
+
+import javax.annotation.Nonnull;
+
+public final class AffinityAdvice {
+
+    /**
+     * Empty utility class private constructor.
+     */
+    private AffinityAdvice() {
+    }
+
+    /**
+     * Byte-Buddy advice method, this is latched on to {@link Thread#run()}, running before the actual method executes.
+     * <p>
+     * Simply put, it search for a matching affinity group, and sets the thread's affintiy according to the group's
+     * {@link org.sheinbergon.needle.AffinityDescriptor} specified values.
+     *
+     * @see NeedleAgentConfiguration
+     * @see org.sheinbergon.needle.AffinityDescriptor
+     */
+    @Advice.OnMethodEnter
+    public static void run() {
+        try {
+            val thread = Thread.currentThread();
+            if (!excluded(thread)) {
+                val group = AffinityGroupMatcher.forThread(thread);
+                Needle.affinity(group.affinity());
+            }
+        } catch (Throwable throwable) {
+            // Do nothing if any exception were thrown, for now.
+            // TODO - Expose these via JUL logging (meant to bridged to other loggers)
+        }
+    }
+
+    /**
+     * Indicates whether this given thread should be exempt from affinity-group matching and setting.
+     * This method is public to satisfy byte-buddy class-loading requirements.
+     *
+     * @param thread The thread to inspect for exclusion from affinity-group matching.
+     * @return Boolean value, indicating if this thread is to be applied affinity-group settings
+     */
+    public static boolean excluded(final @Nonnull Thread thread) {
+        val type = thread.getClass();
+        return Pinned.class.isAssignableFrom(type) || type.isAnnotationPresent(NeedleAffinity.class);
+    }
+}

agent/src/main/java/org/sheinbergon/needle/agent/NeedleAgent.java

@@ -0,0 +1,129 @@
+package org.sheinbergon.needle.agent;
+
+import lombok.val;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.dynamic.loading.ClassInjector;
+import net.bytebuddy.utility.JavaModule;
+import org.sheinbergon.needle.Pinned;
+import org.sheinbergon.needle.agent.util.AffinityGroupMatcher;
+import org.sheinbergon.needle.agent.util.YamlCodec;
+import org.sheinbergon.needle.util.NeedleAffinity;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.io.File;
+import java.lang.instrument.Instrumentation;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.nio.file.Files;
+import java.util.Map;
+import java.util.function.Supplier;
+
+import static net.bytebuddy.matcher.ElementMatchers.is;
+import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+public final class NeedleAgent {
+
+    private NeedleAgent() {
+    }
+
+    /**
+     * Static agent loading endpoint.
+     *
+     * @param arguments       Agent configuration string, if specified, must be a valid JVM URL string pointing to the
+     *                        agent configuration file path (i.e. file:///some/file.yml).
+     * @param instrumentation JVM instrumentation interface
+     * @throws Exception any execution error encountered during instrumentation setup
+     */
+    public static void premain(
+            final String arguments,
+            final Instrumentation instrumentation) throws Exception {
+        val storage = Files.createTempDirectory("needle-agent-instrumentation").toFile();
+        setupBootstrapInjection(storage, instrumentation);
+        agentConfiguration(arguments);
+        val builder = new AgentBuilder.Default()
+                .disableClassFormatChanges()
+                .ignore(nameStartsWith("net.bytebuddy."))
+                .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
+                .with(AgentBuilder.RedefinitionStrategy.REDEFINITION)
+                .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
+                .with(new AgentBuilder.InjectionStrategy.UsingInstrumentation(instrumentation, storage));
+        val narrowable = matchers(builder);
+        narrowable.transform(NeedleAgent::premainTransform)
+                .installOn(instrumentation);
+    }
+
+    /**
+     * Dynamic agent loading endpoint.
+     *
+     * @param arguments       Agent configuration string, if specified, must be a valid JVM URL string pointing to the
+     *                        agent configuration file path (i.e. file:///some/file.yml).
+     * @param instrumentation JVM instrumentation interface
+     * @throws Exception any execution error encountered during instrumentation setup
+     */
+    public static void agentmain(
+            final String arguments,
+            final Instrumentation instrumentation) throws Exception {
+        agentConfiguration(arguments);
+        val builder = new AgentBuilder.Default()
+                .disableClassFormatChanges()
+                .ignore(nameStartsWith("net.bytebuddy."))
+                .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
+                .with(AgentBuilder.RedefinitionStrategy.REDEFINITION)
+                .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE);
+        val narrowable = matchers(builder);
+        narrowable.transform(agentmainTransform())
+                .installOn(instrumentation);
+
+    }
+
+    private static AgentBuilder.Identified.Narrowable matchers(
+            final @Nonnull AgentBuilder builder) {
+        return builder.type(not(isSubTypeOf(Pinned.class))
+                .and(not(isAnnotatedWith(NeedleAffinity.class)))
+                .and(isSubTypeOf(Thread.class).or(is(Thread.class))));
+    }
+
+    private static DynamicType.Builder<?> premainTransform(
+            final @Nonnull DynamicType.Builder<?> builder,
+            final @Nonnull TypeDescription typeDescription,
+            final @Nullable ClassLoader classLoader,
+            final @Nonnull JavaModule module) {
+        return builder.visit(Advice.to(AffinityAdvice.class).on(named("run")));
+    }
+
+    private static AgentBuilder.Transformer agentmainTransform() {
+        return new AgentBuilder.Transformer.ForAdvice()
+                .include(NeedleAgent.class.getClassLoader())
+                .advice(named("run"), AffinityAdvice.class.getName());
+    }
+
+    private static void setupBootstrapInjection(
+            final @Nonnull File storage,
+            final @Nonnull Instrumentation instrumentation) {
+        ClassInjector.UsingInstrumentation
+                .of(storage, ClassInjector.UsingInstrumentation.Target.BOOTSTRAP, instrumentation)
+                .inject(Map.of(
+                        new TypeDescription.ForLoadedType(AffinityAdvice.class),
+                        ClassFileLocator.ForClassLoader.read(AffinityAdvice.class)));
+    }
+
+    private static void agentConfiguration(final @Nullable String arguments) throws MalformedURLException {
+        Supplier<NeedleAgentConfiguration> supplier;
+        if (arguments != null) {
+            val url = new URL(arguments);
+            supplier = () -> YamlCodec.parseConfiguration(url);
+        } else {
+            supplier = () -> NeedleAgentConfiguration.DEFAULT;
+        }
+        AffinityGroupMatcher.setConfigurationSupplier(supplier);
+    }
+}