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?

Chamfer

Just want to know if you thought of using CGLIB? I consider you want something like @Transactional proxy layer generated in Spring Boot – Selffulfillment

From my understanding that modifies the byte code. I was trying to see a more compile time generated language, so that when I debug its clear what is happening, without doing any additional setup – Chamfer

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

Rowden

Thanks for taking the time to answer. I'll test it out in a while & get back. In the code you shared, you've written: LatencyProcessorTester tester = new Latency_process(); Why do you do that? If a method has added an annotation, a caller of the function doesn't need to know about it. They should be able to use the original definition itself. Or does JavaPoet not allow that? – Chamfer

In the LatencyProcessor implementation in the question, a new class named Latency_$methodName that implements the interface(LatencyProcessorTester in this case) which contains the method with the @Latency annotation is being created. Hence this test code to initialize a Latency_process object. If you want to rewrite the process method within the same class, JavaPoet can probably do that, but the LatencyProcessor implementation would have to be changed. – Rowden

Got it. I'm not looking to replace the class that I wrote, but maybe some generated code. I was hoping for something like mapstruct that generates code for a given interface. AspectJ does what I'm looking for as it does some weaving and abstracts this away, but when you want to debug, you've no visibility into what AspectJ has generated. – Chamfer

@Chamfer MapStruct involves writing a mapper interface. It then generates an implementation at compile-time based on the applied annotations. At run-time, MapStruct uses reflection to find and instantiate the mapper implementation. User code is aware of the mapper interface and is written to explicitly use said mapper. Your goal, however, seems to be adding latency measurements transparently. You want it where user code does not have to be rewritten so the @Latency method takes effect. That is not something an annotation processor can do. – Buote

@Buote Thanks for taking time to add a response. Is my only option to use AspectJ in this case? – Chamfer

@Chamfer AspectJ or similar (e.g., a class-transforming agent). Note measuring execution time like this seems like something that should be optionally enabled, so I don't think compile-time weaving would make much sense. Also, if you want to measure execution time, why not just use an existing profiler? – Buote

hello, sir @devatherock, maybe you could help me. I also have a problem with custom annotation processor. its been 5 days now. im trying to fix the issue were there is no generated files and the processor cannot be found. I setup correctly everything. I did manually add the javax.annotation.processing.Processor inside resources/META-INF/services. I still get the same error. I tried the autor service from google. It still the same. – Bebebebeerine

@Bebebebeerine If you can post a new question with the relevant details and a minimal reproducible example, I can take a look. – Rowden