Repository: tipsy/javalin
Pull request #1536 (2022-05-06 09:19:10)

GH-1532 Support SseClient#close


Patch:
diff --git a/javalin/src/main/java/io/javalin/http/HttpResponseException.kt b/javalin/src/main/java/io/javalin/http/HttpResponseException.kt
index a224e8f7..f9abb065 100644
--- a/javalin/src/main/java/io/javalin/http/HttpResponseException.kt
+++ b/javalin/src/main/java/io/javalin/http/HttpResponseException.kt
@@ -6,12 +6,10 @@
 
 package io.javalin.http
 
-import org.eclipse.jetty.http.HttpStatus
-
 open class HttpResponseException @JvmOverloads constructor(val status: Int, message: String, val details: Map<String, String> = mapOf()) : RuntimeException(message)
 
 class RedirectResponse @JvmOverloads constructor(
-    status: Int = HttpStatus.FOUND_302,
+    status: Int = HttpCode.FOUND.status,
     message: String = "Redirected",
     details: Map<String, String> = mapOf()
 ) : HttpResponseException(status, message, details)
@@ -20,54 +18,54 @@ class RedirectResponse @JvmOverloads constructor(
 class BadRequestResponse @JvmOverloads constructor(
     message: String = "Bad request",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.BAD_REQUEST_400, message, details)
+) : HttpResponseException(HttpCode.BAD_REQUEST.status, message, details)
 
 class UnauthorizedResponse @JvmOverloads constructor(
     message: String = "Unauthorized",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.UNAUTHORIZED_401, message, details)
+) : HttpResponseException(HttpCode.UNAUTHORIZED.status, message, details)
 
 class ForbiddenResponse @JvmOverloads constructor(
     message: String = "Forbidden",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.FORBIDDEN_403, message, details)
+) : HttpResponseException(HttpCode.FORBIDDEN.status, message, details)
 
 class NotFoundResponse @JvmOverloads constructor(
     message: String = "Not found",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.NOT_FOUND_404, message, details)
+) : HttpResponseException(HttpCode.NOT_FOUND.status, message, details)
 
 class MethodNotAllowedResponse @JvmOverloads constructor(
     message: String = "Method not allowed",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.METHOD_NOT_ALLOWED_405, message, details)
+) : HttpResponseException(HttpCode.METHOD_NOT_ALLOWED.status, message, details)
 
 class ConflictResponse @JvmOverloads constructor(
     message: String = "Conflict",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.CONFLICT_409, message, details)
+) : HttpResponseException(HttpCode.CONFLICT.status, message, details)
 
 class GoneResponse @JvmOverloads constructor(
     message: String = "Gone",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.GONE_410, message, details)
+) : HttpResponseException(HttpCode.GONE.status, message, details)
 
 class InternalServerErrorResponse @JvmOverloads constructor(
     message: String = "Internal server error",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.INTERNAL_SERVER_ERROR_500, message, details)
+) : HttpResponseException(HttpCode.INTERNAL_SERVER_ERROR.status, message, details)
 
 class BadGatewayResponse @JvmOverloads constructor(
     message: String = "Bad gateway",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.BAD_GATEWAY_502, message, details)
+) : HttpResponseException(HttpCode.BAD_GATEWAY.status, message, details)
 
 class ServiceUnavailableResponse @JvmOverloads constructor(
     message: String = "Service unavailable",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.SERVICE_UNAVAILABLE_503, message, details)
+) : HttpResponseException(HttpCode.SERVICE_UNAVAILABLE.status, message, details)
 
 class GatewayTimeoutResponse @JvmOverloads constructor(
     message: String = "Gateway timeout",
     details: Map<String, String> = mapOf()
-) : HttpResponseException(HttpStatus.GATEWAY_TIMEOUT_504, message, details)
+) : HttpResponseException(HttpCode.GATEWAY_TIMEOUT.status, message, details)
diff --git a/javalin/src/main/java/io/javalin/http/sse/SseClient.kt b/javalin/src/main/java/io/javalin/http/sse/SseClient.kt
index c3383b70..0a5c44e9 100644
--- a/javalin/src/main/java/io/javalin/http/sse/SseClient.kt
+++ b/javalin/src/main/java/io/javalin/http/sse/SseClient.kt
@@ -2,17 +2,23 @@ package io.javalin.http.sse
 
 import io.javalin.http.Context
 import io.javalin.plugin.json.jsonMapper
