apt: implement dispatch annotation

Author: jknack

jooby/src/main/java/io/jooby/Route.java

@@ -21,6 +21,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.TreeMap;
+import java.util.concurrent.Executor;
 
 /**
  * Route contains information about the HTTP method, path pattern, which content types consumes and
@@ -340,6 +341,8 @@ public interface Handler extends Serializable {
 
   private Set<String> supportedMethod;
 
+  private String executorKey;
+
   /**
    * Creates a new route.
    *
@@ -764,6 +767,34 @@ public boolean isHttpHead() {
     return this;
   }
 
+  /**
+   * Specify the name of the executor where the route is going to run.
+   * Default is <code>null</code>.
+   *
+   * @return Executor key.
+   */
+  public @Nullable String getExecutorKey() {
+    return executorKey;
+  }
+
+  /**
+   * Set executor key. The route is going to use the given key to fetch an executor. Possible values
+   * are:
+   *
+   * - <code>null</code>: no specific executor, uses the default Jooby logic to choose one, based
+   *    on the value of {@link ExecutionMode};
+   * - <code>worker</code>: use the executor provided by the server.
+   * - <code>arbitrary name</code>: use an named executor which as registered using
+   *    {@link Router#executor(String, Executor)}.
+   *
+   * @param executorKey Executor key.
+   * @return This route.
+   */
+  public @Nonnull Route setExecutorKey(@Nullable String executorKey) {
+    this.executorKey = executorKey;
+    return this;
+  }
+
   @Override public String toString() {
     return method + " " + pattern;
   }

jooby/src/main/java/io/jooby/annotations/Dispatch.java

@@ -36,5 +36,5 @@
    *
    * @return Name of the executor to use or blank to use the server worker executor.
    */
-  String value() default "";
+  String value() default "worker";
 }

jooby/src/main/java/io/jooby/internal/RouterImpl.java

@@ -6,7 +6,6 @@
 package io.jooby.internal;
 
 import io.jooby.Context;
-import io.jooby.HeadHandler;
 import io.jooby.RegistryException;
 import io.jooby.ServiceKey;
 import io.jooby.StatusCodeException;
@@ -38,7 +37,6 @@
 import javax.annotation.Nonnull;
 import javax.inject.Provider;
 import java.io.FileNotFoundException;
