From 7f800d4b476498fb4a7ceda41ec16c3b7c2a3a2d Mon Sep 17 00:00:00 2001
From: Brice Jaglin <bjaglin@teads.tv>
Date: Tue, 7 Jul 2020 08:53:18 +0200
Subject: [PATCH] share custom toolclasspath classloader across projects &
 configs

Since d81dd72 (see commit description), when `scalafix` or `scalafixAll`
is run on a build root that aggregates many projects, many identical
classloaders may be created when an external rule is provided on the
CLI or when projects use local rules via ScalafixConfig.

This allows aggregated tasks across projects & configs to share a
single, warm classloader when possible (same Scalafix classpath).

Note that cache is bound on purpose to a task value, so two successive
invocations of `scalafix` with the exact same arguments & environment
will not share the same classloader. This is for simplicity, as local
rules can be updated from one run to the other, and to my knowledge,
there is no simple way to invalidate a URLClassLoader if the underlying
class files change.
---
 .../scalafix/internal/sbt/BlockingCache.scala | 21 +++++++++++
 .../scala/scalafix/sbt/ScalafixPlugin.scala   | 37 ++++++++++++++-----
 2 files changed, 49 insertions(+), 9 deletions(-)
 create mode 100644 src/main/scala/scalafix/internal/sbt/BlockingCache.scala

diff --git a/src/main/scala/scalafix/internal/sbt/BlockingCache.scala b/src/main/scala/scalafix/internal/sbt/BlockingCache.scala
new file mode 100644
index 00000000..91a08c4c
--- /dev/null
+++ b/src/main/scala/scalafix/internal/sbt/BlockingCache.scala
@@ -0,0 +1,21 @@
+package scalafix.internal.sbt
+
+import scala.collection.mutable
+
+/** A basic thread-safe cache without any eviction. */
+class BlockingCache[K, V] {
+  private val underlying = new mutable.HashMap[K, V]
+
+  /**
+    * @param value By-name parameter evaluated when the key if missing. Value computation is guaranteed
+    *              to be called only once per key across all invocations.
+    */
+  def getOrElseUpdate(key: K, value: => V): V = {
+    // ConcurrentHashMap does not guarantee that there is only one evaluation of the value, so
+    // we use our own (global) locking, which is OK as the number of keys is expected to be
+    // very small (bound by the number of projects in the sbt build).
+    underlying.synchronized {
+      underlying.getOrElseUpdate(key, value)
+    }
+  }
+}
diff --git a/src/main/scala/scalafix/sbt/ScalafixPlugin.scala b/src/main/scala/scalafix/sbt/ScalafixPlugin.scala
index 1d89b28b..c47660e7 100644
--- a/src/main/scala/scalafix/sbt/ScalafixPlugin.scala
+++ b/src/main/scala/scalafix/sbt/ScalafixPlugin.scala
@@ -137,6 +137,16 @@ object ScalafixPlugin extends AutoPlugin {
       Invisible
     )
 
+  // Memoize ScalafixInterface instances initialized with a custom tool classpath across projects & configurations
+  // during task execution, to amortize the classloading cost when invoking scalafix concurrently on many targets
+  private val scalafixInterfaceCache
+      : TaskKey[BlockingCache[ToolClasspath, ScalafixInterface]] =
+    TaskKey(
+      "scalafixInterfaceCache",
+      "Implementation detail - do not use",
+      Invisible
+    )
+
   override lazy val projectConfigurations: Seq[Configuration] =
     Seq(ScalafixConfig)
 
@@ -173,7 +183,11 @@ object ScalafixPlugin extends AutoPlugin {
       // depend on settings, while local rules classpath must be looked up via tasks
       loadedRules = () => scalafixInterfaceProvider.value().availableRules(),
       terminalWidth = Some(JLineAccess.terminalWidth)
-    ).parser
+    ).parser,
+    scalafixInterfaceCache := new BlockingCache[
+      ToolClasspath,
+      ScalafixInterface
+    ]
   )
 
   override def buildSettings: Seq[Def.Setting[_]] =
@@ -186,6 +200,7 @@ object ScalafixPlugin extends AutoPlugin {
   private def scalafixArgsFromShell(
       shell: ShellArgs,
       scalafixInterface: () => ScalafixInterface,
+      scalafixInterfaceCache: BlockingCache[ToolClasspath, ScalafixInterface],
       projectDepsExternal: Seq[ModuleID],
       baseDepsExternal: Seq[ModuleID],
       baseResolvers: Seq[Repository],
@@ -214,15 +229,18 @@ object ScalafixPlugin extends AutoPlugin {
     val customToolClasspath =
       (projectDepsInternal0 ++ projectDepsExternal ++ rulesDepsExternal).nonEmpty
     val interface =
-      if (customToolClasspath)
-        scalafixInterface().withArgs(
-          ToolClasspath(
-            projectDepsInternal0.map(_.toURI.toURL),
-            baseDepsExternal ++ projectDepsExternal ++ rulesDepsExternal,
-            baseResolvers
-          )
+      if (customToolClasspath) {
+        val toolClasspath = ToolClasspath(
+          projectDepsInternal0.map(_.toURI.toURL),
+          baseDepsExternal ++ projectDepsExternal ++ rulesDepsExternal,
+          baseResolvers
+        )
+        scalafixInterfaceCache.getOrElseUpdate(
+          toolClasspath,
+          // costly: triggers artifact resolution & classloader creation
+          scalafixInterface().withArgs(toolClasspath)
         )
-      else
+      } else
         // if there is nothing specific to the project or the invocation, reuse the default
         // interface which already has the baseDepsExternal loaded
         scalafixInterface()
@@ -290,6 +308,7 @@ object ScalafixPlugin extends AutoPlugin {
         val (shell, mainInterface0) = scalafixArgsFromShell(
           shellArgs,
           scalafixInterfaceProvider.value,
+          scalafixInterfaceCache.value,
           projectDepsExternal,
           scalafixDependencies.in(ThisBuild).value,
           scalafixResolvers.in(ThisBuild).value,