Transaction handling using an annotation processor

Declarative transaction processing is a very popular feature of Java application frameworks. We can see examples from the three most popular frameworks below.

Spring Framework (example of runtime support)
Micronaut (example of compile-time support)
Quarkus

So why not build a simplified version of your own just to learn how it works? The repository Java Own Framework – step-by-step shows how to do it purely in runtime. I want to show you the compile-time version today.

An important reminder is that this is the second post of the series. The first is available here: Build your own framework using annotation processor – introduction and dependency injection. You do not need to know the previous part by heart, but I recommend reading it or at least glancing through it.

In the first article of the series, I jump-started creating an annotation processor-based framework by providing the first feature – dependency injection.

Note: Micronaut has heavily inspired all code examples in this text.

What exactly is declarative transaction support?

At the end of the previous article, we saw logs for the running app:

1Begin transaction
2Participant: 'Participant[]' takes part in event: 'Event[]'
3Commit transaction

As you can guess, the transaction is already there. It is managed by the TransactionManager instance implemented by ManualTransactionParticipationService. The Begin transaction and Commit transaction messages are printed by a fake TransactionManager implementation called TransactionalManagerStub.

If we take a look at the ManualTransactionParticipationService code:

1@Singleton
2public class ManualTransactionParticipationService implements ParticipationService {
3    private final ParticipantRepository participantRepository;
4    private final EventRepository eventRepository;
5    private final TransactionManager transactionManager;
6
7
8    // constructor
9
10    @Override
11    public void participate(ParticipantId participantId, EventId eventId) {
12        try {
13            transactionManager.begin();
14            var participant = participantRepository.getParticipant(participantId);
15            var event = eventRepository.findEvent(eventId);
16            eventRepository.store(event.addParticipant(participant));
17
18            System.out.printf("Participant: '%s' takes part in event: '%s'%n", participant, event);
19
20            transactionManager.commit();
21        } catch (Exception e) {
22            rollback();
23            throw new RuntimeException(e);
24        }
25    }
26
27    private void rollback() {
28        try {
29            transactionManager.rollback();
30        } catch (SystemException e) {
31            throw new RuntimeException(e);
32        }
33    }
34}

We see that the transaction here adds lots of boilerplate. Wouldn’t it be easier to write code like this:

1@Singleton
2public class DeclarativeTransactionsParticipationService implements ParticipationService {
3    private final ParticipantRepository participantRepository;
4    private final EventRepository eventRepository;
5    // constructor
6
7    @Override
8    @Transactional
9    public void participate(ParticipantId participantId, EventId eventId) {
10            var participant = participantRepository.getParticipant(participantId);
11            var event = eventRepository.findEvent(eventId);
12            eventRepository.store(event.addParticipant(participant));
13            
14            System.out.printf("Participant: '%s' takes part in event: '%s'%n", participant, event);
15    }
16}

From now on, being able to write code like that is our target. We want to handle transactions by adding the @Transactional annotation to the method of our interest.

The @Transactional annotation and TransactionManager

First, it would be beneficial to have an annotation to achieve the outcome. As I wanted to use the standard one instead of writing my own, I chose the @Transactional annotation and the TransactionManager interface from the Jakarta EE Transactions 2.0 specification.

How is the transactional handling going to work?

Once a method is annotated with @Transactional, we want the annotation processor to generate transaction handling code. For the sake of simplicity, the processor will generate code only for methods of concrete classes.

Since the processor can only generate new code, it will create a subclass of the class with annotated methods. Therefore, the class cannot be final. The methods cannot be final, private, or static. Non-annotated methods won’t be touched at all.

To get a better idea, please look at the example below.

For the below class:

1@Singleton
2public class RepositoryA {
3
4    @Transactional
5    void voidMethod() {
6    }
7
8    int intMethod() {
9        return 1;
10    }
11}

The annotation processor should generate the following:

1@Singleton
2class RepositoryA$Intercepted extends RepositoryA {
3  private final TransactionManager transactionManager;
4
5  RepositoryA$Intercepted(TransactionManager transactionManager) {
6    super();
7    this.transactionManager = transactionManager;
8  }
9
10  @Override
11  void voidMethod() {
12    // transaction handling code
13  }
14}

