Merge pull request #332 from modelix/MODELIX-608

MODELIX-608 Health check for warming up the MPS cache

Author: Oleksandr Dzhychko

docs/global/modules/core/pages/reference/component-mps-model-server-plugin.adoc

@@ -9,9 +9,21 @@ https://api.modelix.org/3.12.0/mps-model-server-plugin/index.html[API doc^] | ht
 
 == Health checks
 
-The plugin offers a set of health checks via HTTP on port 48305 and path `/health`.
+The plugin offers a set of health checks via HTTP GET on port `48305` and path `/health`.
+
 Health checks can be enabled adding query parameters with the health check name and the value `true` to the request.
+Some health checks require further information that needs to be provided by query parameters.
+
+.Example of combining health check
+[source,text]
+----
+http(s)://<host>:48305/health?indexer=true&loadModels=true&loadModelsModuleNamespacePrefix=foo.bar <1> <2>
+----
+<.> `indexer=true` enables <<indexer>>
+<.> `loadModels=true` enables <<loadModels>>
+* `loadModelsModuleNamespacePrefix` is a parameter related to `loadModels`
 
+[#indexer]
 === indexer
 
 The check fails, if the indexer is currently running for one of the opened projects.
@@ -30,3 +42,26 @@ Reports an unhealthy system whenever no project is loaded.
 
 Reports an unhealthy system when no virtual folders are available.
 This might also be true in case a project without virtual folders is fully loaded.
+
+[#loadModels]
+=== loadModels
+
+Returns after trying to eagerly load a set of specified modules.
+This check can be used to avoid a slow first ModelQl query after launching an MPS instance running this plugin.
+
+[NOTE]
+This health check has the side effect of loading the model data.
+It does not just report whether the model data is loaded or not.
+It always tries to load model data before returning a result.
+
+Multiple `loadModelsModuleNamespacePrefix` parameters can be provided
+to specify the modules from which the models should be loaded.
+
+.Usage example
+[source,text]
+----
+http(s)://<host>:48305/health?loadModels=true&loadModelsModuleNamespacePrefix=org.foo&loadModelsModuleNamespacePrefix=org.bar <.> <.> <.>
+----
+<.> `loadModels=true` enables <<loadModels>>
+<.> `loadModelsModuleNamespacePrefix=org.foo` specifies to load all models from modules starting with `org.foo`
+<.> `loadModelsModuleNamespacePrefix=org.bar` specifies to load all models from modules starting with `org.bar`

model-server-lib/src/main/kotlin/org/modelix/model/server/light/LightModelServer.kt

@@ -35,6 +35,7 @@ import io.ktor.server.websocket.WebSockets
 import io.ktor.server.websocket.pingPeriod
 import io.ktor.server.websocket.timeout
 import io.ktor.server.websocket.webSocket
+import io.ktor.util.toMap
 import io.ktor.websocket.Frame
 import io.ktor.websocket.readText
 import io.ktor.websocket.send
@@ -208,22 +209,27 @@ class LightModelServer @JvmOverloads constructor(val port: Int, val rootNodeProv
                 try {
                     val allChecks = healthChecks.associateBy { it.id }.toMap()
                     val enabledChecks = allChecks.filter { it.value.enabledByDefault }.keys.toMutableSet()
-
-                    call.request.queryParameters.entries().forEach { entry ->
-                        entry.value.forEach { value ->
-                            if (!allChecks.containsKey(entry.key)) throw IllegalArgumentException("Unknown check: ${entry.key}")
-                            if (value.toBooleanStrict()) {
-                                enabledChecks.add(entry.key)
-                            } else {
-                                enabledChecks.remove(entry.key)
+                    val validParameterNames = (allChecks.keys + allChecks.flatMap { it.value.validParameterNames }).toSet()
+
+                    val queryParameters = call.request.queryParameters
+                    val queryParametersMap = queryParameters.toMap()
+                    queryParameters.entries().forEach { entry ->
+                        require(validParameterNames.contains(entry.key)) { "Unknown check: ${entry.key}" }
+                        if (allChecks.containsKey(entry.key)) {
+                            entry.value.forEach { value ->
+                                if (value.toBooleanStrict()) {
+                                    enabledChecks.add(entry.key)
+                                } else {
+                                    enabledChecks.remove(entry.key)
+                                }
                             }
                         }
                     }
                     var isHealthy = true
                     for (healthCheck in allChecks.values) {
                         if (enabledChecks.contains(healthCheck.id)) {
                             output.appendLine("--- running check '${healthCheck.id}' ---")
-                            val result = healthCheck.run(output)
+                            val result = healthCheck.run(output, queryParametersMap)
                             output.appendLine()
                             output.appendLine("-> " + if (result) "successful" else "failed")
                             isHealthy = isHealthy && result
@@ -442,9 +448,11 @@ class LightModelServer @JvmOverloads constructor(val port: Int, val rootNodeProv
     }
 
     interface IHealthCheck {
+        val validParameterNames: Set<String> get() = emptySet()
         val id: String
         val enabledByDefault: Boolean
         fun run(output: StringBuilder): Boolean
+        fun run(output: StringBuilder, parameters: Map<String, List<String>>): Boolean = run(output)
     }
 }
 

mps-model-server-plugin/src/main/kotlin/org/modelix/model/server/mps/MPSModelServer.kt

@@ -39,6 +39,8 @@ class MPSModelServerForProject(private val project: Project) : Disposable {
     }
 }
 
+const val LOAD_MODELS_MODULE_PARAMETER_NAME = "loadModelsModuleNamespacePrefix"
+
 @Service(Service.Level.APP)
 class MPSModelServer : Disposable {
 
@@ -128,6 +130,50 @@ class MPSModelServer : Disposable {
                         return false
                     }
                 })
+                .healthCheck(object : LightModelServer.IHealthCheck {
+                    // Usage example for this health check:
+                    // `/health?loadModels=true&loadModelsModuleNamespacePrefix=org.foo&loadModelsModuleNamespacePrefix=org.bar`
+                    // This should load all models from modules starting with either `org.foo` or `org.bar`.
+                    override val validParameterNames: Set<String> = setOf(LOAD_MODELS_MODULE_PARAMETER_NAME)
+                    override val id: String
+                        get() = "loadModels"
+                    override val enabledByDefault: Boolean
+                        get() = false
+
+                    override fun run(output: StringBuilder): Boolean {
+                        throw UnsupportedOperationException("parameters required")
+                    }
+
+                    override fun run(output: StringBuilder, parameters: Map<String, List<String>>): Boolean {
+                        val projects = getMPSProjects()
+                        val namespaces = parameters[LOAD_MODELS_MODULE_PARAMETER_NAME]?.toSet()
+                        if (namespaces == null) {
+                            output.append("parameter '$LOAD_MODELS_MODULE_PARAMETER_NAME' missing")
+                            return false
+                        }
+
+                        val project = projects.firstOrNull()
+                        if (project == null) {
+                            output.append("no projects loaded")
+                            return false
+                        }
+                        val repository = project.repository
+                        repository.modelAccess.runReadAction {
+                            val modules = repository.modules.filter {
+                                val moduleName = it.moduleName ?: return@filter false
+                                namespaces.any { namespace -> moduleName.startsWith(namespace) }
+                            }
+                            val models = modules.flatMap { it.models }
+                            models.forEach {
+                                // This triggers the loading of the model data which would otherwise slow down the
+                                // first query.
+                                it.rootNodes
+                            }
+                            output.append("${models.size} models in ${modules.size} modules loaded")
+                        }
+                        return true
+                    }
+                })
                 .build()
             server!!.start()
         }