Merge pull request #1124 from alexarchambault/two-step-startup

Add two-step startup

.github/scripts/run/run-its-windows.sh

@@ -9,13 +9,28 @@ if [ "$(expr substr $(uname -s) 1 5 2>/dev/null)" == "MINGW" ]; then
   RUN_APP="$RUN_APP.exe"
 fi
 
+checkResults() {
+  RESULTS="$(jq -r '.[1] | map(.status) | join("\n")' < test-output.json | sort -u)"
+  if [ "$RESULTS" != "Success" ]; then
+    exit 1
+  fi
+}
+
 "$SCALA_CLI" --power package .github/scripts/run --native-image -o "$RUN_APP"
 
-./mill -i show "scala.integration.test.testCommand" > test-args.json
-cat test-args.json
-"$RUN_APP" test-args.json
+trap "jps -mlv" EXIT
 
-RESULTS="$(jq -r '.[1] | map(.status) | join("\n")' < test-output.json | sort -u)"
-if [ "$RESULTS" != "Success" ]; then
-  exit 1
-fi
+./mill -i show "scala.integration.test.testCommand" "almond.integration.KernelTestsTwoStepStartup212.*" > test-args-212.json
+cat test-args-212.json
+"$RUN_APP" test-args-212.json
+checkResults
+
+./mill -i show "scala.integration.test.testCommand" "almond.integration.KernelTestsTwoStepStartup213.*" > test-args-213.json
+cat test-args-213.json
+"$RUN_APP" test-args-213.json
+checkResults
+
+./mill -i show "scala.integration.test.testCommand" "almond.integration.KernelTestsTwoStepStartup3.*" > test-args-3.json
+cat test-args-3.json
+"$RUN_APP" test-args-3.json
+checkResults

.github/workflows/ci.yml

@@ -93,7 +93,7 @@ jobs:
         apps:
     - run: ./mill -i scala.integration.test
 
-  windows-tests:
+  windows-integration-tests:
     runs-on: windows-latest
     steps:
     - name: Don't convert LF to CRLF during checkout
@@ -189,7 +189,7 @@ jobs:
   # job whose name doesn't change when we bump Scala versions, add OSes, …
   # We require this job for auto-merge.
   all-tests:
-    needs: [examples, bincompat, test, integration-tests, windows-tests, website, publishLocal]
+    needs: [examples, bincompat, test, integration-tests, windows-integration-tests, website, publishLocal]
     runs-on: ubuntu-latest
     steps:
     - run: true

build.sc

@@ -9,6 +9,7 @@ import $file.scripts.website0.Website, Website.Relativize
 import $file.project.settings, settings.{
   AlmondModule,
   AlmondRepositories,
+  AlmondSimpleModule,
   AlmondTestModule,
   BootstrapLauncher,
   DependencyListResource,
@@ -348,6 +349,35 @@ class ScalaKernelHelper(val crossScalaVersion: String) extends AlmondModule with
   )
 }
 
