CockroachDB Follower Reads with Spring and AOP

How I built Flight Chop Series

·

4 min read

CockroachDB Follower Reads with Spring and AOP

For the past few years I've been building a small app called Flight Chop using some of my favorite technologies including the Spring Framework and CockroachDB. Flight Chop collects data about current flight and weather conditions from various government and commercial services to provide passengers with a sense of how much turbulence they are likely to encounter while in the air. It is very much a work in progress.

Flight Chop is mostly a read-only workload that does not require up to the second, real-time data. As a result, it's a perfect use case for CockroachDB Follower Reads. In short, Follower Reads allow very fast access to a consistent snapshot of data AS OF some time in the past. One of the many benefits of this feature is that this data can be served from any "follower" replica in a Cockroach cluster instead of from the leaseholder. This can lead to significant performance gains, especially in multi-region or other non-trivial deployments.

The old way...

When I first started building Flight Chop I used handwritten queries and manually included the Follower Read syntax. For example, many of my data access methods looked something like this...

@Transactional(readOnly = true)
public Aircraft getAircraft(String id) {

  final String sql = """
          SELECT *
          FROM aircraft a
               WHERE a.id = ?
          AS OF SYSTEM TIME follower_read_timestamp()
          """;

  return jdbcTemplate.queryForObject(sql, Aircraft.class, id);
}

As you can see I was manually inserting AS OF SYSTEM TIME follower_read_timestamp() into the body of the query. Functionally this worked well but handwriting SQL for every query including simple look ups was time consuming and error prone. I really wanted to leverage JPA as much as possible but didn't want to lose the performance gains offered by Follower Reads.

AOP to the rescue...

Luckily CockroachDB added support for applying the AS OF SYSTEM TIME syntax at the transaction level in addition to the query level. This was a huge win for productivity as it meant I no longer needed to manually write queries and could leverage JPA to generate the queries for me. To do this cleanly, I combined Spring's Aspect Oriented Programming (AOP) support, a custom @interface (aka Annotation) and a custom @Aspect to inject the Follower Read syntax wherever I wanted.

The first step was to create simple @interface or annotation which could be added to an appropriate method signature. This acts as a marker which AOP uses to find and inject behavior.

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

With the annotation in place, I created the FollowerReadAspect which looks for instances of @FollowerRead and executes the followerRead method before executing the annotated getAircraft method.

@Component
@Aspect
public class FollowerReadAspect {

    private final Logger log = LoggerFactory.getLogger(getClass());

    private static final String SQL = "SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()";

    private final JdbcTemplate jdbcTemplate;

    public FollowerReadAspect(JdbcTemplate jdbcTemplate) {
        this.jdbcTemplate = jdbcTemplate;
    }

    @Before("@annotation(FollowerRead)")
    public void followerRead() {
        boolean actualTransactionActive = TransactionSynchronizationManager.isActualTransactionActive();

        if (!actualTransactionActive) {
            throw new RuntimeException("No transaction is active!  Make sure to use the @Transactional annotation with the @FollowerRead annotation.");
        }

        boolean currentTransactionReadOnly = TransactionSynchronizationManager.isCurrentTransactionReadOnly();

        if (!currentTransactionReadOnly) {
            log.warn("Transaction is active but NOT ready-only.  While this may work it is unlikely that this is intentional.  Consider using @Transactional(readOnly = true).");
        }

        jdbcTemplate.execute(SQL);
    }
}

The followerRead method checks the status and type of the current transaction and then uses Spring's JdbcTemplate to enable Follower Read behavior for any subsequent SQL statement executed during the transaction.

With this approach I can now "inject" Follower Read behavior into any method that queries the database. What's great about this approach is that this works equally well for JPA style queries as well as more complex, handwritten queries. Lets look at how my original method has changed...

@Transactional(readOnly = true)
@FollowerRead
public Optional<Aircraft> getAircraft(String id) {
  return aircraftRepository.findById(id);
}
public interface AircraftRepository extends PagingAndSortingRepository<Aircraft, String> { 
}

I've added the @FollowerRead annotation to the method. Since I no longer have to write custom SQL, I can leverage a JPA PagingAndSortingRepository for this simple lookup (see AircraftRepository). There is less code, it's easier to read and takes full advantage of Follower Read optimizations in CockroachDB.