Rewrite InstrumentPlugin to be lazy

buildSrc/src/main/groovy/InstrumentPlugin.groovy

@@ -1,19 +1,16 @@
-import org.gradle.api.DefaultTask
+import org.gradle.api.Action
 import org.gradle.api.Plugin
 import org.gradle.api.Project
 import org.gradle.api.file.ConfigurableFileCollection
-import org.gradle.api.file.Directory
 import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.file.FileCollection
 import org.gradle.api.invocation.BuildInvocationDetails
+import org.gradle.api.logging.Logger
+import org.gradle.api.logging.Logging
 import org.gradle.api.provider.ListProperty
 import org.gradle.api.provider.Property
-import org.gradle.api.tasks.Classpath
-import org.gradle.api.tasks.Input
-import org.gradle.api.tasks.InputDirectory
-import org.gradle.api.tasks.InputFiles
-import org.gradle.api.tasks.Optional
-import org.gradle.api.tasks.OutputDirectory
-import org.gradle.api.tasks.TaskAction
+import org.gradle.api.tasks.SourceSet
+import org.gradle.api.tasks.SourceSetContainer
 import org.gradle.api.tasks.compile.AbstractCompile
 import org.gradle.jvm.toolchain.JavaLanguageVersion
 import org.gradle.jvm.toolchain.JavaToolchainService
@@ -22,6 +19,9 @@ import org.gradle.workers.WorkParameters
 import org.gradle.workers.WorkerExecutor
 
 import javax.inject.Inject
+import java.nio.file.Files
+import java.nio.file.Path
+import java.nio.file.StandardCopyOption
 import java.util.concurrent.ConcurrentHashMap
 import java.util.regex.Matcher
 