-import java.lang.annotation.Annotation;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.nio.file.Path;
@@ -373,9 +371,11 @@ private Route defineRoute(@Nonnull String method, @Nonnull String pattern,
     handler.setRoute(route);
 
     Stack stack = this.stack.peekLast();
+    String executorKey = route.getExecutorKey();
     if (stack.executor != null) {
       routeExecutor.put(route, stack.executor);
     }
+
     String routePattern = normalizePath(basePath == null
         ? safePattern
         : basePath + safePattern, false, true);
@@ -384,7 +384,9 @@ private Route defineRoute(@Nonnull String method, @Nonnull String pattern,
       tree.insert(Router.OPTIONS, routePattern, route);
     } else if (route.isHttpTrace()) {
       tree.insert(Router.TRACE, routePattern, route);
-    } else if (route.isHttpHead() && route.getMethod().equals(GET)) {
+    } else if (route.isHttpHead() && route.getMethod().
+
+        equals(GET)) {
       tree.insert(Router.HEAD, routePattern, route);
     }
     routes.add(route);
@@ -401,9 +403,21 @@ private Route defineRoute(@Nonnull String method, @Nonnull String pattern,
     renderer.add(MessageEncoder.TO_STRING);
     ExecutionMode mode = owner.getExecutionMode();
     for (Route route : routes) {
-      Executor executor = routeExecutor.get(route);
-      if (executor instanceof ForwardingExecutor) {
-        executor = ((ForwardingExecutor) executor).executor;
+      String executorKey = route.getExecutorKey();
+      Executor executor;
+      if (executorKey == null) {
+        executor = routeExecutor.get(route);
+        if (executor instanceof ForwardingExecutor) {
+          executor = ((ForwardingExecutor) executor).executor;
+        }
+      } else {
+        if (executorKey.equals("worker")) {
+          executor = (worker instanceof ForwardingExecutor) ?
+              ((ForwardingExecutor) worker).executor :
+              worker;
+        } else {
+          executor = executor(executorKey);
+        }
       }
       /** Return type: */
       if (route.getReturnType() == null) {

modules/jooby-apt/src/main/java/io/jooby/internal/apt/ModuleCompiler.java

@@ -11,6 +11,7 @@
 import io.jooby.MediaType;
 import io.jooby.Reified;
 import io.jooby.Route;
+import io.jooby.annotations.Dispatch;
 import io.jooby.internal.apt.asm.ArrayWriter;
 import io.jooby.internal.apt.asm.ConstructorWriter;
 import io.jooby.internal.apt.asm.RouteAttributesWriter;
@@ -20,8 +21,13 @@
 import org.objectweb.asm.Type;
 
 import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.AnnotationValue;
+import javax.lang.model.element.ExecutableElement;
 import java.lang.reflect.Method;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 
 import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
 import static org.objectweb.asm.Opcodes.ACC_SUPER;
@@ -122,15 +128,65 @@ private void install(ClassWriter writer, List<HandlerCompiler> handlers) throws
       setContentType(visitor, "setProduces", handler.getProduces());
 
       /**
+       * ******************************************************************************************
        * Annotations as route attributes
+       * ******************************************************************************************
        */
       routeAttributes.process(handler.getExecutable());
+
+      /**
+       * ******************************************************************************************
+       * Dispatch
+       * ******************************************************************************************
+       */
+      setDispatch(visitor, handler.getExecutable());
     }
     visitor.visitInsn(RETURN);
     visitor.visitMaxs(0, 0);
     visitor.visitEnd();
   }
 
+  private void setDispatch(MethodVisitor visitor, ExecutableElement executable)
+      throws NoSuchMethodException {
+    String executorKey = findAnnotation(executable.getAnnotationMirrors(), Dispatch.class.getName())
+        .map(it -> annotationAttribute(it, "value").toString())
+        .orElseGet(() ->
+            findAnnotation(executable.getEnclosingElement().getAnnotationMirrors(),
+                Dispatch.class.getName())
+                .map(it -> annotationAttribute(it, "value").toString())
+                .orElse(null)
+        );
+
+    if (executorKey != null) {
+      Method setExecutor = Route.class.getDeclaredMethod("setExecutorKey", String.class);
+      visitor.visitVarInsn(ALOAD, 2);
+      visitor.visitLdcInsn(executorKey);
+      visitor.visitMethodInsn(INVOKEVIRTUAL, Type.getInternalName(setExecutor.getDeclaringClass()),
+          setExecutor.getName(),
+          Type.getMethodDescriptor(setExecutor), false);
+      visitor.visitInsn(POP);
+    }
+  }
+
+  private Object annotationAttribute(AnnotationMirror annotationMirror, String method) {
+    Map<? extends ExecutableElement, ? extends AnnotationValue> map = env
+        .getElementUtils().getElementValuesWithDefaults(annotationMirror);
+    for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> entry : map.entrySet()) {
+      if (entry.getKey().getSimpleName().toString().equals(method)) {
+        return entry.getValue().getValue();
+      }
+    }
+    throw new IllegalArgumentException(
+        "Missing: " + annotationMirror.getAnnotationType().toString() + "." + method);
+  }
+
+  private Optional<? extends AnnotationMirror> findAnnotation(
+      List<? extends AnnotationMirror> annotationMirrors, String name) {
+    return annotationMirrors.stream()
+        .filter(it -> it.getAnnotationType().toString().equals(name))
+        .findFirst();
+  }
+
   private void setReturnType(MethodVisitor visitor, HandlerCompiler handler)
       throws NoSuchMethodException {
     TypeDefinition returnType = handler.getReturnType();

modules/jooby-apt/src/test/java/output/MvcExtension.java

@@ -3,20 +3,17 @@
 import io.jooby.Extension;
 import io.jooby.Jooby;
 import io.jooby.Route;
-import org.jetbrains.annotations.NotNull;
 
 import javax.annotation.Nonnull;
 import javax.inject.Provider;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.Map;
 
 public class MvcExtension implements Extension {
   public MvcExtension(Provider c) {
 
   }
 
   @Override public void install(@Nonnull Jooby application) throws Exception {
-    application.dispatch(new MvcDispatch(application));
+    Route route = application.get("/", ctx -> "..");
+    route.setExecutorKey("myexecutor");
   }
 }

modules/jooby-apt/src/test/java/source/RouteDispatch.java

@@ -0,0 +1,22 @@
+package source;
+
+import io.jooby.annotations.Dispatch;
+import io.jooby.annotations.GET;
+import io.jooby.annotations.Path;
+
+@Dispatch
+public class RouteDispatch {
+
+  @Path("/toplevel")
+  @GET
+  public void toplevel() {
+
+  }
+
+  @Path("/methodlevel")
+  @GET
+  @Dispatch("single")
+  public void methodlevel() {
+
+  }
+}

modules/jooby-apt/src/test/java/tests/ModuleCompilerTest.java

@@ -1,19 +1,17 @@
 package tests;
 
 import io.jooby.Context;
-import io.jooby.MockContext;
 import io.jooby.MockRouter;
 import io.jooby.Route;
 import io.jooby.StatusCode;
-import io.jooby.apt.MvcHandlerCompilerRunner;
 import io.jooby.apt.MvcModuleCompilerRunner;
-import org.junit.Assert;
 import org.junit.jupiter.api.Test;
 import source.GetPostRoute;
 import source.JavaBeanParam;
 import source.NoPathRoute;
 import source.PrimitiveReturnType;
 import source.RouteAttributes;
+import source.RouteDispatch;
 import source.RouteWithMimeTypes;
 import source.Routes;
 import source.VoidRoute;
@@ -145,4 +143,14 @@ public void routeAttributes() throws Exception {
         })
     ;
   }
+
+  @Test
+  public void routeDispatch() throws Exception {
+    new MvcModuleCompilerRunner(new RouteDispatch())
+        .module(true, app -> {
+          assertEquals("worker", app.getRoutes().get(0).getExecutorKey());
+          assertEquals("single", app.getRoutes().get(1).getExecutorKey());
+        })
+    ;
+  }
 }

tests/src/test/java/io/jooby/MvcTest.java

@@ -378,7 +378,7 @@ public void mvcBody() {
     });
   }
 
-  //  @Test
+  @Test
   public void mvcDispatch() {
     new JoobyRunner(app -> {
       app.executor("single", Executors.newSingleThreadExecutor(r ->