Repository: tipsy/javalin
Pull request #1533 (2022-05-04 12:15:31)

GH-1532 Adapt SseHandler to reimplemented servlet
Changes:
1. I've added `Javalin.sse` with `SseHandler` (also lambda) to Javalin class as it feels strange that I'd need to use `.get()` to do that. It's also a part of change associated to the optional `async timeout` set per each handler. 
  I can revert this, but we'd need to use `.get` instead of `.sse` in tests to explicitly provide custom timeout.
3. SseHandler has now dedicated (optional) timeout value, so people can provide custom timeouts explicitly for handler rather than some magic tricks like it was before with e.g. after handler like it was in the test scenarios.
4. New SseHandler constructor is backwards compatibile due to @JvmOverloads flag.
5. SseHandler uses stub future instance to properly notify servlet about request state
6. Replaced global `shortTimeoutServer()` with dedicated handler timeout in tests
7. Added a new test case that verifies behavior of non-blocking SseHandler implementation provided by user

Overall notes:
* Because SSE is based on top of the `get` request, it blocks request lifecycle until the SSE is completed. It means that e.g. after handlers won't be executed immediately, so it may affect users that had some kind of a logic similar to the one used in `shortTimeoutServer()`. Imo the new behavior is more consistent, the second one was more likely like a workaround.

Further proposals:
* We could add a new method, sth like `SseClient.close()`, that completes the internal `await` future used by `SseHandler`. Thanks to that we can finally provide an official method to close sse client. Currently it's not possible, because only timeout and client can close the context.

tipsy: Thank you very much @dzikoysk !

Patch:
diff --git a/javalin/src/main/java/io/javalin/Javalin.java b/javalin/src/main/java/io/javalin/Javalin.java
index b1c96dbe..9117f2e5 100644
--- a/javalin/src/main/java/io/javalin/Javalin.java
+++ b/javalin/src/main/java/io/javalin/Javalin.java
@@ -498,6 +498,13 @@ public class Javalin implements AutoCloseable {
         return sse(path, client, new RouteRole[0]);
     }
 
+    /**
+     * Adds a lambda handler for a Server Sent Event connection on the specified path.
+     */
+    public Javalin sse(@NotNull String path, @NotNull SseHandler handler) {
+        return get(path, handler);
+    }
+
     /**
      * Adds a lambda handler for a Server Sent Event connection on the specified path.
      * Requires an access manager to be set on the instance.
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 06c79918..13f3d092 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,16 @@ package io.javalin.http.sse
 import io.javalin.core.util.Header
 import io.javalin.http.Context
 import io.javalin.http.Handler
+import java.util.concurrent.CompletableFuture
 import java.util.function.Consumer
 import javax.servlet.AsyncEvent
 import javax.servlet.AsyncListener
 
-class SseHandler(private val clientConsumer: Consumer<SseClient>) : Handler {
+class SseHandler @JvmOverloads constructor(
+    private val timeout: Long = 0,
+    private val clientConsumer: Consumer<SseClient>
+) : Handler {
+
     override fun handle(ctx: Context) {
         if (ctx.header(Header.ACCEPT) == "text/event-stream") {
             ctx.res.apply {
@@ -20,14 +25,19 @@ class SseHandler(private val clientConsumer: Consumer<SseClient>) : Handler {
                 flushBuffer()
             }
             ctx.req.startAsync(ctx.req, ctx.res)
-            ctx.req.asyncContext.timeout = 0
-            clientConsumer.accept(SseClient(ctx))
+            ctx.req.asyncContext.timeout = timeout
+
+            val awaitFuture = CompletableFuture<Void>()
             ctx.req.asyncContext.addListener(object : AsyncListener {
                 override fun onComplete(event: AsyncEvent) {}
                 override fun onStartAsync(event: AsyncEvent) {}
-                override fun onTimeout(event: AsyncEvent) = event.asyncContext.complete()
-                override fun onError(event: AsyncEvent) = event.asyncContext.complete()
+                override fun onTimeout(event: AsyncEvent) { awaitFuture.complete(null) }
+                override fun onError(event: AsyncEvent) { awaitFuture.complete(null) }
             })
+            ctx.future(awaitFuture) { /* do nothing with the future result in callback */ }
+
+            clientConsumer.accept(SseClient(ctx))
         }
     }
+
 }


Test patch:
diff --git a/javalin/src/test/java/io/javalin/TestSse.kt b/javalin/src/test/java/io/javalin/TestSse.kt
index 7cf1d92f..10f45edc 100644
--- a/javalin/src/test/java/io/javalin/TestSse.kt
+++ b/javalin/src/test/java/io/javalin/TestSse.kt
@@ -1,40 +1,44 @@
 package io.javalin
 
 import io.javalin.http.sse.SseClient
