apt: integrate with kotlin coroutines

jknack · jknack · commit 9d36d3cce365 · 2019-08-29T08:41:21.000-03:00
diff --git a/jooby/src/main/kotlin/io/jooby/CoroutineRouter.kt b/jooby/src/main/kotlin/io/jooby/CoroutineRouter.kt
@@ -68,16 +68,12 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {
       val xhandler = CoroutineExceptionHandler { _, x ->
         ctx.sendError(x)
       }
-      coroutineScope.launch(CoroutineRouter.ContextCoroutineName + xhandler, coroutineStart) {
+      coroutineScope.launch(xhandler, coroutineStart) {
         val result = handler(HandlerContext(ctx))
         if (result != ctx) {
           ctx.render(result)
         }
       }
     }.setHandle(handler)
   }
-
-  companion object {
-    private val ContextCoroutineName = CoroutineName("ctx")
-  }
 }
diff --git a/jooby/src/main/kotlin/io/jooby/internal/mvc/CoroutineLauncher.kt b/jooby/src/main/kotlin/io/jooby/internal/mvc/CoroutineLauncher.kt
@@ -0,0 +1,32 @@
+/**
+ * Jooby https://jooby.io
+ * Apache License Version 2.0 https://jooby.io/LICENSE.txt
+ * Copyright 2014 Edgar Espina
+ */
+package io.jooby.internal.mvc
+
+import io.jooby.Context
+import io.jooby.CoroutineRouter
+import io.jooby.Route
+import kotlinx.coroutines.CoroutineExceptionHandler
+import kotlinx.coroutines.launch
+import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
+
+class CoroutineLauncher(val next: Route.Handler) : Route.Handler {
+  override fun apply(ctx: Context): Any {
+    val router = ctx.router.attribute<CoroutineRouter>("coroutineRouter")
+    val exceptionHandler = CoroutineExceptionHandler { _, x ->
+      ctx.sendError(x)
+    }
+    router.coroutineScope.launch(exceptionHandler, router.coroutineStart) {
+      val result = suspendCoroutineUninterceptedOrReturn<Any> {
+        ctx.attribute("___continuation", it)
+        next.apply(ctx)
+      }
+      if (!ctx.isResponseStarted) {
+        ctx.render(result!!)
+      }
+    }
+    return ctx
+  }
+}
diff --git a/modules/jooby-apt/src/main/java/io/jooby/internal/apt/HandlerCompiler.java b/modules/jooby-apt/src/main/java/io/jooby/internal/apt/HandlerCompiler.java
@@ -25,6 +25,7 @@
 
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.inject.Provider;
+import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.TypeElement;
@@ -177,16 +178,36 @@ private void apply(ClassWriter writer) throws Exception {
     apply.visitEnd();
   }
 