+import java.io.Closeable
 import java.io.InputStream
 
-class SseClient(@JvmField val ctx: Context) {
+class SseClient internal constructor(
+    private val closeSse: CloseSseFunction,
+    @JvmField val ctx: Context
+) : Closeable {
 
-    private val emitter: Emitter = Emitter(ctx.req.asyncContext)
-    private var closeCallback: Runnable = Runnable {}
+    private val emitter = Emitter(ctx.req.asyncContext)
+    private var closeCallback = Runnable {}
 
     fun onClose(closeCallback: Runnable) {
         this.closeCallback = closeCallback
     }
 
+    override fun close() = closeSse.close()
+
     fun sendEvent(data: Any) = sendEvent("message", data)
 
     @JvmOverloads
diff --git a/javalin/src/main/java/io/javalin/http/sse/SseHandler.kt b/javalin/src/main/java/io/javalin/http/sse/SseHandler.kt
index 13f3d092..e0a46f75 100644
--- a/javalin/src/main/java/io/javalin/http/sse/SseHandler.kt
+++ b/javalin/src/main/java/io/javalin/http/sse/SseHandler.kt
@@ -3,11 +3,14 @@ package io.javalin.http.sse
 import io.javalin.core.util.Header
 import io.javalin.http.Context
 import io.javalin.http.Handler
+import java.io.Closeable
 import java.util.concurrent.CompletableFuture
 import java.util.function.Consumer
 import javax.servlet.AsyncEvent
 import javax.servlet.AsyncListener
 
+internal fun interface CloseSseFunction : Closeable
+
 class SseHandler @JvmOverloads constructor(
     private val timeout: Long = 0,
     private val clientConsumer: Consumer<SseClient>
@@ -27,16 +30,18 @@ class SseHandler @JvmOverloads constructor(
             ctx.req.startAsync(ctx.req, ctx.res)
             ctx.req.asyncContext.timeout = timeout
 
-            val awaitFuture = CompletableFuture<Void>()
+            val closeSse = CompletableFuture<Void>()
+                .also { ctx.future(it) { /* do nothing with the future result in callback */ } }
+                .let { CloseSseFunction { it.complete(null) } }
+
             ctx.req.asyncContext.addListener(object : AsyncListener {
                 override fun onComplete(event: AsyncEvent) {}
                 override fun onStartAsync(event: AsyncEvent) {}
-                override fun onTimeout(event: AsyncEvent) { awaitFuture.complete(null) }
-                override fun onError(event: AsyncEvent) { awaitFuture.complete(null) }
+                override fun onTimeout(event: AsyncEvent) { closeSse.close() }
+                override fun onError(event: AsyncEvent) { closeSse.close() }
             })
-            ctx.future(awaitFuture) { /* do nothing with the future result in callback */ }
 
-            clientConsumer.accept(SseClient(ctx))
+            clientConsumer.accept(SseClient(closeSse, ctx))
         }
     }
 


Test patch:
diff --git a/javalin/src/test/java/io/javalin/TestSse.kt b/javalin/src/test/java/io/javalin/TestSse.kt
index 10f45edc..56191fad 100644
--- a/javalin/src/test/java/io/javalin/TestSse.kt
+++ b/javalin/src/test/java/io/javalin/TestSse.kt
@@ -6,10 +6,8 @@ import io.javalin.testing.SerializableObject
 import io.javalin.testing.TestUtil
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.Timeout
 import java.util.concurrent.CompletableFuture
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.locks.LockSupport
-
 
 class TestSse {
 
@@ -108,11 +106,11 @@ class TestSse {
 
     @Test
     fun `send async data is properly processed`() = TestUtil.test { app, http ->
-        app.get("/sse", SseHandler(100) {
+        app.sse("/sse", SseHandler(100) {
             it.sendEvent("Sync event")
 
             CompletableFuture.runAsync {
-                LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10))
+                Thread.sleep(10)
                 it.sendEvent("Async event")
             }
         })
@@ -127,4 +125,12 @@ class TestSse {
         """.trimIndent().trim())
     }
 
+    @Timeout(5)
+    @Test
+    fun `should close sse if requested`() = TestUtil.test { app, http ->
+        app.sse("/sse") { it.close() }
+        val body = http.sse("/sse").get().body
+        assertThat(body).isEqualTo("")
+    }
+
 }