Amortize the startup cost of some components (#10451)

Building Engine, Context, ApplicationConfig and Ydoc was a adding a rather large delay during the initial startup step as all of those were blocking operations.
Moving all of those to the resource initialization step hopes to amortize some of that cost since it can be done in parallel. Had to add a `ComponentSupervisor` (open for a different name suggestion) to ensure that such delayed components are properly closed on shutdown.

# Important Notes
Adding Ydoc has added a visible delay during startup. I'm hoping that we can amortize some of that with this PR:
![Screenshot from 2024-07-05 11-12-19](https://github.com/enso-org/enso/assets/292128/fd52f749-b2cb-414d-bd2a-847ea867026c)

Now:
![Screenshot from 2024-07-05 11-25-58](https://github.com/enso-org/enso/assets/292128/9e7c96c9-ee47-46c3-9bdb-8f96bbc4a68f)

engine/language-server/src/main/java/org/enso/languageserver/boot/resource/TruffleContextInitialization.java

@@ -3,15 +3,18 @@
 import akka.event.EventStream;
 import java.util.concurrent.Executor;
 import org.enso.common.LanguageInfo;
+import org.enso.languageserver.boot.ComponentSupervisor;
 import org.enso.languageserver.event.InitializedEvent;
 import org.graalvm.polyglot.Context;
+import org.graalvm.polyglot.Engine;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /** Initialize the Truffle context. */
 public class TruffleContextInitialization extends LockedInitialization {
 
-  private final Context truffleContext;
+  private final Context.Builder truffleContextBuilder;
+  private final ComponentSupervisor supervisor;
   private final EventStream eventStream;
 
   private final Logger logger = LoggerFactory.getLogger(this.getClass());
@@ -21,17 +24,36 @@ public class TruffleContextInitialization extends LockedInitialization {
    *
    * @param executor the executor that runs the initialization
    * @param eventStream the events stream
-   * @param truffleContext the Truffle context
+   * @param truffleContextBuilder the Truffle context builder
    */
   public TruffleContextInitialization(
-      Executor executor, Context truffleContext, EventStream eventStream) {
+      Executor executor,
+      Context.Builder truffleContextBuilder,
+      ComponentSupervisor supervisor,
+      EventStream eventStream) {
     super(executor);
-    this.truffleContext = truffleContext;
+    this.truffleContextBuilder = truffleContextBuilder;
+    this.supervisor = supervisor;
     this.eventStream = eventStream;
   }
 
   @Override
   public void initComponent() {
+    logger.trace("Creating Runtime context.");
+    if (Engine.newBuilder()
+        .allowExperimentalOptions(true)
+        .build()
+        .getLanguages()
+        .containsKey("java")) {
+      truffleContextBuilder
+          .option("java.ExposeNativeJavaVM", "true")
+          .option("java.Polyglot", "true")
+          .option("java.UseBindingsLoader", "true")
+          .allowCreateThread(true);
+    }
+    var truffleContext = truffleContextBuilder.build();
+    supervisor.registerService(truffleContext);
+    logger.trace("Created Runtime context [{}].", truffleContext);
     logger.info("Initializing Runtime context [{}]...", truffleContext);
     truffleContext.initialize(LanguageInfo.ID);
     eventStream.publish(InitializedEvent.TruffleContextInitialized$.MODULE$);

engine/language-server/src/main/java/org/enso/languageserver/boot/resource/YdocInitialization.java

@@ -1,26 +1,34 @@
 package org.enso.languageserver.boot.resource;
 
 import java.util.concurrent.Executor;
+import org.enso.languageserver.boot.ComponentSupervisor;
+import org.enso.languageserver.boot.config.ApplicationConfig;
 import org.enso.ydoc.Ydoc;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public final class YdocInitialization extends LockedInitialization {
 
   private final Logger logger = LoggerFactory.getLogger(this.getClass());
+  private final ComponentSupervisor supervisor;
 
-  private final Ydoc ydoc;
-
-  public YdocInitialization(Executor executor, Ydoc ydoc) {
+  public YdocInitialization(Executor executor, ComponentSupervisor componentSupervisor) {
     super(executor);
-    this.ydoc = ydoc;
+    this.supervisor = componentSupervisor;
   }
 
   @Override
   public void initComponent() {
     logger.info("Starting Ydoc server...");
+    var applicationConfig = ApplicationConfig.load();
+    var ydoc =
+        Ydoc.builder()
+            .hostname(applicationConfig.ydoc().hostname())
+            .port(applicationConfig.ydoc().port())
+            .build();
     try {
       ydoc.start();
+      this.supervisor.registerService(ydoc);
     } catch (Exception e) {
       throw new RuntimeException(e);
     }

engine/language-server/src/main/scala/org/enso/languageserver/boot/ComponentSupervisor.scala

@@ -0,0 +1,20 @@
+package org.enso.languageserver.boot
+
+/** A wrapper around an existing component/resource that can be added with a delay and ensured
+  * to be closed safely.
+  */
+class ComponentSupervisor extends AutoCloseable {
+  private var component: AutoCloseable = null
+
+  def registerService(component: AutoCloseable): Unit = {
+    assert(this.component == null, "can't register component twice")
+    this.component = component
+  }
+
+  override def close(): Unit = {
+    if (this.component != null) {
+      this.component.close()
+      this.component = null
+    }
+  }
+}

engine/language-server/src/main/scala/org/enso/languageserver/boot/MainModule.scala

@@ -50,13 +50,10 @@ import org.enso.logger.masking.Masking
 import org.enso.logger.JulHandler
 import org.enso.logger.akka.AkkaConverter
 import org.enso.common.HostAccessFactory
-import org.enso.languageserver.boot.config.ApplicationConfig
 import org.enso.polyglot.{RuntimeOptions, RuntimeServerInfo}
 import org.enso.profiling.events.NoopEventsMonitor
 import org.enso.searcher.memory.InMemorySuggestionsRepo
 import org.enso.text.{ContentBasedVersioning, Sha3_224VersionCalculator}
-import org.enso.ydoc.Ydoc
-import org.graalvm.polyglot.Engine
 import org.graalvm.polyglot.Context
 import org.graalvm.polyglot.io.MessageEndpoint
 import org.slf4j.event.Level
@@ -84,9 +81,9 @@ class MainModule(serverConfig: LanguageServerConfig, logLevel: Level) {
     logLevel
   )
 
-  private val applicationConfig = ApplicationConfig.load()
-
-  private val utcClock = Clock.systemUTC()
+  private val ydocSupervisor    = new ComponentSupervisor()
+  private val contextSupervisor = new ComponentSupervisor()
+  private val utcClock          = Clock.systemUTC()
 
   val directoriesConfig = ProjectDirectoriesConfig(serverConfig.contentRootPath)
   private val contentRoot = ContentRootWithFile(
@@ -343,30 +340,13 @@ class MainModule(serverConfig: LanguageServerConfig, logLevel: Level) {
         connection
       } else null
     })
-  if (
-    Engine
-      .newBuilder()
-      .allowExperimentalOptions(true)
-      .build
-      .getLanguages()
-      .containsKey("java")
-  ) {
-    builder
-      .option("java.ExposeNativeJavaVM", "true")
-      .option("java.Polyglot", "true")
-      .option("java.UseBindingsLoader", "true")
-      .allowCreateThread(true)
-  }
-
-  val context = builder.build()
-  log.trace("Created Runtime context [{}].", context)
 
   system.eventStream.setLogLevel(AkkaConverter.toAkka(logLevel))
   log.trace("Set akka log level to [{}].", logLevel)
 
   val runtimeKiller =
     system.actorOf(
-      RuntimeKiller.props(runtimeConnector, context),
+      RuntimeKiller.props(runtimeConnector, contextSupervisor),
       "runtime-context"
     )
 
@@ -444,21 +424,16 @@ class MainModule(serverConfig: LanguageServerConfig, logLevel: Level) {
 
   private val jsonRpcProtocolFactory = new JsonRpcProtocolFactory
 
-  private val ydoc = Ydoc
-    .builder()
-    .hostname(applicationConfig.ydoc.hostname)
-    .port(applicationConfig.ydoc.port)
-    .build()
-
   private val initializationComponent =
     ResourcesInitialization(
       system.eventStream,
       directoriesConfig,
       jsonRpcProtocolFactory,
       suggestionsRepo,
-      context,
+      builder,
+      contextSupervisor,
       zioRuntime,
-      ydoc
+      ydocSupervisor
     )(system.dispatcher)
 
   private val jsonRpcControllerFactory = new JsonConnectionControllerFactory(
@@ -529,9 +504,9 @@ class MainModule(serverConfig: LanguageServerConfig, logLevel: Level) {
   /** Close the main module releasing all resources. */
   def close(): Unit = {
     suggestionsRepo.close()
-    context.close()
+    contextSupervisor.close()
     runtimeEventsMonitor.close()
-    ydoc.close()
+    ydocSupervisor.close()
     log.info("Closed Language Server main module.")
   }
 

engine/language-server/src/main/scala/org/enso/languageserver/boot/ResourcesInitialization.scala

@@ -16,7 +16,6 @@ import org.enso.languageserver.boot.resource.{
 import org.enso.languageserver.data.ProjectDirectoriesConfig
 import org.enso.languageserver.effect
 import org.enso.searcher.memory.InMemorySuggestionsRepo
-import org.enso.ydoc.Ydoc
 import org.graalvm.polyglot.Context
 
 import scala.concurrent.ExecutionContextExecutor
@@ -42,9 +41,10 @@ object ResourcesInitialization {
     directoriesConfig: ProjectDirectoriesConfig,
     protocolFactory: ProtocolFactory,
     suggestionsRepo: InMemorySuggestionsRepo,
-    truffleContext: Context,
+    truffleContextBuilder: Context#Builder,
+    truffleContextSupervisor: ComponentSupervisor,
     runtime: effect.Runtime,
-    ydoc: Ydoc
+    ydocSupervisor: ComponentSupervisor
   )(implicit ec: ExecutionContextExecutor): InitializationComponent = {
     new SequentialResourcesInitialization(
       ec,
@@ -58,8 +58,13 @@ object ResourcesInitialization {
           eventStream,
           suggestionsRepo
         ),
-        new TruffleContextInitialization(ec, truffleContext, eventStream),
-        new YdocInitialization(ec, ydoc)
+        new TruffleContextInitialization(
+          ec,
+          truffleContextBuilder,
+          truffleContextSupervisor,
+          eventStream
+        ),
+        new YdocInitialization(ec, ydocSupervisor)
       )
     )
   }

engine/language-server/src/main/scala/org/enso/languageserver/runtime/RuntimeKiller.scala

@@ -1,13 +1,12 @@
 package org.enso.languageserver.runtime
 
 import java.util.UUID
-
 import akka.actor.{Actor, ActorRef, Cancellable, Props}
 import com.typesafe.scalalogging.LazyLogging
+import org.enso.languageserver.boot.ComponentSupervisor
 import org.enso.languageserver.runtime.RuntimeKiller._
 import org.enso.languageserver.util.UnhandledLogging
 import org.enso.polyglot.runtime.Runtime.Api
-import org.graalvm.polyglot.Context
 
 import scala.concurrent.duration._
 import scala.util.control.NonFatal
@@ -19,8 +18,10 @@ import scala.util.control.NonFatal
   * @param runtimeConnector a proxy to the runtime
   * @param truffleContext a Truffle context
   */
-class RuntimeKiller(runtimeConnector: ActorRef, truffleContext: Context)
-    extends Actor
+class RuntimeKiller(
+  runtimeConnector: ActorRef,
+  truffleContextSupervisor: ComponentSupervisor
+) extends Actor
     with LazyLogging
     with UnhandledLogging {
 
@@ -63,19 +64,16 @@ class RuntimeKiller(runtimeConnector: ActorRef, truffleContext: Context)
   private def shutDownTruffle(replyTo: ActorRef, retryCount: Int = 0): Unit = {
     try {
       logger.info(
-        "Shutting down the Truffle context [{}]. " +
-        "Attempt #{}.",
-        truffleContext,
+        "Shutting down the Truffle context. Attempt #{}.",
         retryCount + 1
       )
-      truffleContext.close()
+      truffleContextSupervisor.close()
       replyTo ! RuntimeGracefullyStopped
       context.stop(self)
     } catch {
       case NonFatal(ex) =>
         logger.error(
-          s"An error occurred during stopping Truffle context [{}]. {}",
-          truffleContext,
+          s"An error occurred during stopping Truffle context. {}",
           ex.getMessage
         )
         if (retryCount < MaxRetries) {
@@ -123,7 +121,10 @@ object RuntimeKiller {
     * @param truffleContext a Truffle context
     * @return a configuration object
     */
-  def props(runtimeConnector: ActorRef, truffleContext: Context): Props =
-    Props(new RuntimeKiller(runtimeConnector, truffleContext))
+  def props(
+    runtimeConnector: ActorRef,
+    truffleContextSupervisor: ComponentSupervisor
+  ): Props =
+    Props(new RuntimeKiller(runtimeConnector, truffleContextSupervisor))
 
 }

engine/runner/src/main/java/org/enso/runner/Main.java

@@ -1303,7 +1303,11 @@ private void launch(String[] args) {
   protected CommandLine preprocessArguments(Options options, String[] args) {
     var parser = new DefaultParser();
     try {
+      var startParsing = System.currentTimeMillis();
       var line = parser.parse(options, args);
+      logger.trace(
+          "Parsing Language Server arguments took {0}ms",
+          System.currentTimeMillis() - startParsing);
       return line;
     } catch (Exception e) {
       printHelp(options);