Writing your own lombok annotation

It took me quite a while to get around to writing my own lombok annotation and processor. This took more effort than I expected, hopefully this post will save someone else some.

tl;dr — Look at the source in my Github repo.

Motivation

Reading the excellent ExecutorService - 10 tips and tricks by Tomasz Nurkiewicz, I thought about tip #2, Switch names according to context, which recommends wrapping important methods and code blocks with custom thread names to aid in logging and debugging.

"This is a great use case for annotations!" I thought. The code screams boilerplate:

public void doNiftyThings() {
    final Thread thread = Thread.currentThread();
    final String oldName = thread.getName();
    thread.setName("Boy this is nifty!");
    try {
        // Do those nifty things - the actual work
    } finally {
        thread.setName(oldName);
    }
}

The whole point of the method is indented out of focus, wrapped with bookkeeping. I'd rather write this:

@ThreadNamed("Boy this is nifty!")
public void doNiftyThings() {
    // Do those nifty thing - the actual work
}

Bonus: simple text search finds those places in my code where I change the thread name based on context.

Writing the annotation

Ok, let's make this work. I started with cloning the @Cleanup annotation and processor, and editing from there. First the annotation, the easy bit. I include the javadoc to emphasize the importance of documenting your public APIs.

/**
 * {@code ThreadNamed} sets the thread name during method execution, restoring
 * it when the method completes (normally or exceptionally).
 *
 * @author <a href="mailto:binkley@alumni.rice.edu">B. K. Oxley (binkley)</a>
 */
@Documented
@Retention(SOURCE)
@Target({CONSTRUCTOR, METHOD})
public @interface ThreadNamed {
    /** The name for the thread while the annotated method executes. */
    String value();
}

Nothing special here. I've made the decision to limit the annotation to methods and constructors. Ideally I'd include blocks but that isn't an option (yet) in Java, and you can always refactor out a block to a method.

Writing the processor

This is the serious part. First some preliminaries:

1. I have only implemented support for JDK javac. Lombok also supports the Eclipse compiler, which requires a separate processor class. I have nothing against Eclipse, but it's not in my toolkit.
2. I'll discuss library dependencies below. For now pretend these are already working for you.
3. I'm a big fan of static imports, the diamond operator, etc. I don't like retyping what the compiler is already thinking. You should note List below is not java.util.List; it's com.sun.tools.javac.util.List. Yeah, I don't know this class either.
4. The implementation is hard to follow. Most of us don't spend much time with expression trees, which is how most compilers (including javac) see your source code. A language like LISP lets you write you code as the expression tree directly, which is both nifty and challenging (macros being like annotation processors).

Without further ado:

/**
 * Handles the {@code lombok.ThreadNamed} annotation for javac.
 */
