improve signature of launchContext

Anton Keks, Margo Mitt · Anton Keks, Margo Mitt · commit 8f3aa3315f0d · 2021-06-09T09:49:26.000+03:00
Reasons:

1. Extra coroutine context often needs to access Jooby Context
2. Relying on users to extend the default CoroutineContext can be error prone - users can forget to include the ExceptionHandler, resuling in "hanging" asynchronous requests
diff --git a/jooby/src/main/kotlin/io/jooby/CoroutineRouter.kt b/jooby/src/main/kotlin/io/jooby/CoroutineRouter.kt
@@ -19,6 +19,7 @@ import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.launch
 import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 
 internal class RouterCoroutineScope(override val coroutineContext: CoroutineContext) : CoroutineScope
 
@@ -28,9 +29,9 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
     RouterCoroutineScope(router.worker.asCoroutineDispatcher())
   }
 
-  private var extendCoroutineContext: (CoroutineContext) -> CoroutineContext = { it }
-  fun launchContext(block: (CoroutineContext) -> CoroutineContext) {
-    extendCoroutineContext = block
+  private var extraCoroutineContextProvider: HandlerContext.() -> CoroutineContext = { EmptyCoroutineContext }
+  fun launchContext(provider: HandlerContext.() -> CoroutineContext) {
+    extraCoroutineContextProvider = provider
   }
 
   @RouterDsl
@@ -67,16 +68,18 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
 
   fun route(method: String, pattern: String, handler: suspend HandlerContext.() -> Any): Route =
       router.route(method, pattern) { ctx ->
-        launch(ctx) {
-          val result = handler(HandlerContext(ctx))
+        val handlerContext = HandlerContext(ctx)
+        launch(handlerContext) {
+          val result = handler(handlerContext)
           if (result != ctx) {
             ctx.render(result)
           }
         }
       }.setHandle(handler).attribute("coroutine", true)
 
-  internal fun launch(ctx: Context, block: suspend CoroutineScope.() -> Unit) {
-    val exceptionHandler = CoroutineExceptionHandler { _, x -> ctx.sendError(x) }
-    coroutineScope.launch(extendCoroutineContext(exceptionHandler), coroutineStart, block)
+  internal fun launch(handlerContext: HandlerContext, block: suspend CoroutineScope.() -> Unit) {
+    val exceptionHandler = CoroutineExceptionHandler { _, x -> handlerContext.ctx.sendError(x) }
+    val coroutineContext = exceptionHandler + handlerContext.extraCoroutineContextProvider()
+    coroutineScope.launch(coroutineContext, coroutineStart, block)
   }
 }
diff --git a/jooby/src/main/kotlin/io/jooby/HandlerContext.kt b/jooby/src/main/kotlin/io/jooby/HandlerContext.kt
@@ -9,8 +9,8 @@ class AfterContext(val ctx: Context, val result: Any?, val failure: Any?)
 
 class DecoratorContext(val ctx: Context, val next: Route.Handler)
 
-class HandlerContext (val ctx: Context)
+class HandlerContext(val ctx: Context)
 
-class WebSocketInitContext (val ctx: Context, val configurer: WebSocketConfigurer)
+class WebSocketInitContext(val ctx: Context, val configurer: WebSocketConfigurer)
 
-class ServerSentHandler (val ctx: Context, val sse: ServerSentEmitter)
+class ServerSentHandler(val ctx: Context, val sse: ServerSentEmitter)
diff --git a/jooby/src/main/kotlin/io/jooby/internal/mvc/CoroutineLauncher.kt b/jooby/src/main/kotlin/io/jooby/internal/mvc/CoroutineLauncher.kt
@@ -7,6 +7,7 @@ package io.jooby.internal.mvc
 
 import io.jooby.Context
 import io.jooby.CoroutineRouter
+import io.jooby.HandlerContext
 import io.jooby.Route
 import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
 
@@ -16,7 +17,7 @@ import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
 class CoroutineLauncher(val next: Route.Handler) : Route.Handler {
   override fun apply(ctx: Context) = ctx.also {
     val router = ctx.router.attribute<CoroutineRouter>("coroutineRouter")
-    router.launch(ctx) {
+    router.launch(HandlerContext(ctx)) {
       val result = suspendCoroutineUninterceptedOrReturn<Any> {
         ctx.attribute("___continuation", it)
         next.apply(ctx)
diff --git a/jooby/src/test/kotlin/io/jooby/CoroutineRouterTest.kt b/jooby/src/test/kotlin/io/jooby/CoroutineRouterTest.kt
@@ -1,18 +1,16 @@
 package io.jooby
 
 import io.jooby.Router.GET
-import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.coroutineScope
+import org.junit.jupiter.api.Assertions.assertSame
+import org.junit.jupiter.api.Assertions.assertTrue
 import org.junit.jupiter.api.Test
 import org.mockito.ArgumentCaptor
 import org.mockito.Mockito.RETURNS_DEEP_STUBS
-import org.mockito.Mockito.`when`
-import org.mockito.Mockito.any
-import org.mockito.Mockito.argThat
 import org.mockito.Mockito.eq
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.verify
-import org.mockito.Mockito.verifyNoInteractions
 import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 
@@ -35,26 +33,24 @@ class CoroutineRouterTest {
 
   @Test
   fun launchContext_isRunEveryTime() {
-    val mockCoroutineContext = mock(CoroutineContext::class.java)
-    `when`(mockCoroutineContext.plus(any()
-        ?: mockCoroutineContext)).thenReturn(mockCoroutineContext, ExtraContext())
-
-    CoroutineRouter(CoroutineStart.DEFAULT, router).apply {
-      launchContext { mockCoroutineContext + it + ExtraContext() }
-      get("/path") { "Result" }
+    var coroutineRouteCalled = false
+    CoroutineRouter(CoroutineStart.UNDISPATCHED, router).apply {
+      launchContext { SampleCoroutineContext(ctx) }
+      get("/path") {
+        coroutineScope {
+          assertSame(ctx, coroutineContext[SampleCoroutineContext.Key]!!.ctx)
+          coroutineRouteCalled = true
+        }
+      }
     }
 
     val handlerCaptor = ArgumentCaptor.forClass(Route.Handler::class.java)
     verify(router).route(eq(GET), eq("/path"), handlerCaptor.capture())
-    verifyNoInteractions(mockCoroutineContext)
-
     handlerCaptor.value.apply(ctx)
-    verify(mockCoroutineContext).plus(argThat { it is CoroutineExceptionHandler }
-        ?: mockCoroutineContext)
-    verify(mockCoroutineContext).plus(argThat { it is ExtraContext } ?: mockCoroutineContext)
+    assertTrue(coroutineRouteCalled)
   }
 
-  class ExtraContext : AbstractCoroutineContextElement(Key) {
-    companion object Key : CoroutineContext.Key<ExtraContext>
+  class SampleCoroutineContext(val ctx: Context) : AbstractCoroutineContextElement(Key) {
+    companion object Key : CoroutineContext.Key<SampleCoroutineContext>
   }
 }