The example presents a simplified version of what will be generated, but you probably get the idea. The actual code generation and other issues will be shown later on. The generated code will be simple. It won’t care about transaction propagation. It will wrap checked exceptions into the RuntimeException and rethrow them in that form.

The problem is that if you want transactions, you cannot directly create an instance of the class with annotated methods using new or any other factory method. You must rely on the framework, in our case, the one described in the previous article of the series, to provide it. This is because only the generated class will have the expected transactional code.

The only extra thing worth noticing in the example is the generated class name. For the rest of this project, if the annotation processor ever creates replacements for some classes, their names will include the Intercepted word.

Handling @Transactional

As transaction handling is the main subject of this text, we will get straight to it.

Processing the @Transactional annotation is not a mandatory part of our framework. It should be used based on the user’s decision. Therefore, the code responsible for it will be called TransactionalPlugin, as this is a pluggable feature.

Let’s look at the code below (the code also is available here).

1public class TransactionalPlugin implements ProcessorPlugin { // 7
2   private TransactionalMessenger transactionalMessenger; // 3
3
4   @Override
5    public Collection<JavaFile> process(Set<? extends Element> annotated) { // 1
6        Set<ExecutableElement> transactionalMethods = ElementFilter.methodsIn(annotated); // 2
7        validateMethods(transactionalMethods); // 3
8        Map<TypeElement, List<ExecutableElement>> typeToTransactionalMethods = transactionalMethods.stream() // 4
9                .collect(groupingBy(element -> (TypeElement) element.getEnclosingElement())); // 4
10        return typeToTransactionalMethods.entrySet()
11                .stream()
12                .map(this::writeTransactional) // 5
13                .toList();
14    }
15
16    private void validateMethods(Set<ExecutableElement> transactionalMethods) { // 3
17        raiseForPrivate(transactionalMethods);
18        raiseForStatic(transactionalMethods);
19        raiseForFinalMethods(transactionalMethods);
20        raiseForFinalClass(transactionalMethods);
21    }
22
23    private JavaFile writeTransactional(Map.Entry<TypeElement, List<ExecutableElement>> typeElementListEntry) { // 5
24        var transactionalType = typeElementListEntry.getKey();
25        var transactionalMethods = typeElementListEntry.getValue();
26        PackageElement packageElement = processingEnv.getElementUtils().getPackageOf(transactionalType);
27        return new TransactionalInterceptedWriter(transactionalType, transactionalMethods, packageElement) // 6
28                .createDefinition(processingEnv.getMessager()); // 6
29    }
30    
31   // more methods ...
32}

Now, it is time for us to dive deeply into the provided source.

Set<? extends Element> annotated contains all Elements annotated with @Transactional.
In the first step, we filter all methods out of the annotated set of elements using ElementFilter.
Then, the annotated elements are validated against the previously mentioned rules. I introduced the utility class TransactionalMessenger (code here). Its sole responsibility is to wrap Messager and provide a unified API for raising errors associated with the @Transactional processing. Every raiseForSth method calls TransactionalMessenger providing information about the error. The raiseForSth methods’ code is skipped to keep the example concise and manageable.
Now, we group the annotated methods by classes that the methods are declared. In Java, you can only create a method in a class or interface. However, the plugin accepts only concrete class methods and raises errors for others. Therefore, we can be sure that calling *element.getEnclosingElement(_)* where an element is the annotated method will return class representation – TypeElement.
Once we have the mentioned grouping, we can write the code. We need to intercept classes that are keys in the grouping and write transactional versions of methods that are values of the mapping.
The last part is to write the code. The logic is stored in TransactionalInterceptedWriter, so we can move to see its code.
As the TransactionalPlugin must be somehow plugged into our framework workings, the class implements the ProcessorPlugin interface. How it all works will be described after we finish with the transaction handling, as it is not the main topic here.

Writing the code with TransactionalInterceptedWriter

For code generation, I will use the proven JavaPoet library.

The code of the TransactionalInterceptedWriter is quite complicated. The thing that requires special attention is writing transactional versions of void methods and value-returning methods. Unfortunately, Java language has the void type contrary to Kotlin, Scala, Rust, and others.

