The batching operation implementation should be re-factored to take advantage of the generic event context manager API (see org.hibernate.ogm.dialect.eventstate.impl). This will allow to get rid of some custom code (event listeners, a ThreadLocal etc).