-  private void processArguments(ClassWriter classWriter, MethodVisitor visitor) throws Exception {
-    for (VariableElement var : executable.getParameters()) {
-      visitor.visitVarInsn(ALOAD, 1);
-      ParamDefinition param = ParamDefinition.create(environment, var);
-      ParamWriter writer = param.newWriter();
-      writer.accept(classWriter, getGeneratedInternalClass(), visitor, param);
+  public boolean isSuspendFunction() {
+    List<? extends VariableElement> parameters = executable.getParameters();
+    if (parameters.isEmpty()) {
+      return false;
     }
+    VariableElement last = parameters.get(parameters.size() - 1);
+    return isSuspendFunction(last);
   }
 
+  private boolean isSuspendFunction(VariableElement parameter) {
+    String type = ParamDefinition.create(environment, parameter).getType().getRawType().toString();
+    return type.equals("kotlin.coroutines.Continuation");
+  }
 
+  private void processArguments(ClassWriter classWriter, MethodVisitor visitor) throws Exception {
+    for (VariableElement var : executable.getParameters()) {
+      if (isSuspendFunction(var)) {
+        visitor.visitVarInsn(ALOAD, 1);
+        visitor.visitMethodInsn(INVOKEINTERFACE, "io/jooby/Context", "getAttributes", "()Ljava/util/Map;", true);
+        visitor.visitLdcInsn("___continuation");
+        visitor.visitMethodInsn(INVOKEINTERFACE, "java/util/Map", "remove", "(Ljava/lang/Object;)Ljava/lang/Object;", true);
+        visitor.visitTypeInsn(CHECKCAST, "kotlin/coroutines/Continuation");
+      } else {
+        visitor.visitVarInsn(ALOAD, 1);
+        ParamDefinition param = ParamDefinition.create(environment, var);
+        ParamWriter writer = param.newWriter();
+        writer.accept(classWriter, getGeneratedInternalClass(), visitor, param);
+      }
+    }
+  }
 
   private void processReturnType(MethodVisitor visitor) throws Exception {
     TypeKind kind = executable.getReturnType().getKind();
diff --git a/modules/jooby-apt/src/main/java/io/jooby/internal/apt/ModuleCompiler.java b/modules/jooby-apt/src/main/java/io/jooby/internal/apt/ModuleCompiler.java
@@ -101,13 +101,21 @@ private void install(ClassWriter writer, List<HandlerCompiler> handlers) throws
     for (HandlerCompiler handler : handlers) {
       visitor.visitVarInsn(ALOAD, 1);
       visitor.visitLdcInsn(handler.getPattern());
+      if (handler.isSuspendFunction()) {
+        visitor.visitTypeInsn(NEW, "io/jooby/internal/mvc/CoroutineLauncher");
+        visitor.visitInsn(DUP);
+      }
       visitor.visitTypeInsn(NEW, handler.getGeneratedInternalClass());
       visitor.visitInsn(DUP);
       visitor.visitVarInsn(ALOAD, 0);
       visitor.visitFieldInsn(GETFIELD, moduleInternalName, "provider",
           "Ljavax/inject/Provider;");
       visitor.visitMethodInsn(INVOKESPECIAL, handler.getGeneratedInternalClass(), "<init>",
           "(Ljavax/inject/Provider;)V", false);
+      if (handler.isSuspendFunction()) {
+        visitor.visitMethodInsn(INVOKESPECIAL, "io/jooby/internal/mvc/CoroutineLauncher", "<init>",
+            "(Lio/jooby/Route$Handler;)V", false);
+      }
       visitor.visitMethodInsn(INVOKEVIRTUAL, "io/jooby/Jooby", handler.getHttpMethod(),
           "(Ljava/lang/String;Lio/jooby/Route$Handler;)Lio/jooby/Route;", false);
       visitor.visitVarInsn(ASTORE, 2);
@@ -190,7 +198,9 @@ private Optional<? extends AnnotationMirror> findAnnotation(
   private void setReturnType(MethodVisitor visitor, HandlerCompiler handler)
       throws NoSuchMethodException {
     TypeDefinition returnType = handler.getReturnType();
-    if (returnType.isVoid()) {
+    if (handler.isSuspendFunction()) {
+      visitor.visitLdcInsn(Type.getType("Lkotlin/coroutines/Continuation;"));
+    } else if (returnType.isVoid()) {
       visitor.visitLdcInsn(Type.getType(Context.class));
     } else if (returnType.isPrimitive()) {
       Method wrapper = Primitives.wrapper(returnType);
diff --git a/modules/jooby-apt/src/test/java/output/ASMPrinter.java b/modules/jooby-apt/src/test/java/output/ASMPrinter.java
@@ -1,40 +1,33 @@
 package output;
 
-import examples.CoroutineApp;
 import org.junit.jupiter.api.Test;
 import org.objectweb.asm.util.ASMifier;
-import source.NullRoutes;
 
 public class ASMPrinter {
 
-  @Test
-  public void coroutines() throws Exception {
-    ASMifier.main(new String[] {CoroutineApp.class.getName() + "$1"});
-  }
-
   @Test
   public void mvcModule() throws Exception {
-    ASMifier.main(new String[] {MyMvcModule.class.getName() });
+    ASMifier.main(new String[]{MyMvcModule.class.getName()});
   }
 
   @Test
   public void mvcExtension() throws Exception {
-    ASMifier.main(new String[] {MvcExtension.class.getName() });
+    ASMifier.main(new String[]{MvcExtension.class.getName()});
   }
 
   @Test
   public void mvcDispatch() throws Exception {
-    ASMifier.main(new String[] {MvcDispatch.class.getName() });
+    ASMifier.main(new String[]{MvcDispatch.class.getName()});
   }
 
   @Test
   public void myController() throws Exception {
-    ASMifier.main(new String[] {MyControllerHandler.class.getName() });
+    ASMifier.main(new String[]{MyControllerHandler.class.getName()});
   }
 
   @Test
   public void nullRoutes() throws Exception {
-    ASMifier.main(new String[] {NullRoutes.class.getName() });
+    ASMifier.main(new String[]{"source.SuspendRoute"});
   }
 }
 
diff --git a/modules/jooby-apt/src/test/java/output/MvcExtension.java b/modules/jooby-apt/src/test/java/output/MvcExtension.java
@@ -1,19 +1,23 @@
 package output;
 
+import io.jooby.Context;
 import io.jooby.Extension;
 import io.jooby.Jooby;
 import io.jooby.Route;
+import io.jooby.internal.mvc.CoroutineLauncher;
 
 import javax.annotation.Nonnull;
 import javax.inject.Provider;
 
 public class MvcExtension implements Extension {
-  public MvcExtension(Provider c) {
+  private Provider provider;
 
+  public MvcExtension(Provider c) {
+    this.provider = c;
   }
 
   @Override public void install(@Nonnull Jooby application) throws Exception {
-    Route route = application.get("/", ctx -> "..");
-    route.setExecutorKey("myexecutor");
+    Route route = application.get("/", new CoroutineLauncher(new MyControllerHandler(provider)));
+    route.setReturnType(Context.class);
   }
 }
diff --git a/modules/jooby-apt/src/test/java/output/MyControllerHandler.java b/modules/jooby-apt/src/test/java/output/MyControllerHandler.java
@@ -3,20 +3,22 @@
 import io.jooby.Context;
 import io.jooby.Route;
 import io.jooby.Value;
+import kotlin.coroutines.Continuation;
+import source.SuspendRoute;
 
 import javax.annotation.Nonnull;
 import javax.inject.Provider;
 
 public class MyControllerHandler implements Route.Handler {
 
-  private Provider<MyController> provider;
+  private Provider<SuspendRoute> provider;
 
-  public MyControllerHandler(Provider<MyController> provider) {
+  public MyControllerHandler(Provider<SuspendRoute> provider) {
     this.provider = provider;
   }
 
   @Nonnull @Override public Object apply(@Nonnull Context ctx) throws Exception {
-    return provider.get().doIt(x(ctx.query("x")));
+    return provider.get().suspendFun((Continuation) ctx.getAttributes().remove("__continuation"));
   }
 
   private static Integer x(Value x) {
diff --git a/modules/jooby-apt/src/test/java/source/SuspendRoute.java b/modules/jooby-apt/src/test/java/source/SuspendRoute.java
@@ -0,0 +1,14 @@
+package source;
+
+import io.jooby.annotations.GET;
+import io.jooby.annotations.Path;
+import kotlin.coroutines.Continuation;
+
+@Path("/suspend")
+public class SuspendRoute {
+
+  @GET
+  public Continuation suspendFun(Continuation continuation) {
+    return continuation;
+  }
+}
diff --git a/modules/jooby-apt/src/test/java/tests/ModuleCompilerTest.java b/modules/jooby-apt/src/test/java/tests/ModuleCompilerTest.java
@@ -1,10 +1,14 @@
 package tests;
 
 import io.jooby.Context;
+import io.jooby.CoroutineRouter;
+import io.jooby.MockContext;
 import io.jooby.MockRouter;
 import io.jooby.Route;
 import io.jooby.StatusCode;
 import io.jooby.apt.MvcModuleCompilerRunner;
+import kotlin.coroutines.Continuation;
+import kotlinx.coroutines.GlobalScope;
 import org.junit.jupiter.api.Test;
 import source.GetPostRoute;
 import source.JavaBeanParam;
@@ -14,6 +18,7 @@
 import source.RouteDispatch;
 import source.RouteWithMimeTypes;
 import source.Routes;
+import source.SuspendRoute;
 import source.VoidRoute;
 
 import java.util.Arrays;
@@ -23,6 +28,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
 
 public class ModuleCompilerTest {
   @Test
@@ -147,7 +153,7 @@ public void routeAttributes() throws Exception {
   @Test
   public void routeDispatch() throws Exception {
     new MvcModuleCompilerRunner(new RouteDispatch())
-        .module(true, app -> {
+        .module(app -> {
           assertEquals("worker", app.getRoutes().get(0).getExecutorKey());
           assertEquals("single", app.getRoutes().get(1).getExecutorKey());
         })
diff --git a/modules/jooby-apt/src/test/kotlin/examples/CoroutineApp.kt b/modules/jooby-apt/src/test/kotlin/examples/CoroutineApp.kt
diff --git a/modules/jooby-apt/src/test/kotlin/examples/CoroutineController.kt b/modules/jooby-apt/src/test/kotlin/examples/CoroutineController.kt
diff --git a/modules/jooby-jetty/src/main/java/io/jooby/internal/jetty/JettyContext.java b/modules/jooby-jetty/src/main/java/io/jooby/internal/jetty/JettyContext.java
@@ -82,6 +82,7 @@ public class JettyContext implements Callback, DefaultContext {
   private MediaType responseType;
   private Map<String, String> cookies;
   private HashMap<String, String> responseCookies;
+  private boolean responseStarted;
 
   public JettyContext(Request request, Router router, int bufferSize, long maxRequestSize) {
     this.request = request;
@@ -316,12 +317,14 @@ public Context setResponseType(@Nonnull MediaType contentType, @Nullable Charset
   }
 
   @Nonnull @Override public Sender responseSender() {
+    responseStarted = true;
     ifSetChunked();
     ifStartAsync();
     return new JettySender(this, response.getHttpOutput());
   }
 
   @Nonnull @Override public OutputStream responseStream() {
+    responseStarted = true;
     try {
       ifSetChunked();
       // TODO: session should be safe after response, find a better way
@@ -333,6 +336,7 @@ public Context setResponseType(@Nonnull MediaType contentType, @Nullable Charset
   }
 
   @Nonnull @Override public PrintWriter responseWriter(MediaType type, Charset charset) {
+    responseStarted = true;
     try {
       ifSetChunked();
       // TODO: session should be safe after response, find a better way
@@ -359,12 +363,14 @@ public Context setResponseType(@Nonnull MediaType contentType, @Nullable Charset
   }
 
   @Nonnull @Override public Context send(@Nonnull ByteBuffer data) {
+    responseStarted = true;
     HttpOutput sender = response.getHttpOutput();
     sender.sendContent(data, this);
     return this;
   }
 
   @Nonnull @Override public Context send(@Nonnull ReadableByteChannel channel) {
+    responseStarted = true;
     ifSetChunked();
     ifStartAsync();
     HttpOutput sender = response.getHttpOutput();
@@ -397,6 +403,7 @@ private Context sendStreamInternal(@Nonnull InputStream in) {
         response.setHeader(HttpHeader.TRANSFER_ENCODING, HttpHeaderValue.CHUNKED.asString());
         stream = in;
       }
+      responseStarted = true;
       response.getHttpOutput().sendContent(stream, this);
       return this;
     } catch (IOException x) {
@@ -414,7 +421,7 @@ private Context sendStreamInternal(@Nonnull InputStream in) {
   }
 
   @Override public boolean isResponseStarted() {
-    return response.isCommitted();
+    return responseStarted;
   }
 
   @Override public void succeeded() {
diff --git a/tests/src/test/kotlin/io/jooby/FeaturedKotlinTest.kt b/tests/src/test/kotlin/io/jooby/FeaturedKotlinTest.kt
@@ -122,6 +122,7 @@ class FeaturedKotlinTest {
     }
   }
 
+  @Test
   fun suspendMvc() {
     JoobyRunner { ->
       Kooby {
@@ -149,7 +150,7 @@ class FeaturedKotlinTest {
       }
 
       client.get("/456x") { rsp ->
-        assertEquals("Unable to provision parameter: 'id: int'", rsp.body!!.string())
+        assertEquals("Cannot convert value: 'id', to: 'int'", rsp.body!!.string())
       }
     }
   }
diff --git a/tests/src/test/kotlin/io/jooby/SuspendMvc.kt b/tests/src/test/kotlin/io/jooby/SuspendMvc.kt
@@ -21,6 +21,6 @@ class SuspendMvc {
   @GET
   @Path("/{id}")
   suspend fun pathParam(@PathParam id: Int): Any {
-    return id;
+    return id
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -68,16 +68,12 @@ class CoroutineRouter(val coroutineStart: CoroutineStart, val router: Router) {`
`68`	`68`	`val xhandler = CoroutineExceptionHandler { _, x ->`
`69`	`69`	`ctx.sendError(x)`
`70`	`70`	`}`
`71`		`- coroutineScope.launch(CoroutineRouter.ContextCoroutineName + xhandler, coroutineStart) {`
	`71`	`+ coroutineScope.launch(xhandler, coroutineStart) {`
`72`	`72`	`val result = handler(HandlerContext(ctx))`
`73`	`73`	`if (result != ctx) {`
`74`	`74`	`ctx.render(result)`
`75`	`75`	`}`
`76`	`76`	`}`
`77`	`77`	`}.setHandle(handler)`
`78`	`78`	`}`
`79`		`-`
`80`		`- companion object {`
`81`		`- private val ContextCoroutineName = CoroutineName("ctx")`
`82`		`- }`
`83`	`79`	`}`