We will get to the mentioned part later. Now let’s start with instance fields and constructor.

Instance fields and constants

The Writer constructor is fairly simple, so it can be omitted.

1class TransactionalInterceptedWriter {
2    private static final String TRANSACTION_MANAGER = "transactionManager";
3    private static final Modifier[] PRIVATE_FINAL_MODIFIERS = {Modifier.PRIVATE, Modifier.FINAL};
4    
5    private final TypeElement transactionalElement; // 1 
6    private final List<ExecutableElement> transactionalMethods; // 2
7    private final PackageElement packageElement; // 3
8}

The constants are fairly simple, and their names are self-explanatory.

The class instance fields are more interesting.

1. The transactionalElement stores the TypeElement representation of the class with the annotated methods. The class will be referred to as intercepted class or superclass.

2. The transactionalMethods stores the ExecutableElement representation of the annotated methods of the transactionalElement class.

3. The packageElement stores the PackageElement representation of the package in which transactionalElement is defined.

Intercepting class definition

We will start with the most high-level thing. Let’s see how the intercepting class is written, but without going into details.

1class TransactionalInterceptedWriter {
2    
3    public JavaFile createDefinition(Messager messager) {
4        TypeSpec typeSpec = TypeSpec.classBuilder("%s$Intercepted".formatted(transactionalElement.getSimpleName().toString())) // 1
5                .addAnnotation(Singleton.class) // 2
6                .superclass(transactionalElement.asType()) // 3
7                .addSuperinterface(TypeName.get(Intercepted.class)) // 4
8                .addMethod(interceptedTypeMethod()) // 4
9                .addField(TransactionManager.class, TRANSACTION_MANAGER, PRIVATE_FINAL_MODIFIERS) // 5
10                .addMethod(constructor(messager)) // 6
11                .addMethods(transactionalMethodDefinitions()) // 7
12                .build();
13        return JavaFile.builder(packageElement.getQualifiedName().toString(), typeSpec).build(); // 8
14    }
15    
16}

First of all, the class must have a name. As mentioned before the generated class will be called the old one but with an extra Intercepted part. For example, Repository will be changed into Repository$Intercepted. Therefore, we know that the type before $ is intercepted by the generated class.
The created class must be annotated with @Singleton, so the DI solution from the first part will pick it up.
To fulfill its role, the generated class will extend the class with methods annotated with @Transactional. We have already talked about it above.
The class will also implement the Intercepted interface, which will be covered later. The interface is related to the provisioning of the intercepted instances. This requires the generated class to implement an extra method. I will describe how it works at the end of the article, as this is unrelated to transactions.
To handle transactions, the class needs a TransactionalManager field. Adding the field is very straightforward.
The class must have a constructor that will call super(requiredDependencies) and set the transactionManager field.
The class will override the methods annotated in its superclass.
The generated code will be stored in the JavaFile object to be written in a real file later.

Now, having the high-level view, we can dive into the details where needed. So let’s start with writing the constructor.

Constructor

To provide the transactional capability, the constructor must call the constructor of its superclass via the super keyword, passing the parameters in the correct order. The transactionManager field of the intercepting class also must be populated.

1class TransactionalInterceptedWriter {
2
3    private MethodSpec constructor(Messager messager) {
4        Dependency dependency = new TypeDependencyResolver().resolve(transactionalElement, messager); // 1
5        var typeNames = dependency.dependencies().stream().map(TypeName::get).toList(); // 1
6        
7        var constructorParameters = typeNames.stream() // 2
8                .map(typeName -> ParameterSpec.builder(typeName, "$" + typeNames.indexOf(typeName)).build()) // 2
9                .toList();
10        
11        var superCallParams = IntStream.range(0, typeNames.size()) // 3
12                .mapToObj(integer -> "$" + integer) // 3
13                .collect(Collectors.joining(", ")); // 3
14
15        return MethodSpec.constructorBuilder()
16                .addParameter(ParameterSpec.builder(TransactionManager.class, TRANSACTION_MANAGER).build()) // 2
17                .addParameters(constructorParameters) // 2
18                .addCode(CodeBlock.builder()
19                        .addStatement("super($L)", superCallParams) // 3
20                        .addStatement("this.$L = $L", TRANSACTION_MANAGER, TRANSACTION_MANAGER) // 4
21                        .build())
22                .build();
23    }
24}

