kotlin: Allow Kotlin suspending functions in error handlers

jknack · jknack · commit b015958b92cf · 2024-06-21T13:50:11.000-03:00
- add custom error handler for coroutines - fix #3405
diff --git a/jooby/src/main/java/io/jooby/DefaultErrorHandler.java b/jooby/src/main/java/io/jooby/DefaultErrorHandler.java
@@ -9,11 +9,7 @@
 import static io.jooby.MediaType.json;
 import static io.jooby.MediaType.text;
 
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.Optional;
-import java.util.Set;
-import java.util.stream.Stream;
+import java.util.*;
 
 import org.slf4j.Logger;
 
@@ -28,9 +24,9 @@
  */
 public class DefaultErrorHandler implements ErrorHandler {
 
-  private Set<StatusCode> muteCodes = new HashSet<>();
+  private final Set<StatusCode> muteCodes = new HashSet<>();
 
-  private Set<Class> muteTypes = new HashSet<>();
+  private final Set<Class> muteTypes = new HashSet<>();
 
   /**
    * Generate a log.debug call if any of the status code error occurs as exception.
@@ -39,7 +35,7 @@ public class DefaultErrorHandler implements ErrorHandler {
    * @return This error handler.
    */
   public @NonNull DefaultErrorHandler mute(@NonNull StatusCode... statusCodes) {
-    Stream.of(statusCodes).forEach(muteCodes::add);
+    muteCodes.addAll(List.of(statusCodes));
     return this;
   }
 
@@ -50,7 +46,7 @@ public class DefaultErrorHandler implements ErrorHandler {
    * @return This error handler.
    */
   public @NonNull DefaultErrorHandler mute(@NonNull Class<? extends Exception>... exceptionTypes) {
-    Stream.of(exceptionTypes).forEach(muteTypes::add);
+    muteTypes.addAll(List.of(exceptionTypes));
     return this;
   }
 
@@ -63,7 +59,7 @@ protected void log(Context ctx, Throwable cause, StatusCode code) {
     }
   }
 
-  @NonNull @Override
+  @Override
   public void apply(@NonNull Context ctx, @NonNull Throwable cause, @NonNull StatusCode code) {
     log(ctx, cause, code);
     MediaType type = ctx.accept(Arrays.asList(html, json, text));
diff --git a/jooby/src/main/java/io/jooby/ErrorHandler.java b/jooby/src/main/java/io/jooby/ErrorHandler.java
@@ -22,7 +22,6 @@ public interface ErrorHandler {
    * @param cause Application error.
    * @param code Status code.
    */
-  @NonNull
   void apply(@NonNull Context ctx, @NonNull Throwable cause, @NonNull StatusCode code);
 
   /**
@@ -50,15 +49,13 @@ public interface ErrorHandler {
    * @return Single line message.
    */
   static @NonNull String errorMessage(@NonNull Context ctx, @NonNull StatusCode statusCode) {
-    return new StringBuilder()
-        .append(ctx.getMethod())
-        .append(" ")
-        .append(ctx.getRequestPath())
-        .append(" ")
-        .append(statusCode.value())
-        .append(" ")
-        .append(statusCode.reason())
-        .toString();
+    return ctx.getMethod()
+        + " "
+        + ctx.getRequestPath()
+        + " "
+        + statusCode.value()
+        + " "
+        + statusCode.reason();
   }
 
   /**
diff --git a/modules/jooby-kotlin/src/main/kotlin/io/jooby/kt/CoroutineRouter.kt b/modules/jooby-kotlin/src/main/kotlin/io/jooby/kt/CoroutineRouter.kt
@@ -5,25 +5,13 @@
  */
 package io.jooby.kt
 
-import io.jooby.RequestScope
-import io.jooby.Route
-import io.jooby.Router
-import io.jooby.Router.DELETE
-import io.jooby.Router.GET
-import io.jooby.Router.HEAD
-import io.jooby.Router.OPTIONS
-import io.jooby.Router.PATCH
-import io.jooby.Router.POST
-import io.jooby.Router.PUT
-import io.jooby.Router.TRACE
+import io.jooby.*
+import io.jooby.Router.*
+import java.util.function.Predicate
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
-import kotlinx.coroutines.CoroutineExceptionHandler
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart
-import kotlinx.coroutines.asContextElement
-import kotlinx.coroutines.asCoroutineDispatcher
-import kotlinx.coroutines.launch
+import kotlin.reflect.KClass
+import kotlinx.coroutines.*
 
 internal class RouterCoroutineScope(override val coroutineContext: CoroutineContext) :
   CoroutineScope
@@ -34,6 +22,8 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
     RouterCoroutineScope(router.worker.asCoroutineDispatcher())
   }
 
+  private var errorHandler: suspend ErrorHandlerContext.() -> Unit = FALLBACK_ERROR_HANDLER
+
   private var extraCoroutineContextProvider: HandlerContext.() -> CoroutineContext = {
     EmptyCoroutineContext
   }
@@ -42,6 +32,84 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
     extraCoroutineContextProvider = provider
   }
 
+  /**
+   * Add a custom error handler that matches the given status code.
+   *
+   * @param statusCode Status code.
+   * @param handler Error handler.
+   * @return This router.
+   */
+  @RouterDsl
+  fun error(
+    statusCode: StatusCode,
+    handler: suspend ErrorHandlerContext.() -> Unit
+  ): CoroutineRouter {
+    return error({ it: StatusCode -> statusCode == it }, handler)
+  }
+
+  /**
+   * Add a custom error handler that matches the given exception type.
+   *
+   * @param type Exception type.
+   * @param handler Error handler.
+   * @return This router.
+   */
+  @RouterDsl
+  fun error(
+    type: KClass<Throwable>,
+    handler: suspend ErrorHandlerContext.() -> Unit
+  ): CoroutineRouter {
+    return error {
+      if (type.java.isInstance(cause) || type.java.isInstance(cause.cause)) {
+        handler.invoke(this)
+      }
+    }
+  }
+
+  /**
+   * Add a custom error handler that matches the given predicate.
+   *
+   * @param predicate Status code filter.
+   * @param handler Error handler.
+   * @return This router.
+   */
+  @RouterDsl
+  fun error(
+    predicate: Predicate<StatusCode>,
+    handler: suspend ErrorHandlerContext.() -> Unit
+  ): CoroutineRouter {
+    return error {
+      if (predicate.test(statusCode)) {
+        handler.invoke(this)
+      }
+    }
+  }
+
+  /**
+   * Add a custom error handler.
+   *
+   * @param handler Error handler.
+   * @return This router.
+   */
+  @RouterDsl
+  fun error(handler: suspend ErrorHandlerContext.() -> Unit): CoroutineRouter {
+    val chain =
+      fun(
+        current: suspend ErrorHandlerContext.() -> Unit,
+        next: suspend ErrorHandlerContext.() -> Unit
+      ): suspend ErrorHandlerContext.() -> Unit {
+        return {
+          current(this)
+          if (!ctx.isResponseStarted) {
+            next(this)
+          }
+        }
+      }
+    errorHandler =
+      if (errorHandler == FALLBACK_ERROR_HANDLER) handler else chain(errorHandler, handler)
+    return this
+  }
+
   @RouterDsl
   fun get(pattern: String, handler: suspend HandlerContext.() -> Any) = route(GET, pattern, handler)
 
@@ -77,10 +145,18 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
       .route(method, pattern) { ctx ->
         val handlerContext = HandlerContext(ctx)
         launch(handlerContext) {
-          val result = handler(handlerContext)
-          ctx.route.after?.apply(ctx, result, null)
-          if (result != ctx && !ctx.isResponseStarted) {
-            ctx.render(result)
+          try {
+            val result = handler(handlerContext)
+            ctx.route.after?.apply(ctx, result, null)
+            if (result != ctx && !ctx.isResponseStarted) {
+              ctx.render(result)
+            }
+          } catch (cause: Throwable) {
+            try {
+              ctx.route.after?.apply(ctx, null, cause)
+            } finally {
+              errorHandler.invoke(ErrorHandlerContext(ctx, cause, router.errorCode(cause)))
+            }
           }
         }
         // Return context to mark as handled
@@ -89,6 +165,7 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
       .setHandle(handler)
 
   internal fun launch(handlerContext: HandlerContext, block: suspend CoroutineScope.() -> Unit) {
+    // Global catch-all exception handler
     val exceptionHandler = CoroutineExceptionHandler { _, x ->
       val ctx = handlerContext.ctx
       ctx.route.after?.apply(ctx, null, x)
@@ -99,4 +176,10 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
       exceptionHandler + requestScope + handlerContext.extraCoroutineContextProvider()
     coroutineScope.launch(coroutineContext, coroutineStart, block)
   }
+
+  private companion object {
+    private val FALLBACK_ERROR_HANDLER: suspend ErrorHandlerContext.() -> Unit = {
+      ctx.sendError(cause, statusCode)
+    }
+  }
 }
diff --git a/modules/jooby-kotlin/src/main/kotlin/io/jooby/kt/HandlerContext.kt b/modules/jooby-kotlin/src/main/kotlin/io/jooby/kt/HandlerContext.kt
@@ -5,17 +5,17 @@
  */
 package io.jooby.kt
 
-import io.jooby.Context
-import io.jooby.Route
-import io.jooby.ServerSentEmitter
-import io.jooby.WebSocketConfigurer
+import io.jooby.*
 
 class AfterContext(val ctx: Context, val result: Any?, val failure: Any?)
 
 class FilterContext(val ctx: Context, val next: Route.Handler)
 
 class HandlerContext(val ctx: Context) : java.io.Serializable
 
+class ErrorHandlerContext(val ctx: Context, val cause: Throwable, val statusCode: StatusCode) :
+  java.io.Serializable
+
 class WebSocketInitContext(val ctx: Context, val configurer: WebSocketConfigurer)
 
 class ServerSentHandler(val ctx: Context, val sse: ServerSentEmitter)
diff --git a/tests/src/test/kotlin/i3405/Issue3405.kt b/tests/src/test/kotlin/i3405/Issue3405.kt
@@ -0,0 +1,80 @@
+/*
+ * Jooby https://jooby.io
+ * Apache License Version 2.0 https://jooby.io/LICENSE.txt
+ * Copyright 2014 Edgar Espina
+ */
+package i3405
+
+import io.jooby.StatusCode
+import io.jooby.exception.StatusCodeException
+import io.jooby.junit.ServerTest
+import io.jooby.junit.ServerTestRunner
+import io.jooby.kt.Kooby
+import kotlinx.coroutines.delay
+import org.junit.jupiter.api.Assertions.assertEquals
+
+class Issue3405 {
+  @ServerTest
+  fun coroutineShouldFallbackToNormalErrorHandler(runner: ServerTestRunner) =
+    runner
+      .use {
+        Kooby {
+          error { ctx, cause, code -> ctx.send("normal/global") }
+          coroutine { get("/i3405/normal-error") { ctx.query("q").value() } }
+        }
+      }
+      .ready { client ->
+        client.get("/i3405/normal-error") { rsp ->
+          assertEquals("normal/global", rsp.body!!.string())
+        }
+      }
+
+  @ServerTest
+  fun coroutineSuspendErrorHandler(runner: ServerTestRunner) =
+    runner
+      .use {
+        Kooby {
+          error { ctx, cause, code -> ctx.send("normal/global") }
+          coroutine {
+            get("/i3405/coroutine-error") {
+              delay(10)
+              ctx.query("q").value()
+            }
+
+            error {
+              delay(10)
+              ctx.send("coroutine/suspended")
+            }
+          }
+        }
+      }
+      .ready { client ->
+        client.get("/i3405/coroutine-error") { rsp ->
+          assertEquals("coroutine/suspended", rsp.body!!.string())
+        }
+      }
+
+  @ServerTest
+  fun coroutineChainSuspendErrorHandler(runner: ServerTestRunner) =
+    runner
+      .use {
+        Kooby {
+          error { ctx, cause, code -> ctx.send("normal/global") }
+          coroutine {
+            get("/i3405/chain") {
+              if (ctx.query("q").isMissing) throw StatusCodeException(StatusCode.CONFLICT)
+            }
+            get("/i3405/coroutine-error") { ctx.query("q").value() }
+
+            error(StatusCode.CONFLICT) { ctx.send("conflict") }
+            error { ctx.send("coroutine/suspended") }
+          }
+        }
+      }
+      .ready { client ->
+        client.get("/i3405/coroutine-error") { rsp ->
+          assertEquals("coroutine/suspended", rsp.body!!.string())
+        }
+        client.get("/i3405/chain") { rsp -> assertEquals("conflict", rsp.body!!.string()) }
+      }
+}
