Merge pull request #646 from modelix/use-respondBytesWriter

chore(model-server): use non-blocking method to send version delta

model-server/src/main/kotlin/org/modelix/model/server/Main.kt

@@ -48,6 +48,7 @@ import org.apache.commons.io.FileUtils
 import org.apache.ignite.Ignition
 import org.modelix.authorization.KeycloakUtils
 import org.modelix.authorization.installAuthentication
+import org.modelix.model.InMemoryModels
 import org.modelix.model.server.handlers.ContentExplorer
 import org.modelix.model.server.handlers.DeprecatedLightModelServer
 import org.modelix.model.server.handlers.HistoryHandler
@@ -150,8 +151,9 @@ object Main {
                 i += 2
             }
             val localModelClient = LocalModelClient(storeClient)
+            val inMemoryModels = InMemoryModels()
             val repositoriesManager = RepositoriesManager(localModelClient)
-            val modelServer = KeyValueLikeModelServer(repositoriesManager)
+            val modelServer = KeyValueLikeModelServer(repositoriesManager, storeClient, inMemoryModels)
             val sharedSecretFile = cmdLineArgs.secretFile
             if (sharedSecretFile.exists()) {
                 modelServer.setSharedSecret(
@@ -162,7 +164,7 @@ object Main {
             val repositoryOverview = RepositoryOverview(repositoriesManager)
             val historyHandler = HistoryHandler(localModelClient, repositoriesManager)
             val contentExplorer = ContentExplorer(localModelClient, repositoriesManager)
-            val modelReplicationServer = ModelReplicationServer(repositoriesManager)
+            val modelReplicationServer = ModelReplicationServer(repositoriesManager, localModelClient, inMemoryModels)
             val metricsHandler = MetricsHandler()
 
             val configureNetty: NettyApplicationEngine.Configuration.() -> Unit = {

model-server/src/main/kotlin/org/modelix/model/server/handlers/ContentExplorer.kt

@@ -50,7 +50,7 @@ import org.modelix.model.lazy.RepositoryId
 import org.modelix.model.server.templates.PageWithMenuBar
 import kotlin.collections.set
 
-class ContentExplorer(private val client: IModelClient, private val repoManager: RepositoriesManager) {
+class ContentExplorer(private val client: IModelClient, private val repoManager: IRepositoriesManager) {
 
     fun init(application: Application) {
         application.routing {

model-server/src/main/kotlin/org/modelix/model/server/handlers/HistoryHandler.kt

@@ -55,7 +55,7 @@ import org.modelix.model.server.templates.PageWithMenuBar
 import java.time.LocalDateTime
 import java.time.format.DateTimeFormatter
 
-class HistoryHandler(val client: IModelClient, private val repositoriesManager: RepositoriesManager) {
+class HistoryHandler(val client: IModelClient, private val repositoriesManager: IRepositoriesManager) {
 
     fun init(application: Application) {
         application.routing {

model-server/src/main/kotlin/org/modelix/model/server/handlers/IRepositoriesManager.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2024.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.modelix.model.server.handlers
+
+import org.modelix.model.lazy.BranchReference
+import org.modelix.model.lazy.CLVersion
+import org.modelix.model.lazy.RepositoryId
+
+interface IRepositoriesManager {
+    /**
+     * Used to retrieve the server ID. If needed, the server ID is created and stored.
+     *
+     * If a server ID was not created yet, it is generated and saved in the database.
+     * It gets stored under the current and all legacy database keys.
+     *
+     * If the server ID was created previously but is only stored under a legacy database key,
+     * it also gets stored under the current and all legacy database keys.
+     */
+    suspend fun maybeInitAndGetSeverId(): String
+    fun getRepositories(): Set<RepositoryId>
+    suspend fun createRepository(repositoryId: RepositoryId, userName: String?, useRoleIds: Boolean = true): CLVersion
+    suspend fun removeRepository(repository: RepositoryId): Boolean
+
+    fun getBranches(repositoryId: RepositoryId): Set<BranchReference>
+
+    suspend fun removeBranches(repository: RepositoryId, branchNames: Set<String>)
+
+    /**
+     * Same as [removeBranches] but blocking.
+     * Caller is expected to execute it outside the request thread.
+     */
+    fun removeBranchesBlocking(repository: RepositoryId, branchNames: Set<String>)
+    suspend fun getVersion(branch: BranchReference): CLVersion?
+    suspend fun getVersionHash(branch: BranchReference): String?
+    suspend fun pollVersionHash(branch: BranchReference, lastKnown: String?): String
+    suspend fun mergeChanges(branch: BranchReference, newVersionHash: String): String
+
+    /**
+     * Same as [mergeChanges] but blocking.
+     * Caller is expected to execute it outside the request thread.
+     */
+    fun mergeChangesBlocking(branch: BranchReference, newVersionHash: String): String
+    suspend fun computeDelta(versionHash: String, baseVersionHash: String?): ObjectData
+}
+
+fun IRepositoriesManager.getBranchNames(repositoryId: RepositoryId): Set<String> {
+    return getBranches(repositoryId).map { it.branchName }.toSet()
+}

model-server/src/main/kotlin/org/modelix/model/server/handlers/KeyValueLikeModelServer.kt

@@ -44,6 +44,7 @@ import org.modelix.authorization.checkPermission
 import org.modelix.authorization.getUserName
 import org.modelix.authorization.requiresPermission
 import org.modelix.authorization.toKeycloakScope
+import org.modelix.model.InMemoryModels
 import org.modelix.model.lazy.BranchReference
 import org.modelix.model.lazy.RepositoryId
 import org.modelix.model.persistent.HashUtil
@@ -67,16 +68,22 @@ private class NotFoundException(description: String?) : RuntimeException(descrip
 
 typealias CallContext = PipelineContext<Unit, ApplicationCall>
 
-class KeyValueLikeModelServer(val repositoriesManager: RepositoriesManager) {
+class KeyValueLikeModelServer(
+    private val repositoriesManager: IRepositoriesManager,
+    private val storeClient: IStoreClient,
+    private val inMemoryModels: InMemoryModels,
+) {
+
+    constructor(repositoriesManager: RepositoriesManager) :
+        this(repositoriesManager, repositoriesManager.client.store, InMemoryModels())
+
     companion object {
         private val LOG = LoggerFactory.getLogger(KeyValueLikeModelServer::class.java)
         val HASH_PATTERN = Pattern.compile("[a-zA-Z0-9\\-_]{5}\\*[a-zA-Z0-9\\-_]{38}")
         const val PROTECTED_PREFIX = "$$$"
         val HEALTH_KEY = PROTECTED_PREFIX + "health2"
     }
 
-    val storeClient: IStoreClient get() = repositoriesManager.client.store
-
     fun init(application: Application) {
         // Functionally, it does not matter if the server ID
         // is created eagerly on startup or lazily on the first request,
@@ -100,7 +107,7 @@ class KeyValueLikeModelServer(val repositoriesManager: RepositoriesManager) {
                     ?.getBranchReference(System.getenv("MODELIX_SERVER_MODELQL_WARMUP_BRANCH"))
                 if (branchRef != null) {
                     val version = repositoriesManager.getVersion(branchRef)
-                    if (repositoriesManager.inMemoryModels.getModel(version!!.getTree()).isActive) {
+                    if (inMemoryModels.getModel(version!!.getTree()).isActive) {
                         call.respondText(
                             status = HttpStatusCode.ServiceUnavailable,
                             text = "Waiting for version $version to be loaded into memory",
@@ -346,7 +353,7 @@ class KeyValueLikeModelServer(val repositoriesManager: RepositoriesManager) {
 
         HashUtil.checkObjectHashes(hashedObjects)
 
-        repositoriesManager.client.store.runTransactionSuspendable {
+        storeClient.runTransactionSuspendable {
             storeClient.putAll(hashedObjects)
             storeClient.putAll(userDefinedEntries)
             for ((branch, value) in branchChanges) {

model-server/src/main/kotlin/org/modelix/model/server/handlers/ModelReplicationServer.kt

@@ -27,13 +27,17 @@ import io.ktor.server.resources.get
 import io.ktor.server.resources.post
 import io.ktor.server.resources.put
 import io.ktor.server.response.respond
+import io.ktor.server.response.respondBytesWriter
 import io.ktor.server.response.respondText
-import io.ktor.server.response.respondTextWriter
 import io.ktor.server.routing.Route
 import io.ktor.server.routing.route
 import io.ktor.server.routing.routing
 import io.ktor.server.websocket.webSocket
+import io.ktor.util.cio.use
 import io.ktor.util.pipeline.PipelineContext
+import io.ktor.utils.io.ByteWriteChannel
+import io.ktor.utils.io.close
+import io.ktor.utils.io.writeStringUtf8
 import io.ktor.websocket.send
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
@@ -46,6 +50,7 @@ import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 import org.modelix.api.public.Paths
 import org.modelix.authorization.getUserName
+import org.modelix.model.InMemoryModels
 import org.modelix.model.api.ITree
 import org.modelix.model.api.PBranch
 import org.modelix.model.api.TreePointer
@@ -69,15 +74,21 @@ import org.slf4j.LoggerFactory
  * Implements the endpoints used by the 'model-client', but compared to KeyValueLikeModelServer also understands what
  * client sends. This allows more validations and more responsibilities on the server side.
  */
-class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
-    constructor(modelClient: LocalModelClient) : this(RepositoriesManager(modelClient))
+class ModelReplicationServer(
+    private val repositoriesManager: IRepositoriesManager,
+    private val modelClient: LocalModelClient,
+    private val inMemoryModels: InMemoryModels,
+) {
+    constructor(repositoriesManager: RepositoriesManager) :
+        this(repositoriesManager, repositoriesManager.client, InMemoryModels())
+
+    constructor(modelClient: LocalModelClient) : this(RepositoriesManager(modelClient), modelClient, InMemoryModels())
     constructor(storeClient: IStoreClient) : this(LocalModelClient(storeClient))
 
     companion object {
         private val LOG = LoggerFactory.getLogger(ModelReplicationServer::class.java)
     }
 
-    private val modelClient: LocalModelClient get() = repositoriesManager.client
     private val storeClient: IStoreClient get() = modelClient.store
 
     fun init(application: Application) {
@@ -268,16 +279,16 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
             LOG.trace("Running query on {} @ {}", branchRef, version)
             val initialTree = version!!.getTree()
             val branch = OTBranch(
-                PBranch(initialTree, repositoriesManager.client.idGenerator),
-                repositoriesManager.client.idGenerator,
-                repositoriesManager.client.storeCache,
+                PBranch(initialTree, modelClient.idGenerator),
+                modelClient.idGenerator,
+                modelClient.storeCache,
             )
 
             ModelQLServer.handleCall(call, { writeAccess ->
                 if (writeAccess) {
                     branch.getRootNode() to branch.getArea()
                 } else {
-                    val model = repositoriesManager.inMemoryModels.getModel(initialTree).await()
+                    val model = inMemoryModels.getModel(initialTree).await()
                     model.getNode(ITree.ROOT_ID) to model.getArea()
                 }
             }, {
@@ -286,7 +297,7 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
                 val (ops, newTree) = branch.getPendingChanges()
                 if (newTree != initialTree) {
                     val newVersion = CLVersion.createRegularVersion(
-                        id = repositoriesManager.client.idGenerator.generate(),
+                        id = modelClient.idGenerator.generate(),
                         author = getUserName(),
                         tree = newTree as CLTree,
                         baseVersion = version,
@@ -299,7 +310,7 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
 
         post<Paths.postRepositoryVersionHashQuery> {
             val versionHash = call.parameters["versionHash"]!!
-            val version = CLVersion.loadFromHash(versionHash, repositoriesManager.client.storeCache)
+            val version = CLVersion.loadFromHash(versionHash, modelClient.storeCache)
             val initialTree = version.getTree()
             val branch = TreePointer(initialTree)
             ModelQLServer.handleCall(call, branch.getRootNode(), branch.getArea())
@@ -372,21 +383,47 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
         respond(delta)
     }
 
-    private suspend fun ApplicationCall.respondDeltaAsObjectStream(versionHash: String, baseVersionHash: String?, plainText: Boolean) {
-        respondTextWriter(contentType = if (plainText) ContentType.Text.Plain else VersionDeltaStream.CONTENT_TYPE) {
-            repositoriesManager.computeDelta(versionHash, baseVersionHash).asFlow()
-                .flatten()
-                .withSeparator("\n")
-                .onEmpty { emit(versionHash) }
-                .withIndex()
-                .collect {
-                    if (it.index == 0) check(it.value == versionHash) { "First object should be the version" }
-                    append(it.value)
-                }
+    private suspend fun ApplicationCall.respondDeltaAsObjectStream(
+        versionHash: String,
+        baseVersionHash: String?,
+        plainText: Boolean,
+    ) {
+        // Call `computeDelta` before starting to respond.
+        // It could already throw an exception, and in that case we do not want a successful response status.
+        val objectData = repositoriesManager.computeDelta(versionHash, baseVersionHash)
+        val contentType = if (plainText) ContentType.Text.Plain else VersionDeltaStream.CONTENT_TYPE
+        respondBytesWriter(contentType) {
+            this.useClosingWithoutCause {
+                objectData.asFlow()
+                    .flatten()
+                    .withSeparator("\n")
+                    .onEmpty { emit(versionHash) }
+                    .withIndex()
+                    .collect {
+                        if (it.index == 0) check(it.value == versionHash) { "First object should be the version" }
+                        writeStringUtf8(it.value)
+                    }
+            }
         }
     }
 }
 
+/**
+ * Same as [[ByteWriteChannel.use]] but closing without a cause in case of an exception.
+ *
+ * Calling [[ByteWriteChannel.close]] with a cause results in not closing the connection properly.
+ * See ModelReplicationServerTest.`server closes connection when failing to compute delta after starting to respond`
+ * This will only be fixed in Ktor 3.
+ * See https://youtrack.jetbrains.com/issue/KTOR-4862/Ktor-hangs-if-exception-occurs-during-write-response-body
+ */
+private inline fun ByteWriteChannel.useClosingWithoutCause(block: ByteWriteChannel.() -> Unit) {
+    try {
+        block()
+    } finally {
+        close()
+    }
+}
+
 private fun <T> Flow<Pair<T, T>>.flatten() = flow<T> {
     collect {
         emit(it.first)