@MetaInfServices(JavacAnnotationHandler.class)
@HandlerPriority(value = 1024)
// 2^10; @NonNull must have run first, so that we wrap around the
// statements generated by it.
public class HandleThreadNamed
        extends JavacAnnotationHandler<ThreadNamed> {
    /**
     * lombok configuration: {@code lab.lombok.threadNamed.flagUsage} = {@code
     * WARNING} | {@code ERROR}.
     * <p>
     * If set, <em>any</em> usage of {@code @ThreadNamed} results in a warning
     * / error.
     */
    public static final ConfigurationKey<FlagUsageType>
            THREAD_NAMED_FLAG_USAGE = new ConfigurationKey<FlagUsageType>(
            "lab.lombok.threadNamed.flagUsage",
            "Emit a warning or error if @ThreadNamed is used.") {
    };

    @Override
    public void handle(final AnnotationValues<ThreadNamed> annotation,
            final JCAnnotation ast, final JavacNode annotationNode) {
        handleFlagUsage(annotationNode, THREAD_NAMED_FLAG_USAGE,
                "@ThreadNamed");

        deleteAnnotationIfNeccessary(annotationNode, ThreadNamed.class);
        final String threadName = annotation.getInstance().value();
        if (threadName.isEmpty()) {
            annotationNode.addError("threadName cannot be the empty string.");
            return;
        }

        final JavacNode owner = annotationNode.up();
        switch (owner.getKind()) {
        case METHOD:
            handleMethod(annotationNode, (JCMethodDecl) owner.get(),
                    threadName);
            break;
        default:
            annotationNode.addError(
                    "@ThreadNamed is legal only on methods and constructors"
                            + ".");
            break;
        }
    }

    public void handleMethod(final JavacNode annotation,
            final JCMethodDecl method, final String threadName) {
        final JavacNode methodNode = annotation.up();

        if ((method.mods.flags & Flags.ABSTRACT) != 0) {
            annotation.addError(
                    "@ThreadNamed can only be used on concrete methods.");
            return;
        }

        if (method.body == null || method.body.stats.isEmpty()) {
            generateEmptyBlockWarning(annotation, false);
            return;
        }

        final JCStatement constructorCall = method.body.stats.get(0);
        final boolean isConstructorCall = isConstructorCall(constructorCall);
        List<JCStatement> contents = isConstructorCall
                ? method.body.stats.tail : method.body.stats;

        if (contents == null || contents.isEmpty()) {
            generateEmptyBlockWarning(annotation, true);
            return;
        }

        contents = List
                .of(buildTryFinallyBlock(methodNode, contents, threadName,
                        annotation.get()));

        method.body.stats = isConstructorCall ? List.of(constructorCall)
                .appendList(contents) : contents;
        methodNode.rebuild();
    }

    public void generateEmptyBlockWarning(final JavacNode annotation,
            final boolean hasConstructorCall) {
        if (hasConstructorCall)
            annotation.addWarning(
                    "Calls to sibling / super constructors are always "
                            + "excluded from @ThreadNamed;"
                            + " @ThreadNamed has been ignored because there"
                            + " is no other code in " + "this constructor.");
        else
            annotation.addWarning(
                    "This method or constructor is empty; @ThreadNamed has "
                            + "been ignored.");
    }

    public JCStatement buildTryFinallyBlock(final JavacNode node,
            final List<JCStatement> contents, final String threadName,
            final JCTree source) {
        final String currentThreadVarName = "$currentThread";
        final String oldThreadNameVarName = "$oldThreadName";

        final JavacTreeMaker maker = node.getTreeMaker();
        final Context context = node.getContext();

        final JCVariableDecl saveCurrentThread = createCurrentThreadVar(node,
                maker, currentThreadVarName);
        final JCVariableDecl saveOldThreadName = createOldThreadNameVar(node,
                maker, currentThreadVarName, oldThreadNameVarName);

        final JCStatement changeThreadName = setThreadName(node, maker,
                maker.Literal(threadName), currentThreadVarName);
        final JCStatement restoreOldThreadName = setThreadName(node, maker,
                maker.Ident(node.toName(oldThreadNameVarName)),
                currentThreadVarName);

        final JCBlock tryBlock = setGeneratedBy(maker.Block(0, contents),
                source, context);
        final JCTry wrapMethod = maker.Try(tryBlock, nil(),
                maker.Block(0, List.of(restoreOldThreadName)));

        if (inNetbeansEditor(node)) {
            //set span (start and end position) of the try statement and
            // the main block
            //this allows NetBeans to dive into the statement correctly:
            final JCCompilationUnit top = (JCCompilationUnit) node.top()
                    .get();
            final int startPos = contents.head.pos;
            final int endPos = Javac
                    .getEndPosition(contents.last().pos(), top);
            tryBlock.pos = startPos;
            wrapMethod.pos = startPos;
            Javac.storeEnd(tryBlock, endPos, top);
            Javac.storeEnd(wrapMethod, endPos, top);
        }

        return setGeneratedBy(maker.Block(0,
                        List.of(saveCurrentThread, saveOldThreadName,
                                changeThreadName, wrapMethod)), source,
                context);
    }

    private static JCVariableDecl createCurrentThreadVar(final JavacNode node,
            final JavacTreeMaker maker, final String currentThreadVarName) {
        return maker.VarDef(maker.Modifiers(FINAL),
                node.toName(currentThreadVarName),
                genJavaLangTypeRef(node, "Thread"), maker.Apply(nil(),
                        genJavaLangTypeRef(node, "Thread", "currentThread"),
                        nil()));
    }

    private static JCVariableDecl createOldThreadNameVar(final JavacNode node,
            final JavacTreeMaker maker, final String currentThreadVarName,
            final String oldThreadNameVarName) {
        return maker.VarDef(maker.Modifiers(FINAL),
                node.toName(oldThreadNameVarName),
                genJavaLangTypeRef(node, "String"),
                getThreadName(node, maker, currentThreadVarName));
    }

    private static JCMethodInvocation getThreadName(final JavacNode node,
            final JavacTreeMaker maker, final String currentThreadVarNAme) {
        return maker.Apply(nil(),
                maker.Select(maker.Ident(node.toName(currentThreadVarNAme)),
                        node.toName("getName")), nil());
    }

    private static JCStatement setThreadName(final JavacNode node,
            final JavacTreeMaker maker, final JCExpression threadName,
            final String currentThreadVarName) {
        return maker.Exec(maker.Apply(nil(),
                maker.Select(maker.Ident(node.toName(currentThreadVarName)),
                        node.toName("setName")), List.of(threadName)));
    }
}

Wasn't that easy?

Dependencies

Of course the code depends on lombok. I'm using version 1.14.8. It also needs tools.jar from the JDK for compiler innards like expression trees. (An Eclipse processor needs an equivalent.)

Unfortunately lombok itself uses "mangosdk" to generate a META-INF/services/lombok.javac.JavacAnnotationHandler file for autodiscovery of processors. I say 'unfortunately' because this library is not in maven and is unsupported. Happyily Kohsuke Kawaguchi wrote the excellent metainf-services library a while back, maintains it, and publishes to Maven central. If you're new to annotation processors it's a good project to learn from.

Conclusion

Ok, that was not actually so easy. On the other hand, finding a starting point was the biggest hurdle for me in writing a lombok annotation. Please browse my source and try your hand at one.

UPDATE — A little bonus. This code:

@ThreadNamed("Slot #%2$d")
public void doSomething(final String name, final int slot) {
    // Do something with method params
}

Produces the thread name "Slot #2" when called with "Foo", 2. Strings without formatting or methods with params treat the annotation value as a plain string.