From dd06c59380d8e93abf1aa68d3472e343b0f6c767 Mon Sep 17 00:00:00 2001 From: Hubert Plociniczak Date: Fri, 7 Jul 2023 13:49:23 +0200 Subject: [PATCH] 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) --- .../protocol/json/JsonRpc.scala | 11 +++++--- .../json/JsonRpcProtocolFactory.scala | 25 ++++++++++++++++--- .../jsonrpc/test/JsonRpcServerTestKit.scala | 4 ++- .../org/enso/jsonrpc/JsonRpcServer.scala | 5 ++-- .../org/enso/jsonrpc/MessageHandler.scala | 25 +++++++++++-------- .../jsonrpc/MessageHandlerSupervisor.scala | 6 ++--- .../scala/org/enso/jsonrpc/Protocol.scala | 10 ++++++-- .../org/enso/jsonrpc/ProtocolFactory.scala | 15 ++++++++--- .../org/enso/jsonrpc/MessageHandlerSpec.scala | 14 ++++++++--- .../projectmanager/protocol/JsonRpc.scala | 1 + .../protocol/JsonRpcProtocolFactory.scala | 8 ++++-- 11 files changed, 89 insertions(+), 35 deletions(-) diff --git a/engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpc.scala b/engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpc.scala index a3746346db..450766c8c7 100644 --- a/engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpc.scala +++ b/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() } diff --git a/engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpcProtocolFactory.scala b/engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpcProtocolFactory.scala index 1814b1b6e5..c3a5850d2a 100644 --- a/engine/language-server/src/main/scala/org/enso/languageserver/protocol/json/JsonRpcProtocolFactory.scala +++ b/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 } } diff --git a/lib/scala/json-rpc-server-test/src/main/scala/org/enso/jsonrpc/test/JsonRpcServerTestKit.scala b/lib/scala/json-rpc-server-test/src/main/scala/org/enso/jsonrpc/test/JsonRpcServerTestKit.scala index e26dc4a948..1cf68a1f1a 100644 --- a/lib/scala/json-rpc-server-test/src/main/scala/org/enso/jsonrpc/test/JsonRpcServerTestKit.scala +++ b/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}" } diff --git a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/JsonRpcServer.scala b/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/JsonRpcServer.scala index f9af9a2398..6c8e02408b 100644 --- a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/JsonRpcServer.scala +++ b/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()}" ) diff --git a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandler.scala b/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandler.scala index 34336b63a3..196df7a031 100644 --- a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandler.scala +++ b/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. diff --git a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandlerSupervisor.scala b/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandlerSupervisor.scala index 7db5bd951b..90a1fedfb1 100644 --- a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/MessageHandlerSupervisor.scala +++ b/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) diff --git a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/Protocol.scala b/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/Protocol.scala index 9126137dd8..02bf21d24f 100644 --- a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/Protocol.scala +++ b/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. diff --git a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/ProtocolFactory.scala b/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/ProtocolFactory.scala index 07c55f110a..8d3ea21658 100644 --- a/lib/scala/json-rpc-server/src/main/scala/org/enso/jsonrpc/ProtocolFactory.scala +++ b/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 } diff --git a/lib/scala/json-rpc-server/src/test/scala/org/enso/jsonrpc/MessageHandlerSpec.scala b/lib/scala/json-rpc-server/src/test/scala/org/enso/jsonrpc/MessageHandlerSpec.scala index a1adb2222a..64f993e7e3 100644 --- a/lib/scala/json-rpc-server/src/test/scala/org/enso/jsonrpc/MessageHandlerSpec.scala +++ b/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) } diff --git a/lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpc.scala b/lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpc.scala index f2aa8e90a9..fc7b887efe 100644 --- a/lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpc.scala +++ b/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() } diff --git a/lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpcProtocolFactory.scala b/lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpcProtocolFactory.scala index 17780b85e8..96a6ef6020 100644 --- a/lib/scala/project-manager/src/main/scala/org/enso/projectmanager/protocol/JsonRpcProtocolFactory.scala +++ b/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 }