The first thing that is done to create a constructor is to find out what the dependencies of the intercepted class are. To do it in a convenient way, we will reuse TypeDependencyResolver, created for the DI solution. You can read more about it here.
Having the dependencies of the superclass, we can create parameters for the constructor. The transactionManager is the first parameter, and the rest is provided conveniently as Type ${position in the constructorParameters list}.
Having the intercepting class constructor params, we can prepare the content of the super call. Then it can be added to the super call in the constructor.
The last thing to do is to also set up the transactionManager field.

The generated constructor may look like the code below:

1class TestRepository$Intercepted {
2   TestRepository$Intercepted(TransactionManager transactionManager,
3                              ParticipantRepository $0, 
4                              EventRepository $1) {
5      super($0, $1);
6      this.transactionManager = transactionManager;
7   }
8}

Overriding transactional methods

In the case of generating the methods, we will start with an example.

1@Singleton
2class RepositoryA$Intercepted extends RepositoryA { 
3    
4  @Override
5  void voidMethod() {
6    try {
7      transactionManager.begin();
8      super.voidMethod();
9      transactionManager.commit();
10    }
11    catch (Exception e) {
12      try {
13        transactionManager.rollback();
14      }
15      catch (Exception innerException) {
16        throw new RuntimeException(innerException);
17      }
18      throw new RuntimeException(e);
19    }
20  }
21
22  @Override
23  int intMethod() {
24    try {
25      transactionManager.begin();
26      var intMethodReturnValue = (int) super.intMethod();
27      transactionManager.commit();
28      return intMethodReturnValue;
29    }
30    catch (Exception e) {
31      try {
32        transactionManager.rollback();
33      }
34      catch (Exception innerException) {
35        throw new RuntimeException(innerException);
36      }
37      throw new RuntimeException(e);
38    }
39  }
40}

TransactionManager’s methods are defined with the checked exception. Therefore, transactional methods need to include try/catch blocks. In the try block, the begin and commit must be called, as well as a rollback in the catch clause. If the return type isn’t void, the result of super method call results must be stored in a variable.

So the high-level method to generate such a call looks like this:

1class TransactionalInterceptedWriter {
2   private MethodSpec generateTransactionalMethod(ExecutableElement executableElement) {
3      var methodName = executableElement.getSimpleName().toString();
4      var transactionalMethodCall = transactionalMethodCall(executableElement);
5      var methodCode = tryClause(transactionalMethodCall, catchClause());
6      return MethodSpec.methodBuilder(methodName)
7              .addModifiers(executableElement.getModifiers())
8              .addParameters(executableElement.getParameters().stream().map(ParameterSpec::get).toList())
9              .addAnnotation(Override.class)
10              .addCode(methodCode)
11              .returns(TypeName.get(executableElement.getReturnType()))
12              .addTypeVariables(getTypeVariableIfNeeded(executableElement).stream().toList())
13              .build();
14   }
15}

In this and the previous article of this series, I have shown a lot of code, mostly containing JavaPoet usage. My hope is that now you get how the JavaPoet works. From now on, I will try to minimize the boilerplate JavaPoet code by omitting it in the examples or sharing it as Gists. The full code is still present in the repository, of course.

Catch and try blocks

The try and catch blocks code is quite simple. So as mentioned before, here are the gists:

https://gist.github.com/JacekDubikowski/167bcaaab151f9d4ad6f033ef1543cec

https://gist.github.com/JacekDubikowski/8b691faf3e0aba2e04d03211823794ff

Super method calls

Once we have try and catch blocks handled, we can focus on the actual super method call.

