Session level Hibernate JDBC batching

Quite some time ago, I had a chance to try out one of the coolest features of Hibernate when I had to insert and update multiple entities. By multiple, I mean thousands of entities and obviously this was a slow process. The solution which I was experimenting with was the JDBC Batching.

This is a really great feature of Hibernate. I don’t want to give another detailed explanation about this because there are plenty of them on the web. One of these is made by Vlad Mihalcea. I really suggest reading this one to get familiar with the basics before going forward.

As I mentioned previously, this is one of the most useful features when it comes to performance, related to inserting and updating entities. However, what Hibernate offered until now is to enable the JDBC Batching globally in the whole application. Luckily, there was a feature request to the Hibernate devs to implement the batching in a session basis which is already implemented and released with Hibernate 5.2. The main problem with the global switch was that it caused some performance degradation which was unacceptable in my case. Unfortunately I didn’t have the time to properly analyze what caused the performance issue in specific cases so I stopped experimenting with JDBC Batching. At that time, I though “a session level batching would be a great feature” and now it is in place. This makes it possible to utilize the Hibernate batching for only a set of use cases.

Let’s see how we can use this new session level batch size support.

With Hibernate 5.2, a new method is introduced to the SharedSessionContract interface which is called setJdbcBatchSize. This interface is extended by the Session interface, so if somehow we can get a Session object, this method can be called. I’m going to show you some ideas about using this new feature. The implementation will be based on Spring and is available on GitHub.

I’m starting with the Spring configuration, because I think it’s necessary to give an explanation on it. The current Spring release versions are not supporting Hibernate 5.2 properly, so you need to use some special classes which I’m going to show.

The configuration for Hibernate with EntityManager may look like the following:

<bean id="jpaVendorAdapter" class="org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter">
    <property name="showSql" value="false" />
    <property name="generateDdl" value="false" />
    <property name="database" value="MYSQL" />
</bean>

<bean id="entityManagerFactory" class="org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean">
    <property name="dataSource" ref="dataSource" />
    <property name="jpaVendorAdapter" ref="jpaVendorAdapter" />
    <property name="packagesToScan">
        <list>
            <value>com.arnoldgalovics.blog</value>
        </list>
    </property>
</bean>

If you try to start the application with this bean definitions and with the new Hibernate, you will get the following error:

Caused by: java.lang.NoSuchMethodError: org.hibernate.Session.getFlushMode()Lorg/hibernate/FlushMode;
	at org.springframework.orm.jpa.vendor.HibernateJpaDialect.prepareFlushMode(HibernateJpaDialect.java:187)
	at org.springframework.orm.jpa.vendor.HibernateJpaDialect.beginTransaction(HibernateJpaDialect.java:173)
	at org.springframework.orm.jpa.JpaTransactionManager.doBegin(JpaTransactionManager.java:380)

Currently, none of the Spring releases are supporting this Hibernate version, but with a little trick, this can be bypassed easily. Just a short explanation; the Session interface has been changed and the getFlushMode method now returns a javax.persistence.FlushModeType rather than the expected org.hibernate.FlushMode. Fortunately, there is a method to get the Hibernate FlushMode, called getHibernateFlushMode. To overcome the error, we should create a custom HibernateJpaDialect which acquires the FlushMode according to the new API. This can be done by creating an extension for the HibernateJpaDialect class and override the prepareFlushMode method. Note that to implement the same behavior as originally, we should copy the superclass implementation for this and just call the new getHibernateFlushMode method. The extension should look like this:

public class CustomHibernateJpaDialect extends HibernateJpaDialect {
    private static final long serialVersionUID = -3483592073952499113L;

    @Override
    protected FlushMode prepareFlushMode(final Session session, final boolean readOnly) throws PersistenceException {
        final FlushMode flushMode = session.getHibernateFlushMode();
        if (readOnly) {
            // We should suppress flushing for a read-only transaction.
            if (!flushMode.equals(FlushMode.MANUAL)) {
                session.setFlushMode(FlushMode.MANUAL);
                return flushMode;
            }
        } else {
            // We need AUTO or COMMIT for a non-read-only transaction.
            if (flushMode.lessThan(FlushMode.COMMIT)) {
                session.setFlushMode(FlushMode.AUTO);
                return flushMode;
            }
        }
        // No FlushMode change needed...
        return null;
    }
}

The NoSuchMethodError is fixed with this custom JpaDialect, but now the HibernateJpaVendorAdapter should use this one instead of the default HibernateJpaDialect. Unfortunately this is not configurable, but we can create a new JpaVendorAdapter based on the existing one which returns the correct CustomHibernateJpaDialect. We just need to extend the HibernateJpaDialect and override the getJpaDialect method.

public class CustomJpaVendorAdapter extends HibernateJpaVendorAdapter {
    private final HibernateJpaDialect jpaDialect = new CustomHibernateJpaDialect();

    @Override
    public HibernateJpaDialect getJpaDialect() {
        return jpaDialect;
    }
}

Now, we can register our custom JpaVendorAdapter which we have adapted to the new Hibernate version. After this, the ApplicationContext should start up correctly. Note that this custom fix is required only until Spring releases a new version which supports Hibernate 5.2 (this will be 4.3 GA). This is only covering the batching part of Hibernate, so you can expect more errors due to API changes and I wouldn’t recommend using this in a production environment.

