Basic SSE infrastructure

2026-08-04 22:49:04 +02:00
parent cd8bbe5687
commit 0e57de60b0
7 changed files with 206 additions and 16 deletions
@@ -1,11 +1,14 @@
 package com.jaytux.phoebench.server

 import com.jaytux.phoebench.common.ErrorResponse
+import com.jaytux.phoebench.common.HomeResponse
 import com.jaytux.phoebench.common.Routes
+import com.jaytux.phoebench.common.fold
 import com.jaytux.phoebench.server.db.DB
 import com.jaytux.phoebench.server.db.User
 import com.jaytux.phoebench.server.handlers.AuthHandler
 import com.jaytux.phoebench.server.handlers.ProjectHandler
+import com.jaytux.phoebench.server.handlers.RouteError
 import com.jaytux.phoebench.server.handlers.deleteAdmin
 import com.jaytux.phoebench.server.handlers.deleteAuth
 import com.jaytux.phoebench.server.handlers.get
@@ -30,10 +33,18 @@ import io.ktor.server.plugins.statuspages.StatusPages
 import io.ktor.server.request.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
+import io.ktor.server.sse.SSE
+import io.ktor.server.sse.sse
+import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
+import kotlinx.serialization.serializer
 import org.jetbrains.exposed.v1.jdbc.transactions.transaction
 import java.net.URI
 import kotlin.uuid.Uuid
+import com.jaytux.phoebench.server.handlers.RouteError.Companion.respondJson
+import io.ktor.server.sse.heartbeat
+import io.ktor.sse.ServerSentEvent
+import io.ktor.utils.io.CancellationException

 fun main(args: Array<String>) {
    DB.db
@@ -54,6 +65,10 @@ fun Application.module() {
        }
    }

+    install(SSE) {
+        //
+    }
+
    val allowLocalhost = environment.config.propertyOrNull("ktor.cors.enableLocalhostOn")?.getString() ?: "0"
    val safeOrigin = environment.config.propertyOrNull("ktor.cors.browserOrigin")?.getString()
    install(CORS) {
@@ -106,7 +121,6 @@ fun Application.module() {
            }

            challenge { defaultScheme, realm ->
-//                call.respond(HttpStatusCode.Unauthorized, ErrorResponse("Invalid/expired token"))
                call.respondText(Json.encodeToString(ErrorResponse("Invalid/expired token")), ContentType.Application.Json, HttpStatusCode.Unauthorized)
            }
        }