1class TransactionalInterceptedWriter {
2   private CodeBlock transactionalMethodCall(ExecutableElement executableElement) {
3      return executableElement.getReturnType().getKind() == TypeKind.VOID // 1
4              ? transactionalVoidCall(executableElement)
5              : returningTransactionalMethodCall(executableElement);
6   }
7
8   private CodeBlock transactionalVoidCall(ExecutableElement method) { // 2
9      var params = translateMethodToSuperCallParams(method);
10      return CodeBlock.builder()
11              .addStatement(TRANSACTION_MANAGER + ".begin()")
12              .addStatement("super.$L(%s)".formatted(params), method.getSimpleName())
13              .addStatement(TRANSACTION_MANAGER + ".commit()")
14              .build();
15   }
16
17   private CodeBlock returningTransactionalMethodCall(ExecutableElement method) { // 3
18      var methodName = method.getSimpleName();
19      var params = translateMethodToSuperCallParams(method);
20      return CodeBlock.builder()
21              .addStatement(TRANSACTION_MANAGER + ".begin()")
22              .addStatement("var $LReturnValue = ($L) super.$L(%s)".formatted(params), methodName, method.getReturnType(), methodName)
23              .addStatement(TRANSACTION_MANAGER + ".commit()")
24              .addStatement("return $LReturnValue", methodName)
25              .build();
26   }
27
28   private String translateMethodToSuperCallParams(ExecutableElement method) { 
29      // just code
30   }
31}

The code generation is really simple here.

The first step is deciding upon the method call based on the return type of the super method.
A void call is generated. This is very simple, as there is no need to store the results of the super call.
Finally, the value-returning method call is generated. The result is stored in a variable to be returned after the commit.

The full code

The full code of the TransactionalInterceptedWriter.

That’s all for transactions

This is everything I have prepared for you in transaction handling.

The code below could be used, and support for transactions can be provided.

1@Singleton
2public class DeclarativeTransactionsParticipationService implements ParticipationService {
3    private final ParticipantRepository participantRepository;
4    private final EventRepository eventRepository;
5
6    public DeclarativeTransactionsParticipationService(
7            ParticipantRepository participantRepository, 
8            ventRepository eventRepository
9    ) {
10        this.participantRepository = participantRepository;
11        this.eventRepository = eventRepository;
12    }
13
14    @Override
15    @Transactional
16    public void participate(ParticipantId participantId, EventId eventId) {
17            var participant = participantRepository.getParticipant(participantId);
18            var event = eventRepository.findEvent(eventId);
19            eventRepository.store(event.addParticipant(participant));
20            
21            System.out.printf("Participant: '%s' takes part in event: '%s'%n", participant, event);
22    }
23}

However, we must be sure to get the expected instance during the runtime, right? Let us check how to make it all work within our framework. To reach the goal, we need two more things.

The TransactionPlugin must be used during compilation.
We must make our framework provide only intercepted instances.

Plugging the transaction handling into the framework

We have seen the code that handles transaction processing. Now, we have to make use of the TransactionalPlugin in our framework. To keep everything simple, I created an interface ProcessorPlugin which will be a way to register extensions. Thanks to that, the whole transaction processing code is held in separate classes.

1public interface ProcessorPlugin {
2    void init(ProcessingEnvironment processingEnv); // 1
3
4    Collection<JavaFile> process(Set<? extends Element> annotated); // 2
5
6    Class<? extends Annotation> reactsTo(); // 3
7}

The interface has three methods.

The init method is responsible for the initialization of the plugin.
The process method does the actual processing. Therefore, it returns generated Java files.
The reactsTo method provides information about annotation that the plugin is interested in.

The plugins are hardwired so far and are used as presented in the code:

1public class BeanProcessor extends AbstractProcessor {
2    private List<ProcessorPlugin> plugins = List.of();
3
4    @Override
5    public synchronized void init(ProcessingEnvironment processingEnv) { // 1
6        super.init(processingEnv);
7        plugins = List.of(new TransactionalPlugin());
8        plugins.forEach(processorPlugin -> processorPlugin.init(processingEnv));
9    }
10
11    @Override
12    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) { // 2
13        try {
14            runPluginsProcessing(roundEnv);
15            // rest of the processing 
16        } catch (Exception e) {
17           // exception handling
18        }
19        // return
20    }
21
22    private void runPluginsProcessing(RoundEnvironment roundEnv) { // 3
23        plugins.stream().map(processorPlugin -> processorPlugin.process(roundEnv.getElementsAnnotatedWith(processorPlugin.reactsTo())))
24                .flatMap(Collection::stream)
25                .forEach(this::writeFile); // 4
26    }
27    
28    private void writeFile(JavaFile javaFile) {} // 4
29}

