Generating code for an annotation using Java Poet

Question

I want to create method level annotations & generate code for it. I want to avoid using AspectJ and would prefer a compile time code generator so that if I've to ever debug the code I can actually see what is happening which aspectJ won't allow me to.

I came across JavaPoet as an option to do this.

I want to create a method level annotation called Latency that captures the execution time for a given method.

So essentially if I've a method like:

@Latency
void process();

The generated code should be:

try {
   long startTime = System.currentTimeMillis();
   this.process();
} finally {
   System.out.println("Total execution time" + System.currentTimeMillis() - startTime);
}

My annotation is defined:

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Latency {
}

The Javapoet code is:

package org.example;

import com.google.auto.service.AutoService;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.JavaFile;
import com.squareup.javapoet.MethodSpec;
import com.squareup.javapoet.TypeName;
import com.squareup.javapoet.TypeSpec;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Processor;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.TypeElement;
import java.io.IOException;
import java.util.Set;

@AutoService(Processor.class)
public class LatencyProcessor extends AbstractProcessor {

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return Set.of(Latency.class.getName());
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        System.out.println("LatencyProcessor is running...");
        for (Element element : roundEnv.getElementsAnnotatedWith(Latency.class)) {
            if (element.getKind().equals(ElementKind.METHOD)) {
                generateLatencyCode((ExecutableElement) element);
            }
        }
        return true;
    }

    private void generateLatencyCode(ExecutableElement methodElement) {
        String methodName = methodElement.getSimpleName().toString();

        CodeBlock codeBlock = CodeBlock.builder()
            .beginControlFlow("try")
            .addStatement("long startTime = System.currentTimeMillis()")
            .addStatement("$N.$N()", "this", methodName)
            .addStatement("long endTime = System.currentTimeMillis()")
            .addStatement("System.out.println(\"Method $N execution time: \" + (endTime - startTime) + \" milliseconds\")", methodName)
            .nextControlFlow("catch (Exception e)")
            .addStatement("e.printStackTrace()")
            .endControlFlow()
            .build();

        MethodSpec latencyMethod = MethodSpec.methodBuilder(methodName)
            .addModifiers(Modifier.PUBLIC)
            .addAnnotation(Override.class)
            .returns(TypeName.VOID)
            .addCode(codeBlock)
            .build();

        TypeSpec latencyClass = TypeSpec.classBuilder("Latency_" + methodName)
            .addModifiers(Modifier.PUBLIC)
            .addSuperinterface(ClassName.get(methodElement.getEnclosingElement().asType()))
            .addMethod(latencyMethod)
            .build();

        JavaFile javaFile = JavaFile.builder("generated", latencyClass)
            .build();

        try {
            javaFile.writeTo(processingEnv.getFiler());
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}

My build.gradle file is:

dependencies {
    implementation 'com.squareup:javapoet:1.13.0'
    implementation 'com.google.auto.service:auto-service:1.1.1'


    annotationProcessor 'com.google.auto.service:auto-service:1.1.1'

}

configurations {
    annotationProcessor
}

I don't see any generated code & LatencyProcessor is never invoked. Am I misunderstanding the use of Javapoet or is it just that I have not set it up correctly?

Answer 1

The problem seems to be that you are trying to use the custom annotation processor, LatencyProcessor, in the same project that it is a part of. Instead you need to build it in advance, either in a separate project or in a sub project and then use that jar or sub project as an annotation processor in the project that contains the method annotated with @Latency. I tested it out using a subproject. For simplicity, I left out auto-service and added the javax.annotation.processing.Processor file that it creates manually. Also needed to remove the line .addStatement("$N.$N()", "this", methodName) from LatencyProcessor as otherwise invoking the generated process method resulted in a StackOverflowError due to it calling the process method again resulting in an infinite loop.

build.gradle of subproject containing annotation processor:

plugins {
    id 'java'
}

repositories {
    mavenCentral()
}

dependencies {
    implementation 'com.squareup:javapoet:1.13.0'
}

build.gradle of main project that uses the annotation processor:

plugins {
    id 'java'
}

repositories {
    mavenCentral()
}

dependencies {
    // The subproject is named as annotation-processor in this case
    annotationProcessor project(':annotation-processor')
    compileOnly project(':annotation-processor')
}

jar {
    manifest {
        attributes 'Main-Class': 'io.github.devatherock.app.App'
    }
}

Output when executing the generated method was something like the below:

Method process execution time: 0 milliseconds

You can find the complete working code on github