@@ -144,6 +158,57 @@ fun Application.module() {
            postAuth(Routes.Entry.new, ProjectHandler::createEntry)
            patchAuth(Routes.Entry.update, ProjectHandler::updateEntry)
            deleteAuth(Routes.Entry.delete, ProjectHandler::deleteEntry)
+
+            val projectBus = SSEBus.register<HomeResponse.ProjectSummary>("/rt/project")
+            sse("/rt/project") {
+                try {
+                    println("Attempt to set up SSE")
+                    val principal = call.principal<JWTPrincipal>()
+                    val userId = principal?.payload?.getClaim(com.jaytux.phoebench.common.Auth.JWT_CLAIM)?.asString()
+                        ?: throw RouteError("Missing user claim", HttpStatusCode.Unauthorized)
+                    println(" -- SSE: userId = $userId")
+                    val user = transaction {
+                        User.findById(Uuid.parse(userId)) ?: throw RouteError(
+                            "User not found",
+                            HttpStatusCode.Unauthorized
+                        )
+                    }
+                    println(" -- SSE: user = ${user.username}")
+
+                    heartbeat {}
+
+                    val flow = projectBus.register(user.id.value)
+
+                    try {
+                        flow.collect {
+                            it.fold({
+                                throw Exception() // force end of collecting
+                            }) { event ->
+                                send(ServerSentEvent(data = Json.encodeToString(projectBus.serializer, event)))
+                            }
+                        }
+                    }
+                    catch(e: CancellationException) {
+                        projectBus.disconnect(user.id.value)
+                        throw e
+                    }
+                }
+                catch(e: RouteError) {
+                    call.respondText(
+                        status = e.status,
+                        text = Json.encodeToString(ErrorResponse(e.message ?: "Unknown error")),
+                        contentType = ContentType.Application.Json
+                    )
+                }
+
+                println("--- SSE Session ended ---")
+            }
+
+            sse("/rt/users") {}
+
+            sse("/rt/invites") {}
+
+            sse("/rt/project/{project-id}") {}
        }

        get("{...}") {
@@ -0,0 +1,75 @@
+package com.jaytux.phoebench.server
+
+import com.jaytux.phoebench.common.Either
+import com.jaytux.phoebench.common.error
+import com.jaytux.phoebench.common.value
+import io.ktor.util.reflect.instanceOf
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.SharedFlow
+import kotlinx.serialization.KSerializer
+import kotlinx.serialization.serializer
+import java.util.concurrent.ConcurrentHashMap
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
+import kotlin.uuid.Uuid
+
+class SSEBus<T> private constructor (private val _containedType: KType, val serializer: KSerializer<T>) {
+    object Cancellation
+
+    private val _flows = ConcurrentHashMap<Uuid, Pair<Int, MutableSharedFlow<Either<Cancellation, T>>>>()
+
+    fun register(user: Uuid): SharedFlow<Either<Cancellation, T>> =
+        _flows.computeIfAbsent(user) { 1 to MutableSharedFlow(extraBufferCapacity = 64) }.second
+
+    suspend fun send(user: Uuid, data: T) {
+        _flows[user]?.second?.emit(data.value())
+    }
+
+    suspend fun sendAll(data: T) {
+        _flows.forEach { it.value.second.emit(data.value()) }
+    }
+
+    fun disconnect(user: Uuid) {
+        _flows.compute(user) { _, data ->
+            if(data == null) null
+            else {
+                val (refCount, flow) = data
+                val newRef = refCount - 1
+                if(newRef == 0) null
+                else newRef to flow
+            }
+        }
+    }
+
+    suspend fun forceDisconnect(user: Uuid) {
+        _flows[user]?.second?.emit(Cancellation.error())
+        _flows.remove(user)
+    }
+
+    companion object {
+        private val _busCache = mutableMapOf<String, SSEBus<*>>()
+
+        fun <T> register(topic: String, contained: KType, serializer: KSerializer<T>): SSEBus<T> {
+            val bus = SSEBus<T>(contained, serializer)
+            _busCache.compute(topic) { k, existing ->
+                if(existing != null) throw IllegalArgumentException("Bus for $topic exists already")
+                bus
+            }
+            return bus
+        }
+
+        inline fun <reified T> register(topic: String) =
+            register<T>(topic, typeOf<T>(), serializer<T>())
+
+        fun <T> getBus(topic: String, contained: KType): SSEBus<T>? {
+            val bus = _busCache[topic] ?: return null
+            if(contained != bus._containedType) throw IllegalArgumentException("Type mismatch for bus for $topic")
+
+            @Suppress("UNCHECKED_CAST")
+            return bus as SSEBus<T>
+        }
+
+        inline fun <reified T> getBus(topic: String) =
+            getBus<T>(topic, typeOf<T>())
+    }
+}
@@ -13,6 +13,7 @@ import com.jaytux.phoebench.common.PartialLabelRequest
 import com.jaytux.phoebench.common.PartialProjectRequest
 import com.jaytux.phoebench.common.ProjectRequest
 import com.jaytux.phoebench.common.ProjectResponse
+import com.jaytux.phoebench.server.SSEBus
 import com.jaytux.phoebench.server.db.Entries
 import com.jaytux.phoebench.server.db.Entry
 import com.jaytux.phoebench.server.db.Label
@@ -64,7 +65,7 @@ object ProjectHandler {
        success(HomeResponse(user.username, user.isAdmin, user.projectLimit, own, publics))
    }

-    fun createProject(user: User, req: ProjectRequest) = transaction {
+    suspend fun createProject(user: User, req: ProjectRequest) = transaction {
        if(user.projectLimit != -1 && (user.projects.count() >= user.projectLimit))
            throw RouteError("You have reached your project limit.", HttpStatusCode.Conflict)

@@ -75,6 +76,11 @@ object ProjectHandler {
        }

        success(proj.toResponse(user))
+    }.also { (_, proj) ->
+        val bus = SSEBus.getBus<HomeResponse.ProjectSummary>("/rt/project") ?: return@also
+        val summary = HomeResponse.ProjectSummary(proj.id, proj.name, proj.isPublic, NamedID(user.username, user.id.value))
+        if(proj.isPublic) bus.sendAll(summary)
+        else bus.send(user.id.value, summary)
    }

    fun getProject(user: User, req: Uuid) = transaction {