write pipelined tasty in parallel.

bishabosha · bishabosha · commit a1a2c1049953 · 2024-04-10T17:53:33.000+02:00
diff --git a/compiler/src/dotty/tools/backend/jvm/GenBCode.scala b/compiler/src/dotty/tools/backend/jvm/GenBCode.scala
@@ -10,6 +10,10 @@ import Symbols.*
 import dotty.tools.io.*
 import scala.collection.mutable
 import scala.compiletime.uninitialized
+import java.util.concurrent.TimeoutException
+
+import scala.concurrent.duration.given
+import scala.concurrent.Await
 
 class GenBCode extends Phase { self =>
 
@@ -90,6 +94,20 @@ class GenBCode extends Phase { self =>
     try
       val result = super.runOn(units)
       generatedClassHandler.complete()
+      for holder <- ctx.asyncTastyPromise do
+        try
+          val asyncState = Await.result(holder.promise.future, 5.seconds)
+          for reporter <- asyncState.pending do
+            reporter.relayReports(frontendAccess.backendReporting)
+        catch
+          case _: TimeoutException =>
+            report.error(
+              """Timeout (5s) in backend while waiting for async writing of TASTy files to -Yearly-tasty-output,
+                |  this may be a bug in the compiler.
+                |
+                |Alternatively consider turning off pipelining for this project.""".stripMargin
+            )
+      end for
       result
     finally
       // frontendAccess and postProcessor are created lazilly, clean them up only if they were initialized
diff --git a/compiler/src/dotty/tools/dotc/Driver.scala b/compiler/src/dotty/tools/dotc/Driver.scala
@@ -66,6 +66,13 @@ class Driver {
 
   protected def command: CompilerCommand = ScalacCommand
 
+  private def setupAsyncTasty(ictx: FreshContext): Unit = inContext(ictx):
+    ictx.settings.YearlyTastyOutput.value match
+      case earlyOut if earlyOut.isDirectory && earlyOut.exists =>
+        ictx.setInitialAsyncTasty()
+      case _ =>
+        () // do nothing
+
   /** Setup context with initialized settings from CLI arguments, then check if there are any settings that
    *  would change the default behaviour of the compiler.
    *
@@ -82,6 +89,7 @@ class Driver {
     Positioned.init(using ictx)
 
     inContext(ictx) {
+      setupAsyncTasty(ictx)
       if !ctx.settings.YdropComments.value || ctx.settings.YreadComments.value then
         ictx.setProperty(ContextDoc, new ContextDocstrings)
       val fileNamesOrNone = command.checkUsage(summary, sourcesRequired)(using ctx.settings)(using ctx.settingsState)
diff --git a/compiler/src/dotty/tools/dotc/core/Contexts.scala b/compiler/src/dotty/tools/dotc/core/Contexts.scala
@@ -31,9 +31,11 @@ import StdNames.nme
 import compiletime.uninitialized
 
 import scala.annotation.internal.sharable
+import scala.concurrent.Promise
 
 import DenotTransformers.DenotTransformer
 import dotty.tools.dotc.profile.Profiler
+import dotty.tools.dotc.transform.Pickler.AsyncTastyHolder
 import dotty.tools.dotc.sbt.interfaces.{IncrementalCallback, ProgressCallback}
 import util.Property.Key
 import util.Store
@@ -54,8 +56,9 @@ object Contexts {
   private val (importInfoLoc,        store9) = store8.newLocation[ImportInfo | Null]()
   private val (typeAssignerLoc,     store10) = store9.newLocation[TypeAssigner](TypeAssigner)
   private val (progressCallbackLoc, store11) = store10.newLocation[ProgressCallback | Null]()
+  private val (tastyPromiseLoc,     store12) = store11.newLocation[Option[AsyncTastyHolder]](None)
 
-  private val initialStore = store11
+  private val initialStore = store12
 
   /** The current context */
   inline def ctx(using ctx: Context): Context = ctx
@@ -197,6 +200,8 @@ object Contexts {
     /** The current settings values */
     def settingsState: SettingsState = store(settingsStateLoc)
 
+    def asyncTastyPromise: Option[AsyncTastyHolder] = store(tastyPromiseLoc)
+
     /** The current compilation unit */
     def compilationUnit: CompilationUnit = store(compilationUnitLoc)
 
@@ -685,6 +690,10 @@ object Contexts {
       updateStore(compilationUnitLoc, compilationUnit)
     }
 
+    def setInitialAsyncTasty(): this.type =
+      assert(store(tastyPromiseLoc) == None, "trying to set async tasty promise twice!")
+      updateStore(tastyPromiseLoc, Some(AsyncTastyHolder(settings.YearlyTastyOutput.value, Promise())))
+
     def setCompilerCallback(callback: CompilerCallback): this.type = updateStore(compilerCallbackLoc, callback)
     def setIncCallback(callback: IncrementalCallback): this.type = updateStore(incCallbackLoc, callback)
     def setProgressCallback(callback: ProgressCallback): this.type = updateStore(progressCallbackLoc, callback)
diff --git a/compiler/src/dotty/tools/dotc/sbt/ExtractAPI.scala b/compiler/src/dotty/tools/dotc/sbt/ExtractAPI.scala
@@ -70,19 +70,13 @@ class ExtractAPI extends Phase {
 
   override def runOn(units: List[CompilationUnit])(using Context): List[CompilationUnit] =
     val doZincCallback = ctx.runZincPhases
-    val sigWriter: Option[Pickler.EarlyFileWriter] = ctx.settings.YearlyTastyOutput.value match
-      case earlyOut if earlyOut.isDirectory && earlyOut.exists =>
-        Some(Pickler.EarlyFileWriter(earlyOut))
-      case _ =>
-        None
     val nonLocalClassSymbols = new mutable.HashSet[Symbol]
     val units0 =
       if doZincCallback then
         val ctx0 = ctx.withProperty(NonLocalClassSymbolsInCurrentUnits, Some(nonLocalClassSymbols))
         super.runOn(units)(using ctx0)
       else
         units // still run the phase for the side effects (writing TASTy files to -Yearly-tasty-output)
-    sigWriter.foreach(writeSigFiles(units0, _))
     if doZincCallback then
       ctx.withIncCallback(recordNonLocalClasses(nonLocalClassSymbols, _))
     if ctx.settings.YjavaTasty.value then
@@ -91,57 +85,19 @@ class ExtractAPI extends Phase {
       units0
   end runOn
 
-  // Why we only write to early output in the first run?
-  // ===================================================
-  // TL;DR the point of pipeline compilation is to start downstream projects early,
-  // so we don't want to wait for suspended units to be compiled.
-  //
-  // But why is it safe to ignore suspended units?
-  // If this project contains a transparent macro that is called in the same project,
-  // the compilation unit of that call will be suspended (if the macro implementation
-  // is also in this project), causing a second run.
-  // However before we do that run, we will have already requested sbt to begin
-  // early downstream compilation. This means that the suspended definitions will not
-  // be visible in *early* downstream compilation.
-  //
-  // However, sbt will by default prevent downstream compilation happening in this scenario,
-  // due to the existence of macro definitions. So we are protected from failure if user tries
-  // to use the suspended definitions.
-  //
-  // Additionally, it is recommended for the user to move macro implementations to another project
-  // if they want to force early output. In this scenario the suspensions will no longer occur, so now
-  // they will become visible in the early-output.
-  //
-  // See `sbt-test/pipelining/pipelining-scala-macro` and `sbt-test/pipelining/pipelining-scala-macro-force`
-  // for examples of this in action.
-  //
-  // Therefore we only need to write to early output in the first run. We also provide the option
-  // to diagnose suspensions with the `-Yno-suspended-units` flag.
-  private def writeSigFiles(units: List[CompilationUnit], writer: Pickler.EarlyFileWriter)(using Context): Unit = {
-    try
-      for
-        unit <- units
-        (cls, pickled) <- unit.pickled
-        if cls.isDefinedInCurrentRun
-      do
-        val internalName =
-          if cls.is(Module) then cls.binaryClassName.stripSuffix(str.MODULE_SUFFIX).nn
-          else cls.binaryClassName
-        val _ = writer.writeTasty(internalName, pickled())
-    finally
-      writer.close()
-      if ctx.settings.verbose.value then
-        report.echo("[sig files written]")
-    end try
-  }
-
   private def recordNonLocalClasses(nonLocalClassSymbols: mutable.HashSet[Symbol], cb: interfaces.IncrementalCallback)(using Context): Unit =
     for cls <- nonLocalClassSymbols do
       val sourceFile = cls.source
       if sourceFile.exists && cls.isDefinedInCurrentRun then
         recordNonLocalClass(cls, sourceFile, cb)
-    cb.apiPhaseCompleted()
-    cb.dependencyPhaseCompleted()
+    for holder <- ctx.asyncTastyPromise do
+      import scala.concurrent.ExecutionContext.Implicits.global
+      // do not expect to be completed with failure
+      holder.promise.future.foreach: state =>
+        if !state.hasErrors then
+          // We also await the promise at GenBCode to emit warnings/errors
+          cb.apiPhaseCompleted()
+          cb.dependencyPhaseCompleted()
 
   private def recordNonLocalClass(cls: Symbol, sourceFile: SourceFile, cb: interfaces.IncrementalCallback)(using Context): Unit =
     def registerProductNames(fullClassName: String, binaryClassName: String) =
diff --git a/compiler/src/dotty/tools/dotc/transform/Pickler.scala b/compiler/src/dotty/tools/dotc/transform/Pickler.scala
@@ -9,7 +9,7 @@ import tasty.*
 import config.Printers.{noPrinter, pickling}
 import config.Feature
 import java.io.PrintStream
-import io.FileWriters.TastyWriter
+import io.FileWriters.{TastyWriter, ReadOnlyContext}
 import StdNames.{str, nme}
 import Periods.*
 import Phases.*
@@ -22,6 +22,11 @@ import compiletime.uninitialized
 import dotty.tools.io.{JarArchive, AbstractFile}
 import dotty.tools.dotc.printing.OutlinePrinter
 import scala.annotation.constructorOnly
+import scala.concurrent.Promise
+import dotty.tools.dotc.transform.Pickler.writeSigFilesAsync
+
+import scala.util.chaining.given
+import dotty.tools.io.FileWriters.BufferingDelayedReporting
 
 object Pickler {
   val name: String = "pickler"
@@ -33,8 +38,62 @@ object Pickler {
    */
   inline val ParallelPickling = true
 
+  class AsyncTastyHolder(val earlyOut: AbstractFile, val promise: Promise[AsyncTastyState])
+  class AsyncTastyState(val hasErrors: Boolean, val pending: Option[BufferingDelayedReporting])
+
+  // Why we only write to early output in the first run?
+  // ===================================================
+  // TL;DR the point of pipeline compilation is to start downstream projects early,
+  // so we don't want to wait for suspended units to be compiled.
+  //
+  // But why is it safe to ignore suspended units?
+  // If this project contains a transparent macro that is called in the same project,
+  // the compilation unit of that call will be suspended (if the macro implementation
+  // is also in this project), causing a second run.
+  // However before we do that run, we will have already requested sbt to begin
+  // early downstream compilation. This means that the suspended definitions will not
+  // be visible in *early* downstream compilation.
+  //
+  // However, sbt will by default prevent downstream compilation happening in this scenario,
+  // due to the existence of macro definitions. So we are protected from failure if user tries
+  // to use the suspended definitions.
+  //
+  // Additionally, it is recommended for the user to move macro implementations to another project
+  // if they want to force early output. In this scenario the suspensions will no longer occur, so now
+  // they will become visible in the early-output.
+  //
+  // See `sbt-test/pipelining/pipelining-scala-macro` and `sbt-test/pipelining/pipelining-scala-macro-force`
+  // for examples of this in action.
+  //
+  // Therefore we only need to write to early output in the first run. We also provide the option
+  // to diagnose suspensions with the `-Yno-suspended-units` flag.
+  def writeSigFilesAsync(
+      tasks: List[(String, Array[Byte])],
+      writer: EarlyFileWriter,
+      promise: Promise[AsyncTastyState])(using ctx: ReadOnlyContext): Unit = {
+    try
+      for (internalName, pickled) <- tasks do
+        val _ = writer.writeTasty(internalName, pickled)
+    finally
+      try
+        writer.close()
+      finally
+        promise.success(
+          AsyncTastyState(
+            hasErrors = ctx.reporter.hasErrors,
+            pending = (
+              ctx.reporter match
+                case buffered: BufferingDelayedReporting => Some(buffered)
+                case _ => None
+            )
+          )
+        )
+      end try
+    end try
+  }
+
   class EarlyFileWriter private (writer: TastyWriter, origin: AbstractFile):
-    def this(dest: AbstractFile)(using @constructorOnly ctx: Context) = this(TastyWriter(dest), dest)
+    def this(dest: AbstractFile)(using @constructorOnly ctx: ReadOnlyContext) = this(TastyWriter(dest), dest)
 
     export writer.writeTasty
 
@@ -50,13 +109,15 @@ object Pickler {
 class Pickler extends Phase {
   import ast.tpd.*
 
+  def doAsyncTasty(using Context): Boolean = ctx.asyncTastyPromise.isDefined
+
   override def phaseName: String = Pickler.name
 
   override def description: String = Pickler.description
 
   // No need to repickle trees coming from TASTY
   override def isRunnable(using Context): Boolean =
-    super.isRunnable && !ctx.settings.fromTasty.value
+    super.isRunnable && (!ctx.settings.fromTasty.value || doAsyncTasty)
 
   // when `-Yjava-tasty` is set we actually want to run this phase on Java sources
   override def skipIfJava(using Context): Boolean = false
@@ -86,11 +147,20 @@ class Pickler extends Phase {
    */
   object serialized:
     val scratch = new ScratchData
+    private val buf = mutable.ListBuffer.empty[(String, Array[Byte])]
     def run(body: ScratchData => Array[Byte]): Array[Byte] =
       synchronized {
         scratch.reset()
         body(scratch)
       }
+    def commit(internalName: String, tasty: Array[Byte]): Unit = synchronized {
+      buf += ((internalName, tasty))
+    }
+    def result(): List[(String, Array[Byte])] = synchronized {
+      val res = buf.toList
+      buf.clear()
+      res
+    }
 
   private val executor = Executor[Array[Byte]]()
 
@@ -100,10 +170,29 @@ class Pickler extends Phase {
     if isOutline then ctx.fresh.setPrinterFn(OutlinePrinter(_))
     else ctx
 
+  /** only ran under -Ypickle-write and -from-tasty */
+  private def runFromTasty(unit: CompilationUnit)(using Context): Unit = {
+    val pickled = unit.pickled
+    for (cls, bytes) <- pickled do
+      serialized.commit(computeInternalName(cls), bytes())
+  }
+
+  private def computeInternalName(cls: ClassSymbol)(using Context): String =
+    if cls.is(Module) then cls.binaryClassName.stripSuffix(str.MODULE_SUFFIX).nn
+    else cls.binaryClassName
+
   override def run(using Context): Unit = {
     val unit = ctx.compilationUnit
     pickling.println(i"unpickling in run ${ctx.runId}")
 
+    if ctx.settings.fromTasty.value then
+      // skip the rest of the phase, as tasty is already "pickled",
+      // however we still need to set up tasks to write TASTy to
+      // early output when pipelining is enabled.
+      if doAsyncTasty then
+        runFromTasty(unit)
+      return ()
+
     for
       cls <- dropCompanionModuleClasses(topLevelClasses(unit.tpdTree))
       tree <- sliceTopLevel(unit.tpdTree, cls)
@@ -137,6 +226,8 @@ class Pickler extends Phase {
       val positionWarnings = new mutable.ListBuffer[Message]()
       def reportPositionWarnings() = positionWarnings.foreach(report.warning(_))
 
+      val internalName = if doAsyncTasty then computeInternalName(cls) else ""
+
       def computePickled(): Array[Byte] = inContext(ctx.fresh) {
         serialized.run { scratch =>
           treePkl.compactify(scratch)
@@ -166,6 +257,10 @@ class Pickler extends Phase {
             println(i"**** pickled info of $cls")
             println(TastyPrinter.showContents(pickled, ctx.settings.color.value == "never"))
             println(i"**** end of pickled info of $cls")
+
+          if doAsyncTasty then
+            serialized.commit(internalName, pickled)
+
           pickled
         }
       }
@@ -194,13 +289,27 @@ class Pickler extends Phase {
   }
 
   override def runOn(units: List[CompilationUnit])(using Context): List[CompilationUnit] = {
+    val isConcurrent = useExecutor
+
+    val writeTask: Option[() => Unit] = ctx.asyncTastyPromise.map: holder =>
+      () =>
+        given ReadOnlyContext = if isConcurrent then ReadOnlyContext.buffered else ReadOnlyContext.eager
+        val writer = Pickler.EarlyFileWriter(holder.earlyOut)
+        writeSigFilesAsync(serialized.result(), writer, holder.promise)
+
+    def runPhase(writeCB: (doWrite: () => Unit) => Unit) =
+      super.runOn(units).tap(_ => writeTask.foreach(writeCB))
+
     val result =
-      if useExecutor then
+      if isConcurrent then
         executor.start()
-        try super.runOn(units)
+        try
+          runPhase: doWrite =>
+            // unless we redesign executor to have "Unit" schedule overload, we need some sentinel value.
+            executor.schedule(() => { doWrite(); Array.emptyByteArray })
         finally executor.close()
       else
-        super.runOn(units)
+        runPhase(_())
     if ctx.settings.YtestPickler.value then
       val ctx2 = ctx.fresh
         .setSetting(ctx.settings.YreadComments, true)
diff --git a/compiler/src/dotty/tools/io/FileWriters.scala b/compiler/src/dotty/tools/io/FileWriters.scala
diff --git a/compiler/src/dotty/tools/io/JarArchive.scala b/compiler/src/dotty/tools/io/JarArchive.scala