+import io.javalin.http.sse.SseHandler
 import io.javalin.testing.SerializableObject
 import io.javalin.testing.TestUtil
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.locks.LockSupport
+
 
 class TestSse {
 
     private val event = "HI"
     private val data = "Hello, world!"
-    private fun shortTimeoutServer() = Javalin.create().after { it.req.asyncContext.timeout = 10 }
 
     @Test
-    fun `sending events works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendEvent(event, data) }
+    fun `sending events works`() = TestUtil.test() { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendEvent(event, data) })
         val body = http.sse("/sse").get().body
         assertThat(body).contains("event: $event")
         assertThat(body).contains("data: $data")
     }
 
     @Test
-    fun `sending input stream works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendEvent(event, "MY DATA".byteInputStream()) }
+    fun `sending input stream works`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendEvent(event, "MY DATA".byteInputStream()) })
         assertThat(http.sse("/sse").get().body).contains("data: MY DATA")
     }
 
     @Test
-    fun `sending json works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendEvent(event, SerializableObject()) }
+    fun `sending json works`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendEvent(event, SerializableObject()) })
         assertThat(http.sse("/sse").get().body).contains("""data: {"value1":"FirstValue","value2":"SecondValue"}""")
     }
 
     @Test
-    fun `sending events with ID works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendEvent(event, data, id = "SOME_ID") }
+    fun `sending events with ID works`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendEvent(event, data, id = "SOME_ID") })
         val body = http.sse("/sse").get().body
         assertThat(body).contains("id: SOME_ID")
         assertThat(body).contains("event: $event")
@@ -43,12 +47,12 @@ class TestSse {
 
     @Test
     fun `sending events to multiple clients works`() {
-        TestUtil.test(shortTimeoutServer()) { app, http ->
+        TestUtil.test { app, http ->
             val eventSources: MutableList<SseClient> = ArrayList()
-            app.sse("/sse") { sse ->
-                eventSources.add(sse)
-                sse.sendEvent(event, data + eventSources.size)
-            }
+            app.sse("/sse", SseHandler(10) {
+                eventSources.add(it)
+                it.sendEvent(event, data + eventSources.size)
+            })
             val bodyClient1 = http.sse("/sse").get().body
             val bodyClient2 = http.sse("/sse").get().body
             assertThat(bodyClient1).isNotEqualTo(bodyClient2)
@@ -57,8 +61,8 @@ class TestSse {
     }
 
     @Test
-    fun `all headers are correctly configured`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendEvent(event, data) }
+    fun `all headers are correctly configured`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendEvent(event, data) })
         val headers = http.sse("/sse").get().headers // Headers extends HashMap<String, List<String>>
         assertThat(headers.getFirst("Connection")).containsIgnoringCase("close")
         assertThat(headers.getFirst("Content-Type")).containsIgnoringCase("text/event-stream")
@@ -67,39 +71,60 @@ class TestSse {
     }
 
     @Test
-    fun `default http status is 200`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendEvent(event, data) }
+    fun `default http status is 200`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendEvent(event, data) })
         val status = http.sse("/sse").get().status
         assertThat(status).isEqualTo(200)
     }
 
     @Test
-    fun `getting queryParam in sse handler works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendEvent(event, it.ctx.queryParam("qp")!!) }
+    fun `getting queryParam in sse handler works`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendEvent(event, it.ctx.queryParam("qp")!!) })
         val body = http.sse("/sse?qp=my-qp").get().body
         assertThat(body).contains("event: $event")
         assertThat(body).contains("data: " + "my-qp")
     }
 
     @Test
-    fun `sending Comment works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendComment("test comment works") }
+    fun `sending Comment works`() = TestUtil.test() { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendComment("test comment works") })
         val body = http.sse("/sse?qp=my-qp").get().body
         assertThat(body).isEqualTo(": test comment works\n")
     }
 
     @Test
-    fun `sending empty Comment works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendComment("") }
+    fun `sending empty Comment works`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendComment("") })
         val body = http.sse("/sse?qp=my-qp").get().body
         assertThat(body).isEqualTo(": \n")
     }
 
     @Test
-    fun `sending multi line Comment works`() = TestUtil.test(shortTimeoutServer()) { app, http ->
-        app.sse("/sse") { it.sendComment("a\nb") }
+    fun `sending multi line Comment works`() = TestUtil.test { app, http ->
+        app.sse("/sse", SseHandler(10) { it.sendComment("a\nb") })
         val body = http.sse("/sse?qp=my-qp").get().body
         assertThat(body).isEqualTo(": a\n: b\n")
     }
 
+    @Test
+    fun `send async data is properly processed`() = TestUtil.test { app, http ->
+        app.get("/sse", SseHandler(100) {
+            it.sendEvent("Sync event")
+
+            CompletableFuture.runAsync {
+                LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10))
+                it.sendEvent("Async event")
+            }
+        })
+
+        val body = http.sse("/sse").get().body
+        assertThat(body.trim()).isEqualTo("""
+        event: message
+        data: Sync event
+
+        event: message
+        data: Async event
+        """.trimIndent().trim())
+    }
+
 }