Initialize Language Server's JSON RPC protocol asynchronously (#7232)

As discovered in #7224, Json RPC protocol was added to the asynchronous resource initialization stage, as part of #6306, but was not in fact initialized at that point.
Instead it was initialized when the server was started to be able to serve correctly the initialization messages. A classic Catch-22. It was really hard to discover this just by looking at the code, but the profiling clearly showed where the time was spent.

This change splits Language Server's protocol into two:
- the first one accepts `heartbeat/init` and `session/initProtocolConnection`
- the second one enriches it with the full set of supported messages

This shifts the initialization from blocking for 0.5 sec to only ~30ms, and performing the second stage asynchronously.

Closes #7224.

# Important Notes
Before the change (blocking server startup):
![Screenshot from 2023-07-05 18-53-24](https://github.com/enso-org/enso/assets/292128/bcfa9043-d00a-4b36-a44c-782a388a16b9)
![Screenshot from 2023-07-05 18-53-10](https://github.com/enso-org/enso/assets/292128/54927787-4c95-46db-bd68-f3a3b82367d5)

After the change (1st stage):
![Screenshot from 2023-07-06 14-02-34](https://github.com/enso-org/enso/assets/292128/d7a7bc34-39dc-46f1-9e64-6d350697c30b)

After the change (2nd, asynchronous initialization, stage):
![Screenshot from 2023-07-06 14-21-17](https://github.com/enso-org/enso/assets/292128/def8c0a1-f211-4fc0-9df0-7c1634312166)

engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpc.scala

@@ -29,12 +29,14 @@ import org.enso.languageserver.workspace.WorkspaceApi.ProjectInfo
 
 object JsonRpc {
 
-  /** A description of supported JSON RPC messages.
-    */
-  val protocol: Protocol = Protocol.empty
-    .registerRequest(Ping)
+  /** A description of JSON RPC messages support during the initialization stage */
+  val initProtocol: Protocol = Protocol.empty
     .registerRequest(InitialPing)
     .registerRequest(InitProtocolConnection)
+
+  /** A description of supported JSON RPC messages at a post-initialization stage */
+  def fullProtocol(init: Protocol): Protocol = init
+    .registerRequest(Ping)
     .registerRequest(AcquireCapability)
     .registerRequest(ReleaseCapability)
     .registerRequest(WriteFile)
@@ -116,4 +118,5 @@ object JsonRpc {
     .registerNotification(WaitingForStandardInput)
     .registerNotification(SuggestionsDatabaseUpdates)
     .registerNotification(VisualizationEvaluationFailed)
+    .finalized()
 }

engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpcProtocolFactory.scala

@@ -1,16 +1,33 @@
 package org.enso.languageserver.protocol.json
 
-import org.enso.jsonrpc.{Protocol, ProtocolFactory}
+import org.enso.jsonrpc
+import org.enso.jsonrpc.{Errors, Protocol, ProtocolFactory}
+import org.enso.languageserver.session.SessionApi.SessionNotInitialisedError
 
 /** Factory creating JSON-RPC protocol. */
 final class JsonRpcProtocolFactory extends ProtocolFactory {
 
+  private var _protocol: Protocol = _
+
   /** @inheritdoc */
-  def getProtocol: Protocol =
-    JsonRpc.protocol
+  def getProtocol(): Protocol = {
+    if (_protocol == null) {
+      _protocol = JsonRpc.initProtocol
+    }
+    _protocol
+  }
 
   /** @inheritdoc */
   override def init(): Unit = {
-    val _ = JsonRpc.protocol
+    if (_protocol == null) {
+      _protocol = JsonRpc.initProtocol
+    }
+    _protocol = JsonRpc.fullProtocol(_protocol)
+  }
+
+  /** Error returned when a requested method is not recognized */
+  override def onMissingMethod(): jsonrpc.Error = {
+    if (_protocol != null && _protocol.initialized) Errors.MethodNotFound
+    else SessionNotInitialisedError
   }
 }

lib/scala/json-rpc-server-test/src/main/scala/org/enso/jsonrpc/test/JsonRpcServerTestKit.scala

@@ -59,7 +59,9 @@ abstract class JsonRpcServerTestKit
 
   override def beforeEach(): Unit = {
     super.beforeEach()
-    server  = new JsonRpcServer(protocolFactory, clientControllerFactory)
+    val factory = protocolFactory
+    factory.init()
+    server  = new JsonRpcServer(factory, clientControllerFactory)
     binding = Await.result(server.bind(interface, port = 0), 3.seconds)
     address = s"ws://$interface:${binding.localAddress.getPort}"
   }

lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/JsonRpcServer.scala

@@ -37,10 +37,11 @@ class JsonRpcServer(
   implicit val ec: ExecutionContext = system.dispatcher
 
   private def newUser(): Flow[Message, Message, NotUsed] = {
-    val protocol = protocolFactory.getProtocol
     val messageHandler =
       system.actorOf(
-        Props(new MessageHandlerSupervisor(clientControllerFactory, protocol)),
+        Props(
+          new MessageHandlerSupervisor(clientControllerFactory, protocolFactory)
+        ),
         s"message-handler-supervisor-${UUID.randomUUID()}"
       )
 

lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandler.scala

@@ -6,14 +6,16 @@ import org.enso.jsonrpc.Errors.InvalidParams
 
 /** An actor responsible for passing parsed massages between the web and
   * a controller actor.
-  * @param protocol a protocol object describing supported messages and their
+  * @param protocol a factory for retrieving protocol object describing supported messages and their
   *                 serialization modes.
   * @param controller the controller actor, handling parsed messages.
   */
-class MessageHandler(val protocol: Protocol, val controller: ActorRef)
+class MessageHandler(protocolFactory: ProtocolFactory, controller: ActorRef)
     extends Actor
     with Stash {
 
+  private def getProtocol(): Protocol = protocolFactory.getProtocol()
+
   /** A pre-initialization behavior, awaiting a to-web connection end.
     * @return the actor behavior.
     */
@@ -53,7 +55,7 @@ class MessageHandler(val protocol: Protocol, val controller: ActorRef)
     response: ResponseResult[Method, Any],
     webConnection: ActorRef
   ): Unit = {
-    val responseDataJson: Json = protocol.payloadsEncoder(response.data)
+    val responseDataJson: Json = getProtocol().payloadsEncoder(response.data)
     val bareResp               = JsonProtocol.ResponseResult(response.id, responseDataJson)
     webConnection ! MessageHandler.WebMessage(JsonProtocol.encode(bareResp))
   }
@@ -77,7 +79,7 @@ class MessageHandler(val protocol: Protocol, val controller: ActorRef)
     webConnection: ActorRef,
     awaitingResponses: Map[Id, Method]
   ): Unit = {
-    val paramsJson = protocol.payloadsEncoder(req.params)
+    val paramsJson = getProtocol().payloadsEncoder(req.params)
     val bareReq    = JsonProtocol.Request(req.method.name, req.id, paramsJson)
     webConnection ! MessageHandler.WebMessage(JsonProtocol.encode(bareReq))
     context.become(
@@ -89,7 +91,7 @@ class MessageHandler(val protocol: Protocol, val controller: ActorRef)
     notification: Notification[Method, Any],
     webConnection: ActorRef
   ): Unit = {
-    val paramsJson = protocol.payloadsEncoder(notification.params)
+    val paramsJson = getProtocol().payloadsEncoder(notification.params)
     val bareNotification =
       JsonProtocol.Notification(notification.method.name, paramsJson)
     webConnection ! MessageHandler.WebMessage(
@@ -131,14 +133,14 @@ class MessageHandler(val protocol: Protocol, val controller: ActorRef)
       case Some(JsonProtocol.ResponseResult(id, result)) =>
         val maybeDecoded: Option[Any] = for {
           method   <- awaitingResponses.get(id)
-          decoder  <- protocol.getResultDecoder(method)
+          decoder  <- getProtocol().getResultDecoder(method)
           response <- decoder.buildResponse(id, result)
         } yield response
         maybeDecoded.foreach(controller ! _)
         context.become(established(webConnection, awaitingResponses - id))
 
       case Some(JsonProtocol.ResponseError(mayId, bareError)) =>
-        val error = protocol
+        val error = getProtocol()
           .resolveError(bareError.code)
           .getOrElse(
             Errors
@@ -165,15 +167,16 @@ class MessageHandler(val protocol: Protocol, val controller: ActorRef)
 
   private def resolveDecoder(
     methodName: String
-  ): Either[Error, ParamsDecoder[Method, Any]] =
+  ): Either[Error, ParamsDecoder[Method, Any]] = {
     for {
-      method <- protocol
+      method <- getProtocol()
         .resolveMethod(methodName)
-        .toRight(Errors.MethodNotFound)
-      decoder <- protocol
+        .toRight(protocolFactory.onMissingMethod())
+      decoder <- getProtocol()
         .getParamsDecoder(method)
         .toRight(Errors.InvalidRequest)
     } yield decoder
+  }
 }
 
 /** Control messages for the [[MessageHandler]] actor.

lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandlerSupervisor.scala

@@ -14,12 +14,12 @@ import java.util.UUID
 
 /** An actor responsible for supervising the [[MessageHandler]].
   *
-  * @param protocol a protocol supported be the server
+  * @param protocolFactory a factory used to create a protocol supported be the server
   * @param clientControllerFactory a factory used to create a client controller
   */
 final class MessageHandlerSupervisor(
   clientControllerFactory: ClientControllerFactory,
-  protocol: Protocol
+  protocolFactory: ProtocolFactory
 ) extends Actor
     with LazyLogging
     with Stash {
@@ -55,7 +55,7 @@ final class MessageHandlerSupervisor(
 
       val messageHandler =
         context.actorOf(
-          Props(new MessageHandler(protocol, clientActor)),
+          Props(new MessageHandler(protocolFactory, clientActor)),
           s"message-handler-$clientId"
         )
       clientActor ! JsonRpcServer.WebConnect(messageHandler)

lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/Protocol.scala

@@ -132,7 +132,8 @@ object Protocol {
       customErrors   = Map(),
       payloadsEncoder = { payload =>
         throw InexhaustivePayloadsSerializerError(payload)
-      }
+      },
+      initialized = false
     )
 }
 
@@ -185,13 +186,15 @@ class ResultDecoder[+M <: Method, +Result](method: M)(implicit
   * @param customErrors custom datatypes used for error codes.
   * @param payloadsEncoder an encoder for any payload (i.e. params or results)
   *                        used within this protocol.
+  * @param initialized indicates if the protocol has been fully initialized with all supported methods
   */
 case class Protocol(
   methods: Set[Method],
   paramsDecoders: Map[Method, ParamsDecoder[Method, Any]],
   resultDecoders: Map[Method, ResultDecoder[Method, Any]],
   customErrors: Map[Int, Error],
-  payloadsEncoder: Encoder[Any]
+  payloadsEncoder: Encoder[Any],
+  initialized: Boolean
 ) {
 
   private val builtinErrors: Map[Int, Error] = List(
@@ -292,6 +295,9 @@ case class Protocol(
       }
     )
 
+  /** Indicates that the protocol has registered all supported methods and notifications */
+  def finalized(): Protocol = copy(initialized = true)
+
   /** Adds a new error code to this protocol.
     * @param error the error to add.
     * @return a new [[Protocol]], recognizing `error` code.

lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/ProtocolFactory.scala

@@ -3,9 +3,18 @@ package org.enso.jsonrpc
 /** Factory that creates [[Protocol]]. */
 trait ProtocolFactory {
 
-  /** @return the [[Protocol]] instance. */
-  def getProtocol: Protocol
+  /** Returns the [[Protocol]] instance.
+    * If the factory has not been properly initialized yet, returns only a minimal set of messages
+    * supported during the initialization period. Returns a full set of supported messages in the
+    * post-initialization stage.
+    *
+    * @return the [[Protocol]] instance.
+    */
+  def getProtocol(): Protocol
 
-  /** Initialize the protocol. */
+  /** Initialize the protocol with the full set of supported messages. */
   def init(): Unit
+
+  /** Error returned when a requested method is not recognized */
+  def onMissingMethod(): Error
 }

lib/scala/json-rpc-server/src/test/scala/org/enso/jsonrpc/MessageHandlerSpec.scala

@@ -66,15 +66,23 @@ class MessageHandlerSpec
 
   case object MyError extends Error(15, "Test error")
 
-  object MyProtocol {
+  object MyProtocolFactory extends ProtocolFactory {
     import io.circe.generic.auto._
 
-    val protocol: Protocol = Protocol.empty
+    private val protocol: Protocol = Protocol.empty
       .registerNotification(MyNotification)
       .registerNotification(MyEmptyNotification)
       .registerRequest(MyRequest)
       .registerRequest(MyEmptyRequest)
       .registerError(MyError)
+      .finalized()
+
+    override def getProtocol(): Protocol = protocol
+
+    override def init(): Unit = ()
+
+    /** Error returned when a requested method is not recognized */
+    override def onMissingMethod(): Error = Errors.MethodNotFound
   }
 
   var out: TestProbe        = _
@@ -85,7 +93,7 @@ class MessageHandlerSpec
     out        = TestProbe()
     controller = TestProbe()
     handler = system.actorOf(
-      Props(new MessageHandler(MyProtocol.protocol, controller.ref))
+      Props(new MessageHandler(MyProtocolFactory, controller.ref))
     )
     handler ! Connected(out.ref)
   }

lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpc.scala

@@ -34,5 +34,6 @@ object JsonRpc {
       .registerRequest(ConfigSet)
       .registerRequest(ConfigDelete)
       .registerRequest(LoggingServiceGetEndpoint)
+      .finalized()
 
 }

lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpcProtocolFactory.scala

@@ -1,16 +1,20 @@
 package org.enso.projectmanager.protocol
 
-import org.enso.jsonrpc.{Protocol, ProtocolFactory}
+import org.enso.jsonrpc
+import org.enso.jsonrpc.{Errors, Protocol, ProtocolFactory}
 
 /** Factory creating JSON-RPC protocol. */
 final class JsonRpcProtocolFactory extends ProtocolFactory {
 
   /** @inheritdoc */
-  override def getProtocol: Protocol =
+  override def getProtocol(): Protocol =
     JsonRpc.protocol
 
   /** @inheritdoc */
   override def init(): Unit = {
     val _ = JsonRpc.protocol
   }
+
+  /** Error returned when a requested method is not recognized */
+  override def onMissingMethod(): jsonrpc.Error = Errors.MethodNotFound
 }