+trait Launcher extends AlmondSimpleModule with BootstrapLauncher with PropertyFile
+    with Bloop.Module {
+  def supports3    = true
+  private def sv   = ScalaVersions.scala3Latest
+  def scalaVersion = sv
+  def moduleDeps = Seq(
+    shared.kernel(ScalaVersions.scala3Compat)
+  )
+  def ivyDeps = Agg(
+    Deps.caseApp,
+    Deps.coursierApi,
+    Deps.coursierLauncher,
+    Deps.directiveHandler,
+    Deps.fansi,
+    Deps.scalaparse
+  )
+
+  def propertyFilePath = "almond/launcher/launcher.properties"
+  def propertyExtra = T {
+    val mainClass = scala.`scala-kernel`(ScalaVersions.scala3Latest).mainClass().getOrElse {
+      sys.error("No main class found")
+    }
+    Seq(
+      "kernel-main-class"     -> mainClass,
+      "default-scala-version" -> ScalaVersions.scala3Latest
+    )
+  }
+}
+
 class ScalaKernelApiHelper(val crossScalaVersion: String) extends AlmondModule with ExternalSources
     with Bloop.Module {
   def skipBloop             = !ScalaVersions.binaries.contains(crossScalaVersion)
@@ -423,6 +453,7 @@ object scala extends Module {
   object `scala-kernel`      extends Cross[ScalaKernel](ScalaVersions.all: _*)
   object `scala-kernel-helper`
       extends Cross[ScalaKernelHelper](ScalaVersions.all.filter(_.startsWith("3.")): _*)
+  object launcher         extends Launcher
   object `almond-scalapy` extends Cross[AlmondScalaPy](ScalaVersions.binaries: _*)
   object `almond-rx`      extends Cross[AlmondRx](ScalaVersions.scala212, ScalaVersions.scala213)
 
@@ -483,7 +514,15 @@ class KernelLocalRepo(val testScalaVersion: String) extends LocalRepo {
       scala.`scala-kernel-api`(testScalaVersion),
       scala.`jupyter-api`(ScalaVersions.binary(testScalaVersion)),
       scala.`scala-interpreter`(testScalaVersion),
-      scala.`toree-hooks`(ScalaVersions.binary(testScalaVersion))
+      scala.`toree-hooks`(ScalaVersions.binary(testScalaVersion)),
+      scala.launcher,
+      shared.kernel(ScalaVersions.binary(ScalaVersions.scala3Latest)),
+      shared.interpreter(ScalaVersions.binary(ScalaVersions.scala3Latest)),
+      shared.`interpreter-api`(ScalaVersions.binary(ScalaVersions.scala3Latest)),
+      shared.protocol(ScalaVersions.binary(ScalaVersions.scala3Latest)),
+      shared.channels(ScalaVersions.binary(ScalaVersions.scala3Latest)),
+      shared.logger(ScalaVersions.binary(ScalaVersions.scala3Latest)),
+      shared.`logger-scala2-macros`(ScalaVersions.binary(ScalaVersions.scala3Latest))
     ) ++ extra
   }
   def version = scala.`scala-kernel`(testScalaVersion).publishVersion()
@@ -530,11 +569,10 @@ trait Integration extends SbtModule {
     )
 
     // based on https://github.com/com-lihaoyi/mill/blob/cfbafb806351c3e1664de4e2001d3d1ddda045da/scalalib/src/TestModule.scala#L98-L164
-    def testCommand: T[Map[String, Seq[String]]] =
-      T.task {
+    def testCommand(args: String*) =
+      T.command {
         import mill.testrunner.TestRunner
 
-        val args          = Nil
         val globSelectors = Nil
         val outputPath    = os.pwd / "test-output.json"
         val useArgsFile   = testUseArgsFile()
@@ -749,13 +787,17 @@ def jupyter0(args: Seq[String], fast: Boolean, console: Boolean = false) = {
   val launcher =
     if (fast) scala.`scala-kernel`(sv).fastLauncher
     else scala.`scala-kernel`(sv).launcher
+  val specialLauncher =
+    if (fast) scala.launcher.fastLauncher
+    else scala.launcher.launcher
   T.command {
-    val jupyterDir = T.ctx().dest / "jupyter"
-    val launcher0  = launcher().path.toNIO
+    val jupyterDir       = T.ctx().dest / "jupyter"
+    val launcher0        = launcher().path.toNIO
+    val specialLauncher0 = specialLauncher().path.toNIO
     if (console)
-      jupyterConsole0(launcher0, jupyterDir.toNIO, args0)
+      jupyterConsole0(launcher0, specialLauncher0, jupyterDir.toNIO, args0)
     else
-      jupyterServer(launcher0, jupyterDir.toNIO, args0)
+      jupyterServer(launcher0, specialLauncher0, jupyterDir.toNIO, args0)
   }
 }
 
@@ -806,6 +848,11 @@ def launcher(scalaVersion: String = ScalaVersions.scala213) = T.command {
   println(launcher)
 }
 
+def specialLauncher(scalaVersion: String = ScalaVersions.scala213) = T.command {
+  val launcher = scala.launcher.launcher().path.toNIO
+  println(launcher)
+}
+
 private val examplesDir = os.pwd / "examples"
 def exampleNotebooks = T.sources {
   os.list(examplesDir)
@@ -932,6 +979,11 @@ def launcherFast(scalaVersion: String = ScalaVersions.scala213) = T.command {
   println(launcher)
 }
 
+def specialLauncherFast(scalaVersion: String = ScalaVersions.scala213) = T.command {
+  val launcher = scala.launcher.fastLauncher().path.toNIO
+  println(launcher)
+}
+
 def ghOrg  = "almond-sh"
 def ghName = "almond"
 object ci extends Module {

modules/scala/integration/src/test/scala/almond/integration/KernelLauncher.scala

@@ -86,7 +86,7 @@ object KernelLauncher {
   }
 }
 
-class KernelLauncher(testScalaVersion: String) {
+class KernelLauncher(isTwoStepStartup: Boolean, defaultScalaVersion: String) {
 
   import KernelLauncher._
 
@@ -111,6 +111,17 @@ class KernelLauncher(testScalaVersion: String) {
       "-r",
       "jitpack"
     )
+    val launcherArgs =
+      if (isTwoStepStartup)
+        Seq(s"sh.almond:launcher_3:$almondVersion")
+      else
+        Seq(
+          s"sh.almond:::scala-kernel:$almondVersion",
+          "--shared",
+          "sh.almond:::scala-kernel-api",
+          "--scala",
+          defaultScalaVersion
+        )
     val res = os.proc(
       cs,
       "bootstrap",
@@ -121,17 +132,13 @@ class KernelLauncher(testScalaVersion: String) {
       repoArgs,
       "-o",
       jarDest,
-      s"sh.almond:::scala-kernel:$almondVersion",
-      "--shared",
-      "sh.almond:::scala-kernel-api",
-      "--scala",
-      testScalaVersion,
+      launcherArgs,
       extraOptions
     )
       .call(stdin = os.Inherit, stdout = os.Inherit, check = false)
     if (res.exitCode != 0)
       sys.error(
-        s"""Error generating an Almond $almondVersion launcher for Scala $testScalaVersion
+        s"""Error generating an Almond $almondVersion launcher for Scala $defaultScalaVersion
            |
            |If that error is unexpected, you might want to:
            |- remove out/repo
@@ -185,16 +192,20 @@ class KernelLauncher(testScalaVersion: String) {
   def runner(): Runner with AutoCloseable =
     new Runner with AutoCloseable {
 
+      override def differedStartUp = isTwoStepStartup
+
       var proc: os.SubProcess = null
       var sessions            = List.empty[Session with AutoCloseable]
 
-      def withSession[T](options: String*)(f: Session => T): T =
+      def withSession[T](options: String*)(f: Session => T)(implicit sessionId: SessionId): T =
         withRunnerSession(options, Nil, Nil)(f)
-      def withExtraClassPathSession[T](extraClassPath: String*)(options: String*)(f: Session => T)
-        : T =
+      def withExtraClassPathSession[T](extraClassPath: String*)(options: String*)(f: Session => T)(
+        implicit sessionId: SessionId
+      ): T =
         withRunnerSession(options, Nil, extraClassPath)(f)
-      def withLauncherOptionsSession[T](launcherOptions: String*)(options: String*)(f: Session => T)
-        : T =
+      def withLauncherOptionsSession[T](launcherOptions: String*)(options: String*)(
+        f: Session => T
+      )(implicit sessionId: SessionId): T =
         withRunnerSession(options, launcherOptions, Nil)(f)
 
       def apply(options: String*): Session =
@@ -208,9 +219,9 @@ class KernelLauncher(testScalaVersion: String) {
         options: Seq[String],
         launcherOptions: Seq[String],
         extraClassPath: Seq[String]
-      )(f: Session => T): T = {
-        val sess    = runnerSession(options, launcherOptions, extraClassPath)
-        var running = true
+      )(f: Session => T)(implicit sessionId: SessionId): T = {
+        implicit val sess = runnerSession(options, launcherOptions, extraClassPath)
+        var running       = true
 
         val currentThread = Thread.currentThread()
 
@@ -230,9 +241,20 @@ class KernelLauncher(testScalaVersion: String) {
           }
 
         t.start()
-        try f(sess)
-        finally
+        try {
+          val t = f(sess)
+          if (Properties.isWin)
+            // On Windows, exit the kernel manually from the inside, so that all involved processes
+            // exit cleanly. A call to Process#destroy would only destroy the first kernel process,
+            // not any of its sub-processes (which would stay around, and such processes would end up
+            // filling up memory on Windows).
+            exit()
+          t
+        }
+        finally {
           running = false
+          close()
+        }
       }
 
       private def runnerSession(
@@ -252,29 +274,65 @@ class KernelLauncher(testScalaVersion: String) {
         os.write(connFile, writeToArray(connDetails))
 
         val (jarLauncher0, launcher0) =
-          if (launcherOptions.isEmpty)
+          if (isTwoStepStartup || launcherOptions.isEmpty)
             (jarLauncher, launcher)
           else
             generateLauncher(launcherOptions)
 
+        val baseCp =
+          if (isTwoStepStartup)
+            jarLauncher0.toString
+          else
+            (extraClassPath :+ jarLauncher0.toString)
+              .filter(_.nonEmpty)
+              .mkString(File.pathSeparator)
         val baseCmd: os.Shellable =
           Seq[os.Shellable](
             "java",
-            "-Xmx512m",
+            "-Xmx1g",
             "-cp",
-            (extraClassPath :+ jarLauncher0.toString).mkString(File.pathSeparator),
+            baseCp,
             "coursier.bootstrap.launcher.Launcher"
           )
 
-        proc = os.proc(
+        val extraStartupClassPathOpts =
+          if (isTwoStepStartup)
+            extraClassPath.flatMap(elem => Seq("--extra-startup-class-path", elem)) ++
+              launcherOptions ++
+              Seq("--scala", defaultScalaVersion)
+          else
+            Nil
+
+        val proc0 = os.proc(
           baseCmd,
           "--log",
           "debug",
           "--color=false",
           "--connection-file",
           connFile,
+          extraStartupClassPathOpts,
           options
-        ).spawn(cwd = dir, stdin = os.Inherit, stdout = os.Inherit)
+        )
+        System.err.println(s"Running ${proc0.command.flatMap(_.value).mkString(" ")}")
+        val extraEnv =
+          if (isTwoStepStartup) {
+            val baseRepos = sys.env.getOrElse(
+              "COURSIER_REPOSITORIES",
+              "ivy2Local|central"
+            )
+            Map(
+              "COURSIER_REPOSITORIES" ->
+                s"$baseRepos|ivy:${localRepoRoot.toNIO.toUri.toASCIIString.stripSuffix("/")}/[defaultPattern]"
+            )
+          }
+          else
+            Map.empty[String, String]
+        proc = proc0.spawn(
+          cwd = dir,
+          env = extraEnv,
+          stdin = os.Inherit,
+          stdout = os.Inherit
+        )
 
         val conn = params.channels(
           bind = false,
@@ -296,9 +354,13 @@ class KernelLauncher(testScalaVersion: String) {
         if (proc != null) {
           if (proc.isAlive()) {
             proc.close()
-            proc.waitFor(3.seconds.toMillis)
-            if (proc.isAlive())
+            val timeout = 3.seconds
+            if (!proc.waitFor(timeout.toMillis)) {
+              System.err.println(
+                s"Test kernel still running after $timeout, destroying it forcibly"
+              )
               proc.destroyForcibly()
+            }
           }
           proc = null
         }

modules/scala/integration/src/test/scala/almond/integration/KernelTestsDefinitions.scala

@@ -2,9 +2,12 @@ package almond.integration
 
 import almond.testkit.Dsl._
 
-abstract class KernelTestsDefinitions(scalaVersion: String) extends munit.FunSuite {
+abstract class KernelTestsDefinitions(
+  scalaVersion: String,
+  isTwoStepStartup: Boolean
+) extends munit.FunSuite {
 
-  val kernelLauncher = new KernelLauncher(scalaVersion)
+  val kernelLauncher = new KernelLauncher(isTwoStepStartup = isTwoStepStartup, scalaVersion)
 
   test("jvm-repr") {
     kernelLauncher.withKernel { implicit runner =>

modules/scala/integration/src/test/scala/almond/integration/KernelTestsSimple212.scala

@@ -1,3 +1,3 @@
 package almond.integration
 
-class KernelTestsSimple212 extends KernelTestsDefinitions("2.12.17")
+class KernelTestsSimple212 extends KernelTestsDefinitions("2.12.17", isTwoStepStartup = false)