Now we have a spring context which is working, let’s take a look at on the session level batching usage. The naive implementation may look like this for setting this property value.

public class BatchingService {
    @PersistenceContext
    private EntityManager entityManager;

    @Transactional
    public void updateWithSessionBatching() {
        final Session session = entityManager.unwrap(Session.class);
        session.setJdbcBatchSize(15);
        // do something
    }
}

This is working fine, however for every use case, we have to use the same boilerplate code. From code quality point of view, we want to prevent this and make it somehow reusable. One can think of a method which sets this up, but that would be the same boilerplate code as this (or just a little bit better :-)).

What I can imagine is an annotation which has a field to set the batch size and it’s automatically handled. Let’s say we have an annotation called @JdbcBatchSize with the following definition:

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface JdbcBatchSize {
    int value();
}

The usage would be similar to the following:

@Transactional
@JdbcBatchSize(15)
public void updateWithSessionBatching() {
    // do something
}

This looks much cleaner than the previous one but now we have to process it somehow. This can be done by Spring’s AOP framework. I won’t go into details here as well because you can check out Spring’s documentation about this part.

There are multiple choices here to implement the processing, I chose to use a MethodInterceptor with a AbstractPointcutAdvisor. The MethodInterceptor looks like the following:

public class JdbcBatchSizeMethodInterceptor implements MethodInterceptor {
    @PersistenceContext
    private EntityManager entityManager;

    @Override
    public Object invoke(final MethodInvocation invocation) throws Throwable {
        final Session session = entityManager.unwrap(Session.class);
        session.setJdbcBatchSize(invocation.getMethod().getAnnotation(JdbcBatchSize.class).value());
        return invocation.proceed();
    }
}

The interceptor basically does the same as in the naive implementation, sets the batch size for the current session.

Edit: Oliver Gierke just mentioned, that the jdbc batch size should be reset after an exceptional or successful execution. Unfortunately, this is not possible at the moment because of a wrong signature of the org.hibernate.SharedSessionContract#setJdbcBatchSize method. However I’ve created a ticket for the modification and will post the modified code.

Edit 2: Until the fix is released, I came up with a Reflection based solution, which is available on GitHub. The modified MethodInterceptor is the following:

public class JdbcBatchSizeMethodInterceptor implements MethodInterceptor {
    private static final Logger logger = LoggerFactory.getLogger(JdbcBatchSizeMethodInterceptor.class);

    @PersistenceContext
    private EntityManager entityManager;

    @Override
    public Object invoke(final MethodInvocation invocation) throws Throwable {
        final Session session = entityManager.unwrap(Session.class);
        final Integer originalBatchSize = session.getJdbcBatchSize();
        session.setJdbcBatchSize(invocation.getMethod().getAnnotation(JdbcBatchSize.class).value());
        Object result = null;
        try {
            result = invocation.proceed();
            entityManager.flush();
        } finally {
            try {
                final Field field = AbstractSharedSessionContract.class.getDeclaredField("jdbcBatchSize");
                field.setAccessible(true);
                ReflectionUtils.setField(field, session, originalBatchSize);
            } catch (final Exception e) {
                logger.error("Error when setting original jdbc batch size for the session.", e);
            }
        }
        return result;
    }
}

Advisor:

public class JdbcBatchSizeAdvisor extends AbstractPointcutAdvisor {
    private static final long serialVersionUID = -8466931120118556794L;

    private final StaticMethodMatcherPointcut pointcut = new StaticMethodMatcherPointcutAdvisor() {
        private static final long serialVersionUID = 4458743456797695036L;

        @Override
        public boolean matches(final Method method, final Class<?> targetClass) {
            return method.isAnnotationPresent(JdbcBatchSize.class) && method.isAnnotationPresent(Transactional.class);
        }
    };

    @Autowired
    private JdbcBatchSizeMethodInterceptor interceptor;

    @Override
    public Pointcut getPointcut() {
        return pointcut;
    }

    @Override
    public Advice getAdvice() {
        return interceptor;
    }

    @Override
    public int getOrder() {
        return Ordered.LOWEST_PRECEDENCE;
    }
}

The main point of this class is the annotation check for the method calls and the ordering of the aspects. If you think about it, it will be clear. To acquire the current session, we should be in a transactional context. The @Transactional annotation works with the Spring AOP framework also which is using proxies. This means that if our aspect has the priority over the transactional aspect, we won’t be in transactional context, thus we can’t get the current session so we have to somehow make sure that the ordering is correct. Setting the ordering for the transactional aspect can be done with the order attribute in the transaction-manager tag.

<tx:annotation-driven transaction-manager="transactionManager" order="0" proxy-target-class="true" />

Lower order value means it has more priority over the bigger one. You can check out the javadoc of the Ordered interface. This ordering is also defined for the JdbcBatchSizeAdvisor with the lowest precedence to make sure that the transactional context is set up.

After this, the new @JdbcBatchSize annotation can be used together with the @Transactional annotation to set up the batch size for the current session.

The implementation is available on my GitHub page.

Edit 3: With Hibernate 5.2.1 and Spring 4.3.1 available, the Reflection part for setting the JDBC Batch size back can be removed as well as the CustomJpaVendorAdapter. The updated code is available on GitHub.