@@ -30,18 +30,65 @@ import java.util.regex.Matcher
  */
 @SuppressWarnings('unused')
 class InstrumentPlugin implements Plugin<Project> {
+  public static final String DEFAULT_JAVA_VERSION = 'default'
+  public static final String INSTRUMENT_PLUGIN_CLASSPATH_CONFIGURATION = 'instrumentPluginClasspath'
+  private final Logger logger = Logging.getLogger(InstrumentPlugin)
+
   @Override
   void apply(Project project) {
     InstrumentExtension extension = project.extensions.create('instrument', InstrumentExtension)
 
-    project.tasks.matching {
-      it.name in ['compileJava', 'compileScala', 'compileKotlin', 'compileGroovy'] ||
-        it.name =~ /compileMain_.+Java/
-    }.all {
-      AbstractCompile compileTask = it as AbstractCompile
-      Matcher versionMatcher = it.name =~ /compileMain_(.+)Java/
-      project.afterEvaluate {
-        if (!compileTask.source.empty) {
+    project.pluginManager.withPlugin("java") { configurePostCompilationInstrumentation("java", project, extension) }
+    project.pluginManager.withPlugin("kotlin") { configurePostCompilationInstrumentation("kotlin", project, extension) }
+    project.pluginManager.withPlugin("scala") { configurePostCompilationInstrumentation("scala", project, extension) }
+    project.pluginManager.withPlugin("groovy") { configurePostCompilationInstrumentation("groovy", project, extension) }
+  }
+
+  private void configurePostCompilationInstrumentation(String language, Project project, InstrumentExtension extension) {
+    project.extensions.configure(SourceSetContainer) { SourceSetContainer sourceSets ->
+      // For any "main" source-set configure its compile task
+      sourceSets.configureEach { SourceSet sourceSet ->
+        def sourceSetName = sourceSet.name
+        logger.info("[InstrumentPlugin] source-set: $sourceSetName, language: $language")
+
+        if (!sourceSetName.startsWith(SourceSet.MAIN_SOURCE_SET_NAME)) {
+          logger.debug("[InstrumentPlugin] Skipping non-main source set {} for language {}", sourceSetName, language)
+          return
+        }
+
+        def compileTaskName = sourceSet.getCompileTaskName(language)
+        logger.info("[InstrumentPlugin] compile task name: " + compileTaskName)
+
+        // For each compile task, append an instrumenting post-processing step
+        // Examples of compile tasks:
+        // - compileJava,
+        // - compileMain_java17Java,
+        // - compileMain_jetty904Java,
+        // - compileMain_play25Java,
+        // - compileKotlin,
+        // - compileScala,
+        // - compileGroovy,
+        def compileTasks = project.tasks.withType(AbstractCompile).matching {
+          it.name == compileTaskName && !it.source.isEmpty()
+        }
+
+        project.configurations.whenObjectAdded { pluginConfig ->
+          if (pluginConfig.name == INSTRUMENT_PLUGIN_CLASSPATH_CONFIGURATION) {
+            logger.info('[InstrumentPlugin] instrumentPluginClasspath configuration was created')
+            compileTasks.configureEach {
+              it.inputs.files(pluginConfig)
+            }
+          }
+        }
+
+        compileTasks.configureEach {
+          if (it.source.isEmpty()) {
+            logger.debug("[InstrumentPlugin] Skipping $compileTaskName for source set $sourceSetName as it has no source files")
+            return
+          }
+
+          // Compute optional Java version
+          Matcher versionMatcher = compileTaskName =~ /compileMain_(.+)Java/
           String sourceSetSuffix = null
           String javaVersion = null
           if (versionMatcher.matches()) {
@@ -50,109 +97,112 @@ class InstrumentPlugin implements Plugin<Project> {
               javaVersion = sourceSetSuffix[4..-1]
             }
           }
+          javaVersion = javaVersion ?: DEFAULT_JAVA_VERSION // null not accepted
+          it.inputs.property("javaVersion", javaVersion)
 
-          // insert intermediate 'raw' directory for unprocessed classes
-          Directory classesDir = compileTask.destinationDirectory.get()
-          Directory rawClassesDir = classesDir.dir(
-            "../raw${sourceSetSuffix ? "_$sourceSetSuffix" : ''}/")
-          compileTask.destinationDirectory.set(rawClassesDir.asFile)
-
-          // insert task between compile and jar, and before test*
-          String instrumentTaskName = compileTask.name.replace('compile', 'instrument')
-          def instrumentTask = project.tasks.register(instrumentTaskName, InstrumentTask) {
-            // Task configuration
-            it.group = 'Byte Buddy'
-            it.description = "Instruments the classes compiled by ${compileTask.name}"
-            it.inputs.dir(compileTask.destinationDirectory)
-            it.outputs.dir(classesDir)
-            // Task inputs
-            it.javaVersion = javaVersion
-            def instrumenterConfiguration = project.configurations.named('instrumentPluginClasspath')
-            if (instrumenterConfiguration.present) {
-              it.pluginClassPath.from(instrumenterConfiguration.get())
-            }
-            it.plugins = extension.plugins
-            it.instrumentingClassPath.from(
-              findCompileClassPath(project, it.name) +
-                rawClassesDir +
-                findAdditionalClassPath(extension, it.name)
+          it.inputs.property("plugins", extension.plugins)
+
+          it.inputs.files(extension.additionalClasspath)
+
+          // Temporary location for raw (un-instrumented) classes
+          DirectoryProperty tmpUninstrumentedClasses = project.objects.directoryProperty().value(
+            project.layout.buildDirectory.dir("tmp/${it.name}-raw-classes")
+          )
+
+          // Class path to use for instrumentation post-processing
+          ConfigurableFileCollection instrumentingClassPath = project.objects.fileCollection()
+          instrumentingClassPath.setFrom(
+            it.classpath,
+            extension.additionalClasspath,
+            tmpUninstrumentedClasses
+          )
+
+          // This were the post processing happens, i.e. were the instrumentation is applied
+          it.doLast(
+            "instrumentClasses",
+            project.objects.newInstance(
+              InstrumentPostProcessingAction,
+              javaVersion,
+              extension.plugins,
+              instrumentingClassPath,
+              it.destinationDirectory,
+              tmpUninstrumentedClasses
             )
-            it.sourceDirectory = rawClassesDir
-            // Task output
-            it.targetDirectory = classesDir
-          }
-          if (javaVersion) {
-            project.tasks.named(project.sourceSets."main_java${javaVersion}".classesTaskName) {
-              it.dependsOn(instrumentTask)
-            }
-          } else {
-            project.tasks.named(project.sourceSets.main.classesTaskName) {
-              it.dependsOn(instrumentTask)
-            }
-          }
+          )
+          logger.info("[InstrumentPlugin] Configured post-compile instrumentation for $compileTaskName for source-set $sourceSetName")
         }
       }
     }
   }
-
-  static findCompileClassPath(Project project, String taskName) {
-    def matcher = taskName =~ /instrument([A-Z].+)Java/
-    def cfgName = matcher.matches() ? "${matcher.group(1).uncapitalize()}CompileClasspath" : 'compileClasspath'
-    project.configurations.named(cfgName).findAll {
-      it.name != 'previous-compilation-data.bin' && !it.name.endsWith('.gz')
-    }
-  }
-
-  static findAdditionalClassPath(InstrumentExtension extension, String taskName) {
-    extension.additionalClasspath.getOrDefault(taskName, []).collect {
-      // insert intermediate 'raw' directory for unprocessed classes
-      def fileName = it.get().asFile.name
-      it.get().dir("../${fileName.replaceFirst('^main', 'raw')}")
-    }
-  }
 }
 
 abstract class InstrumentExtension {
   abstract ListProperty<String> getPlugins()
-  Map<String /* taskName */, List<DirectoryProperty>> additionalClasspath = [:]
+  abstract ListProperty<DirectoryProperty> getAdditionalClasspath()
 }
 
-abstract class InstrumentTask extends DefaultTask {
-  @Input @Optional
-  String javaVersion
-  @InputFiles @Classpath
-  abstract ConfigurableFileCollection getPluginClassPath()
-  @Input
-  ListProperty<String> plugins
-  @InputFiles @Classpath
-  abstract ConfigurableFileCollection getInstrumentingClassPath()
-  @InputDirectory
-  Directory sourceDirectory
-
-  @OutputDirectory
-  Directory targetDirectory
+abstract class InstrumentPostProcessingAction implements Action<AbstractCompile> {
+  private final Logger logger = Logging.getLogger(InstrumentPostProcessingAction)
+
+  @Inject
+  abstract Project getProject()
 
   @Inject
   abstract JavaToolchainService getJavaToolchainService()
+
   @Inject
   abstract BuildInvocationDetails getInvocationDetails()
+
   @Inject
   abstract WorkerExecutor getWorkerExecutor()
 
-  @TaskAction
-  instrument() {
+  // Those cannot be private other wise Groovy will fail at runtime with a missing property ex
+  final String javaVersion
+  final ListProperty<String> plugins
+  final FileCollection instrumentingClassPath
+  final DirectoryProperty compilerOutputDirectory
+  final DirectoryProperty tmpDirectory
+
+  @Inject
+  InstrumentPostProcessingAction(
+    String javaVersion,
+    ListProperty<String> plugins,
+    FileCollection instrumentingClassPath,
+    DirectoryProperty compilerOutputDirectory,
+    DirectoryProperty tmpDirectory
+  ) {
+    this.javaVersion = javaVersion
+    this.plugins = plugins
+    this.instrumentingClassPath = instrumentingClassPath
+    this.compilerOutputDirectory = compilerOutputDirectory
+    this.tmpDirectory = tmpDirectory
+  }
+
+  @Override
+  void execute(AbstractCompile task) {
+    logger.info(
+      "[InstrumentPostProcessingAction] About to instrument classes \n" +
+        "  javaVersion=${javaVersion}, \n" +
+        "  plugins=${plugins.get()}, \n" +
+        "  instrumentingClassPath=${instrumentingClassPath.files}, \n" +
+        "  rawClassesDirectory=${compilerOutputDirectory.get().asFile}"
+    )
+    def postCompileAction = this
     workQueue().submit(InstrumentAction.class, parameters -> {
-      parameters.buildStartedTime.set(this.invocationDetails.buildStartedTime)
-      parameters.pluginClassPath.from(this.pluginClassPath)
-      parameters.plugins.set(this.plugins)
-      parameters.instrumentingClassPath.setFrom(this.instrumentingClassPath)
-      parameters.sourceDirectory.set(this.sourceDirectory.asFile)
-      parameters.targetDirectory.set(this.targetDirectory.asFile)
+      parameters.buildStartedTime.set(invocationDetails.buildStartedTime)
+      parameters.pluginClassPath.from(
+        project.configurations.findByName(InstrumentPlugin.INSTRUMENT_PLUGIN_CLASSPATH_CONFIGURATION)
+          ?: project.files()
+      )
+      parameters.plugins.set(postCompileAction.plugins)
+      parameters.instrumentingClassPath.setFrom(postCompileAction.instrumentingClassPath)
+      parameters.compilerOutputDirectory.set(postCompileAction.compilerOutputDirectory)
+      parameters.tmpDirectory.set(postCompileAction.tmpDirectory)
     })
   }
 
   private workQueue() {
-    if (this.javaVersion) {
+    if (this.javaVersion != InstrumentPlugin.DEFAULT_JAVA_VERSION) {
       def javaLauncher = this.javaToolchainService.launcherFor { spec ->
         spec.languageVersion.set(JavaLanguageVersion.of(this.javaVersion))
       }.get()
@@ -172,8 +222,8 @@ interface InstrumentWorkParameters extends WorkParameters {
   ConfigurableFileCollection getPluginClassPath()
   ListProperty<String> getPlugins()
   ConfigurableFileCollection getInstrumentingClassPath()
-  DirectoryProperty getSourceDirectory()
-  DirectoryProperty getTargetDirectory()
+  DirectoryProperty getCompilerOutputDirectory()
+  DirectoryProperty getTmpDirectory()
 }
 
 abstract class InstrumentAction implements WorkAction<InstrumentWorkParameters> {
@@ -199,10 +249,29 @@ abstract class InstrumentAction implements WorkAction<InstrumentWorkParameters>
         }
       }
     }
-    File sourceDirectory = parameters.getSourceDirectory().get().asFile
-    File targetDirectory = parameters.getTargetDirectory().get().asFile
+    Path classesDirectory = parameters.compilerOutputDirectory.get().asFile.toPath()
+    Path tmpUninstrumentedDir = parameters.tmpDirectory.get().asFile.toPath()
+
+    // Delete previous tmpSourceDir contents recursively
+    if (Files.exists(tmpUninstrumentedDir)) {
+      Files.walk(tmpUninstrumentedDir)
+        .sorted(Comparator.reverseOrder())
+        .forEach { p ->
+          if (!p.equals(tmpUninstrumentedDir)) {
+            Files.deleteIfExists(p)
+          }
+        }
+    }
+
+    Files.move(
+      classesDirectory,
+      tmpUninstrumentedDir,
+      StandardCopyOption.REPLACE_EXISTING,
+      StandardCopyOption.ATOMIC_MOVE,
+    )
+
     ClassLoader instrumentingCL = createClassLoader(parameters.instrumentingClassPath, pluginCL)
-    InstrumentingPlugin.instrumentClasses(plugins, instrumentingCL, sourceDirectory, targetDirectory)
+    InstrumentingPlugin.instrumentClasses(plugins, instrumentingCL, tmpUninstrumentedDir.toFile(), classesDirectory.toFile())
   }
 
   static ClassLoader createClassLoader(cp, parent = InstrumentAction.classLoader) {

buildSrc/src/main/groovy/MuzzlePlugin.groovy

@@ -15,13 +15,7 @@ import org.eclipse.aether.spi.connector.transport.TransporterFactory
 import org.eclipse.aether.transport.http.HttpTransporterFactory
 import org.eclipse.aether.util.version.GenericVersionScheme
 import org.eclipse.aether.version.Version
-import org.gradle.api.Action
-import org.gradle.api.DefaultTask
-import org.gradle.api.GradleException
-import org.gradle.api.NamedDomainObjectProvider
-import org.gradle.api.Plugin
-import org.gradle.api.Project
-import org.gradle.api.Task
+import org.gradle.api.*
 import org.gradle.api.artifacts.Configuration
 import org.gradle.api.file.ConfigurableFileCollection
 import org.gradle.api.file.FileCollection
@@ -39,7 +33,6 @@ import org.gradle.workers.WorkerExecutor
 import java.lang.reflect.Method
 import java.util.function.BiFunction
 import java.util.regex.Pattern
-
 /**
  * muzzle task plugin which runs muzzle validation against a range of dependencies.
  */
@@ -110,8 +103,15 @@ class MuzzlePlugin implements Plugin<Project> {
 
     // compileMuzzle compiles all projects required to run muzzle validation.
     // Not adding group and description to keep this task from showing in `gradle tasks`.
+
+
+    def mainSourceSetOutputs = project.providers.provider {
+      project.sourceSets.matching { sourceSetName.startsWith(SourceSet.MAIN_SOURCE_SET_NAME) }.collect {
+        it.output
+      }
+    }
     TaskProvider<Task> compileMuzzle = project.tasks.register('compileMuzzle') {
-      it.dependsOn(project.tasks.withType(InstrumentTask))
+      it.inputs.files(mainSourceSetOutputs)
       it.dependsOn bootstrapProject.tasks.named("compileJava")
       it.dependsOn bootstrapProject.tasks.named("compileMain_java11Java")
       it.dependsOn toolingProject.tasks.named("compileJava")

buildSrc/src/test/groovy/InstrumentPluginTest.groovy

@@ -26,7 +26,7 @@ class InstrumentPluginTest extends Specification {
       compileOnly group: 'net.bytebuddy', name: 'byte-buddy', version: '1.17.5' // just to build TestPlugin
     }
 
-    apply plugin: 'instrument'
+    // apply plugin: 'instrument'
 
     instrument.plugins = [
       'TestPlugin'