As you can see:

Once the processor is initialized, it also initializes its plugins. In the more real-world code, the plugin discovery could possibly be run here.
The processor starts its processing by running the plugins.
All the plugins are run with elements annotated with the annotation that the plugin reacts to.
The files generated by plugins are written to the actual files in some /generated directory.

Implementation of the ProcessorPlugin for the TransactionalPlugin

In the main part of the article, I omitted some code of the TransactionalPlugin related to the ProcessorPlugin implementation. Now, you can see the missing parts of the code below.

1public class TransactionalPlugin implements ProcessorPlugin { 
2    @Override
3    public void init(ProcessingEnvironment processingEnv) { // 1
4        this.processingEnv = processingEnv;
5        this.transactionalMessenger = new TransactionalMessenger(processingEnv.getMessager());
6    }
7
8    @Override
9    public Class<? extends Annotation> reactsTo() { //2
10        return Transactional.class;
11    }
12}

The init method implementation is fairly simple. It just requires setting processingEnv and creating TransactionalMessenger.
The reactsTo method implementation states that the plugin is interested in @Transactional annotation. Who would guess, right?

The provided code is nothing big. It is easy to notice that the most interesting thing was the process method shown before.

Provisioning of intercepted class

In the “production” code, the framework must be able to provision the intercepted instances. To make this possible, I introduced the interface below.

1public interface Intercepted {
2    Class<?> interceptedType();
3}

This is very simple, yet very important. Thanks to the interface, we can be sure which type has its intercepted version. You may have remembered from the main part that our $Intercepted classes have implemented the interface. So how was this done?

Implementing the interface

The implementation of the interface is quite simple. For the RepositoryA:

1@Singleton
2public class RepositoryA {
3    // some @Transactional methods
4}

It will be implemented as:

1@Singleton
2class RepositoryA$Intercepted extends RepositoryA {
3   @Override
4   public Class interceptedType() {
5      return RepositoryA.class;
6   }
7   // Overridden transactional methods
8}

In the source code of TransactionalInterceptedWriter, it would just add a few extra lines:

1class TransactionalInterceptedWriter {
2   private MethodSpec interceptedTypeMethod() {
3      return MethodSpec.methodBuilder("interceptedType")
4              .addAnnotation(Override.class)
5              .addModifiers(PUBLIC)
6              .addStatement("return $T.class", TypeName.get(transactionalElement.asType()))
7              .returns(ClassName.get(Class.class))
8              .build();
9   }
10}

Now, we can differentiate Intercepted classes from regular ones and point out types that have their intercepted versions.

Using only intercepting classes during provisioning

To get only the intercepted version and not the original one, we need to update the BaseBeanProvider. The simplified code was shown in the previous part about DI. Now, it needs an extra step.

1class BaseBeanProvider implements BeanProvider {
2   @Override
3   public <T> List<T> provideAll(Class<T> beanType) {
4      var allBeans = definitions.stream().filter(def -> beanType.isAssignableFrom(def.type()))
5              .map(def -> beanType.cast(def.create(this)))
6              .toList(); // 1
7      var interceptedTypes = allBeans.stream().filter(bean -> Intercepted.class.isAssignableFrom(bean.getClass()))
8              .map(bean -> ((Intercepted) bean).interceptedType())
9              .toList(); // 2
10      return allBeans.stream().filter(not(bean -> interceptedTypes.contains(bean.getClass()))).toList(); // 3
11   }
12}

Firstly, we find all the beans matching the needed type.
Then, we find the beans that implement the Intercepted type.
In the end, we return a list of the matching beans filtering out the beans that are among the types with their intercepted version.

It works

Now the whole solution works as expected. The framework provides the $Intercepted instances that handle transactions for us. I think that is enough for today. Class dismissed.

In the next and final part, we will a look at RestControllers, so stay tuned!