Using Aspect Oriented Programming (AOP) to automatically batch method calls

Introduction

A while ago we ran into an interesting problem with one of our DAO calls. A seemingly innocent function started misbehaving:

The Postgres driver we used in our code has a limit of 32767 arguments that it may bind in 1 statement. A worst case scenario had broken that limit.

We would have liked to have a database table structure that meant we didn’t need to query thousands of objects at once — but that’s not the topic of this article.

There were plenty of similar examples in our code base that lacked the logic to partition big input collections and we found ourselves looking for the easiest path to address this problem.

A naive approach of editing one method at a time proved to be a repetitive task. We ended up adding very similar logic to method calls, to partition the input correctly and aggregate the results.

Rather than repeating this step, we started looking for a solution that would allow us to intercept the call at the right time and partition the input, without having to create delegating classes to handle this case. Below we discuss how we solved the problem using AOP.

AOP to the rescue

AOP stands for Aspect Oriented Programming. It is similar to ‘middleware’. An Aspect wraps the real implementation, and can execute code before and after the ‘real’ method is called.

Using Spring AOP, we can implement an Aspect that can partition any collection input, and delegate the resulting partitions correctly to the implementation. This will solve the problem of partitioning across different implementations.

It consist of several parts, we will discuss the following for our solution:

• Aspect The implementation that we want to execute
• JoinPoint A point during code execution. In Spring, this is always a method execution
• Pointcut Defines the predicate that enables the Aspect. It matches a specific JoinPoint, for example a method execution or an annotation.
• Advice This defines when an action is taken, for example after a method returns. It is associated with a Pointcut.
• Proxy Bean A proxy object that is created by the AOP framework. It allows the framework to implement the contracts of the Aspect.

In Spring, we can register a specific Aspect that will take effect when some defined criteria is met. Internally, Spring creates a proxy of our implementation, allowing us to intercept a call to this bean.

We will discuss the different steps in more detail next.

Defining the Aspect

We needed to define the logic we want to execute when we intercept our method call. We wanted to be able to execute our Aspect regardless of the underlying implementation, so we needed to split up the work into generic tasks that apply to all implementations.

Our Aspect has 4 different responsibilities:

1. Find the correct input variable
2. Partition the input
3. Invoke the implementation
4. Aggregate the results for the partition

Finding the correct Join Points

Our code base is large and we have many different methods taking a Collection input, not all of them qualifying to be partitioned. AOP would allow us to match all of these, however this would present us with a different problem: how would we choose which ones actually need partitioning?

We decided to go with an annotation driven approach. We define 2 separate annotations:

• Batch This marks our JoinPoint. Any method annotated with @Batch will be proxied by Spring and intercepted with our Aspect.
• BatchInput This is an optional annotation and marks the specific input variable to partition. A function call can have any number of arguments, multiple of which can be a Collection. In this case, the annotation provides the necessary information.

The following code implements our annotations:

With these 2 annotations, the above function may look like this:

Finding the correct input variable

Before, we discussed adding a BatchInput annotation to mark the correct input in our method signature. We could be intercepting a range of different calls where the input may be in different positions, for example:

Our implementation takes advantage of the annotation to determine the correct input argument. The following implementation has a default for methods with only 1 argument, while searching all arguments in case there are many. It returns the index of the argument we need to partition.

Partition the input

Our code makes the reasonable assumption that if a method signature is annotated with a Batch annotation, it provides an input of the type Collection. To partition the input, we may therefore simply cast our variable and invoke Guava's excellent Iterables utility to partition the input.

Invoking the implementation

The partitions need to be submitted to the implementation. We replace the argument found at the index with the correct partition and call on our Pointcut to proceed with this input.

Aggregating the results

We can easily partition the input, and call the real method implementation multiple times. However, our code is expecting a single return value. We need to aggregate the results of the individual calls to the underlying implementation. We define an interface that is responsible to abstract this detail away from our Aspect:

The method calls we needed to intercept have different return types. We found that we needed to be able to aggregate results of Map and List types. For these, we implement two different aggregators:

With these 2 as our defaults in place, we need to determine the correct one to use. The method signature of our PointCut provides this information to the Aspect. We can inspect the method and determine the correct aggregator:

Bringing it together

With our 4 parts in place, we only need to bring them together. We define our PointCut and Advice and publish this to the Spring context:

We define that any call to a method annotated with the Batch annotation should be intercepted by our batchProcess function.

Conclusion

Aspect Oriented Programming can help solve generic problems across implementation. In our case, it allowed us to quickly apply batch processing across all our database access code without invasive code changes.

Full example

You can find the full source here:

https://github.com/pandaadb/aop-batcher

Want to work with us? Here’s our careers page 🎉