Add support for FlushEntity listener (cashapp#2220)

* Make how we add listeners configurable

 * This introduces three policies for adding a listener: prepend, append,
   and replace.

* Add FlushEntity listener

 * A FlushEntity listener also needs CallbackRegistryConsumer support,
   so we support both.

* Update misk-hibernate.api

misk-hibernate/api/misk-hibernate.api

@@ -1,3 +1,11 @@
+public final class misk/hibernate/BindPolicy : java/lang/Enum {
+	public static final field APPEND Lmisk/hibernate/BindPolicy;
+	public static final field PREPEND Lmisk/hibernate/BindPolicy;
+	public static final field REPLACE Lmisk/hibernate/BindPolicy;
+	public static fun valueOf (Ljava/lang/String;)Lmisk/hibernate/BindPolicy;
+	public static fun values ()[Lmisk/hibernate/BindPolicy;
+}
+
 public abstract interface annotation class misk/hibernate/Constraint : java/lang/annotation/Annotation {
 	public abstract fun operator ()Lmisk/hibernate/Operator;
 	public abstract fun path ()Ljava/lang/String;
@@ -72,7 +80,8 @@ public abstract class misk/hibernate/HibernateEntityModule : misk/inject/KAbstra
 	protected final fun addEntities ([Lkotlin/reflect/KClass;)V
 	protected final fun addEntity (Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;)V
 	public static synthetic fun addEntity$default (Lmisk/hibernate/HibernateEntityModule;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;Lkotlin/reflect/KClass;ILjava/lang/Object;)V
-	protected final fun bindListener (Lorg/hibernate/event/spi/EventType;)Lcom/google/inject/binder/LinkedBindingBuilder;
+	protected final fun bindListener (Lorg/hibernate/event/spi/EventType;Lmisk/hibernate/BindPolicy;)Lcom/google/inject/binder/LinkedBindingBuilder;
+	public static synthetic fun bindListener$default (Lmisk/hibernate/HibernateEntityModule;Lorg/hibernate/event/spi/EventType;Lmisk/hibernate/BindPolicy;ILjava/lang/Object;)Lcom/google/inject/binder/LinkedBindingBuilder;
 	protected fun configure ()V
 	public abstract fun configureHibernate ()V
 	protected final fun installHibernateAdminDashboardWebActions ()V

misk-hibernate/src/main/kotlin/misk/hibernate/AggregateListener.kt

@@ -3,6 +3,8 @@ package misk.hibernate
 import com.google.common.collect.LinkedHashMultimap
 import org.hibernate.event.service.spi.EventListenerRegistry
 import org.hibernate.event.spi.EventType
+import org.hibernate.event.spi.FlushEntityEvent
+import org.hibernate.event.spi.FlushEntityEventListener
 import org.hibernate.event.spi.PostDeleteEvent
 import org.hibernate.event.spi.PostDeleteEventListener
 import org.hibernate.event.spi.PostInsertEvent
@@ -21,6 +23,9 @@ import org.hibernate.event.spi.PreUpdateEvent
 import org.hibernate.event.spi.PreUpdateEventListener
 import org.hibernate.event.spi.SaveOrUpdateEvent
 import org.hibernate.event.spi.SaveOrUpdateEventListener
+import org.hibernate.jpa.event.spi.Callback
+import org.hibernate.jpa.event.spi.CallbackRegistry
+import org.hibernate.jpa.event.spi.CallbackRegistryConsumer
 import org.hibernate.persister.entity.EntityPersister
 import javax.inject.Provider
 
@@ -38,19 +43,29 @@ internal class AggregateListener(
   PostUpdateEventListener,
   PreInsertEventListener,
   PostInsertEventListener,
-  SaveOrUpdateEventListener {
+  SaveOrUpdateEventListener,
+  FlushEntityEventListener,
+  CallbackRegistryConsumer {
   private val multimap = LinkedHashMultimap.create<EventType<*>, Provider<*>>()!!
+  private val listenerAddPolicy = LinkedHashMap<EventType<*>, BindPolicy>()
 
   init {
     for (eventTypeAndListener in registrations) {
       multimap.put(eventTypeAndListener.type, eventTypeAndListener.provider)
+      listenerAddPolicy[eventTypeAndListener.type] = eventTypeAndListener.policy
     }
   }
 
   fun registerAll(eventListenerRegistry: EventListenerRegistry) {
     for (eventType in multimap.keySet()) {
       @Suppress("UNCHECKED_CAST") // We don't have static type information for the event type.
-      eventListenerRegistry.appendListeners(eventType as EventType<Any>, this)
+      eventType as EventType<Any>
+
+      when (listenerAddPolicy[eventType]) {
+        BindPolicy.PREPEND -> eventListenerRegistry.prependListeners(eventType, this)
+        BindPolicy.REPLACE -> eventListenerRegistry.setListeners(eventType, this)
+        BindPolicy.APPEND -> eventListenerRegistry.appendListeners(eventType, this)
+      }
     }
   }
 
@@ -130,4 +145,18 @@ internal class AggregateListener(
       (provider.get() as SaveOrUpdateEventListener).onSaveOrUpdate(event)
     }
   }
+
+  override fun onFlushEntity(event: FlushEntityEvent?) {
+    for (provider in multimap[EventType.FLUSH_ENTITY]) {
+      (provider.get() as FlushEntityEventListener).onFlushEntity(event)
+    }
+  }
+
+  override fun injectCallbackRegistry(callbackRegistry: CallbackRegistry?) {
+    for (provider in multimap.values()) {
+      val handler = provider.get() as? CallbackRegistryConsumer ?: continue
+      handler.injectCallbackRegistry(callbackRegistry)
+    }
+  }
+
 }

misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEntityModule.kt

@@ -98,7 +98,10 @@ abstract class HibernateEntityModule(
     addEntity(T::class, Q::class, AA::class)
   }
 
-  protected fun <T> bindListener(type: EventType<T>): LinkedBindingBuilder<in T> {
+  protected fun <T> bindListener(
+    type: EventType<T>,
+    policy: BindPolicy = BindPolicy.APPEND
+  ): LinkedBindingBuilder<in T> {
     // Bind the listener as an anonymous key. We can get the provider for this before its bound!
     val key = Key.get(
       Any::class.java,
@@ -107,7 +110,7 @@ abstract class HibernateEntityModule(
 
     // Create a multibinding for a ListenerRegistration that uses the above key.
     multibind<ListenerRegistration>(qualifier)
-      .toInstance(ListenerRegistration(type, getProvider(key)))
+      .toInstance(ListenerRegistration(type, getProvider(key), policy))
 
     // Start the binding.
     return bind(key)

misk-hibernate/src/main/kotlin/misk/hibernate/HibernateEventListener.kt

@@ -3,14 +3,24 @@ package misk.hibernate
 import org.hibernate.event.spi.EventType
 import javax.inject.Provider
 
+/**
+ * Control how we register listeners.
+ */
+enum class BindPolicy {
+  PREPEND,
+  REPLACE,
+  APPEND
+}
+
 /**
  * A registration of a listener for one of many Hibernate event types. This class uses providers to
  * get a new listener instance each time it is needed. This is intended to prevent circular
  * dependencies.
  */
 internal class ListenerRegistration(
   val type: EventType<*>,
-  val provider: Provider<*>
+  val provider: Provider<*>,
+  val policy: BindPolicy
 ) {
   init {
     when (type) {
@@ -23,6 +33,7 @@ internal class ListenerRegistration(
       EventType.PRE_DELETE -> Unit
       EventType.POST_DELETE -> Unit
       EventType.SAVE_UPDATE -> Unit
+      EventType.FLUSH_ENTITY -> Unit
       else -> throw UnsupportedOperationException("$type not currently supported")
     }
   }