Basic SSE infrastructure

clients/compose/src/commonMain/kotlin/com/jaytux/phoebench/clients/Client.kt

@@ -5,10 +5,14 @@ import io.ktor.client.*
 import io.ktor.client.plugins.auth.*
 import io.ktor.client.plugins.auth.providers.*
 import io.ktor.client.plugins.contentnegotiation.*
+import io.ktor.client.plugins.sse.*
 import io.ktor.serialization.kotlinx.json.*
 import io.ktor.utils.io.*
 import kotlinx.atomicfu.locks.ReentrantLock
 import kotlinx.atomicfu.locks.withLock
+import kotlinx.serialization.KSerializer
+import kotlinx.serialization.json.Json
+import kotlin.time.Duration.Companion.seconds
 
 class Client private constructor(private val _auth: AuthProvider) {
     private val _authClient = platformClient {
@@ -17,6 +21,13 @@ class Client private constructor(private val _auth: AuthProvider) {
     private val _client = platformClient {
         var tryingRefresh = false
         install(ContentNegotiation) { json() }
+
+        install(SSE) {
+            maxReconnectionAttempts = 10
+            reconnectionTime = 1.seconds
+            bufferPolicy = SSEBufferPolicy.LastEvents(5)
+        }
+
         install(Auth) {
             bearer {
                 cacheTokens = false
@@ -31,8 +42,7 @@ class Client private constructor(private val _auth: AuthProvider) {
                     if (tryingRefresh) {
                         _auth.onLogout()
                         null
-                    }
+                    } else {
-                    else {
                         tryingRefresh = true
                         val ref = _auth.refresh.value ?: return@refreshTokens null
                         println("Trying to re-authenticate using $ref")
@@ -57,27 +67,44 @@ class Client private constructor(private val _auth: AuthProvider) {
 
     private var _counter = 0
 
-    private suspend fun <TReq: Any, TRes: Any> callRoute(using: HttpClient, route: ApiRoute<TReq, TRes>, body: TReq, wasInternal: Boolean = false): Either<ErrorResponse, TRes> {
+    private suspend fun <TReq : Any, TRes : Any> callRoute(
+        using: HttpClient,
+        route: ApiRoute<TReq, TRes>,
+        body: TReq,
+        wasInternal: Boolean = false
+    ): Either<ErrorResponse, TRes> {
         val ctr = _counter++
         return try {
             println("Calling route ${route.pattern} with client $using, request ID=$ctr (internal request: $wasInternal)")
             val client = IClient.Default(using, _auth.server.value ?: throw IllegalStateException("No server URL set."))
             val res = route.call(client, body, ctr)
             res
-        }
+        } catch (e: CancellationException) {
-        catch(e: CancellationException) {
             println("Call to ${route.pattern} [$ctr] was cancelled")
             ErrorResponse(COROUTINE_CANCELLED).error()
-        }
+        } catch (e: Exception) {
-        catch(e: Exception) {
             println("Call to ${route.pattern} [$ctr] ran into an exception")
             ErrorResponse(e.message ?: "Unknown error.").error()
         }
     }
 
-    suspend fun <TReq: Any, TRes: Any> callRoute(route: ApiRoute<TReq, TRes>, body: TReq): Either<ErrorResponse, TRes> =
+    suspend fun <TReq : Any, TRes : Any> callRoute(
+        route: ApiRoute<TReq, TRes>,
+        body: TReq
+    ): Either<ErrorResponse, TRes> =
         callRoute(_client, route, body)
 
+    suspend fun <T> connectSSE(route: String, serializer: KSerializer<T>, onEvent: suspend (T) -> Unit) {
+        println("Client tries to set up SSE to $route")
+        val server = _auth.server.value ?: throw IllegalStateException("No server URL set.")
+        _client.sse(urlString = "$server$route", showCommentEvents = true, showRetryEvents = true) {
+            incoming.collect {
+                println("RECEIVE: $it")
+                if(it.data != null) onEvent(Json.decodeFromString(serializer, it.data!!))
+            }
+        }
+    }
+
     companion object {
         private val _sync = ReentrantLock()
         private var _instance: Client? = null
@@ -88,8 +115,7 @@ class Client private constructor(private val _auth: AuthProvider) {
                 if (_instance == null) {
                     _instance = Client(auth)
                     return _instance!!
-                }
+                } else throw IllegalStateException("Client already constructed")
-                else throw IllegalStateException("Client already constructed")
             }
         }
 

clients/compose/src/commonMain/kotlin/com/jaytux/phoebench/clients/data/HomeVM.kt

@@ -14,6 +14,8 @@ import com.jaytux.phoebench.clients.withScope
 import com.jaytux.phoebench.common.HomeResponse
 import com.jaytux.phoebench.common.InviteListResponse
 import com.jaytux.phoebench.common.UserListResponse
+import kotlinx.coroutines.Job
+import kotlinx.serialization.serializer
 import kotlin.uuid.Uuid
 
 class HomeVM(
@@ -42,6 +44,8 @@ class HomeVM(
     val invites = _invites.immutable()
     val users = _users.immutable()
 
+    private var _listenJob: Job? = null
+
     init {
         refresh()
     }
@@ -52,6 +56,7 @@ class HomeVM(
     }
 
     fun reset() {
+        _listenJob?.cancel()
         resetAdmin()
         _username.value = null
         _isAdmin.value = false
@@ -67,6 +72,17 @@ class HomeVM(
     }
 
     fun refresh() {
+        _listenJob = withScope {
+            _client.connectSSE("/rt/project", serializer<HomeResponse.ProjectSummary>()) {
+                if(it.owner.name == _username.value) {
+                    if(_ownProjects.value.none { p -> p.id == it.id }) _ownProjects.value += it
+                }
+                if(it.isPublic) {
+                    if(_publicProjects.value.none { p -> p.id == it.id }) _publicProjects.value += it
+                }
+            }
+        }
+
         withScope {
             resetAdmin()
             _repo.getHome().snackOr {

gradle/libs.versions.toml

@@ -62,6 +62,7 @@ ktor-server-auth = { module = "io.ktor:ktor-server-auth", version.ref = "ktor" }
 ktor-server-auth-jwt = { module = "io.ktor:ktor-server-auth-jwt", version.ref = "ktor" }
 ktor-server-status-pages = { module = "io.ktor:ktor-server-status-pages", version.ref = "ktor" }
 ktor-server-cors = { module = "io.ktor:ktor-server-cors", version.ref = "ktor" }
+ktor-server-sse = { module = "io.ktor:ktor-server-sse", version.ref = "ktor" }
 
 json = { module = "org.json:json", version.ref = "json" }
 kotlinx-datetime = { module = "org.jetbrains.kotlinx:kotlinx-datetime", version.ref = "datetime" }

server/build.gradle.kts

@@ -35,6 +35,7 @@ dependencies {
     implementation(libs.ktor.server.call.logging)
     implementation(libs.ktor.server.cors)
     implementation(libs.ktor.server.status.pages)
+    implementation(libs.ktor.server.sse)
 
     implementation(libs.ktor.serialization.kotlinx.json)
 

server/src/main/kotlin/com/jaytux/phoebench/server/Main.kt

@@ -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("{...}") {

server/src/main/kotlin/com/jaytux/phoebench/server/SSEBus.kt

@@ -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>())
+    }
+}

server/src/main/kotlin/com/jaytux/phoebench/server/handlers/ProjectHandler.kt

@@ -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 {