Add ActionScopeListener (#3668)

`ActionScopeListener` allows developers to hook in to the `close()` method of `ActionScope`, and perform actions immediately before the scope is closed. For example, you could buffer logs while the scope is open, and then dedupe and flush them before it's closed. Currently `ActionScopeListener` only supports `onClose()`, but in the future could support `onEnter()` as well.

Author: JGulbronson

misk-action-scopes/api/misk-action-scopes.api

@@ -42,6 +42,10 @@ public final class misk/scope/ActionScope$Instance : java/lang/AutoCloseable {
 	public final fun inScope (Lkotlin/jvm/functions/Function0;)Ljava/lang/Object;
 }
 
+public abstract interface class misk/scope/ActionScopeListener {
+	public abstract fun onClose ()V
+}
+
 public abstract class misk/scope/ActionScopedProviderModule : misk/inject/KAbstractModule {
 	public static final field Companion Lmisk/scope/ActionScopedProviderModule$Companion;
 	public fun <init> ()V
@@ -51,6 +55,7 @@ public abstract class misk/scope/ActionScopedProviderModule : misk/inject/KAbstr
 	public final fun bindConstant (Lkotlin/reflect/KClass;Ljava/lang/Object;Ljava/lang/annotation/Annotation;)V
 	public static synthetic fun bindConstant$default (Lmisk/scope/ActionScopedProviderModule;Lcom/google/inject/TypeLiteral;Ljava/lang/Object;Ljava/lang/annotation/Annotation;ILjava/lang/Object;)V
 	public static synthetic fun bindConstant$default (Lmisk/scope/ActionScopedProviderModule;Lkotlin/reflect/KClass;Ljava/lang/Object;Ljava/lang/annotation/Annotation;ILjava/lang/Object;)V
+	public final fun bindListener (Lcom/google/inject/TypeLiteral;)V
 	public final fun bindProvider (Lcom/google/inject/TypeLiteral;Lkotlin/reflect/KClass;)V
 	public final fun bindProvider (Lcom/google/inject/TypeLiteral;Lkotlin/reflect/KClass;Ljava/lang/Class;)V
 	public final fun bindProvider (Lcom/google/inject/TypeLiteral;Lkotlin/reflect/KClass;Ljava/lang/annotation/Annotation;)V

misk-action-scopes/src/main/kotlin/misk/scope/ActionScope.kt

@@ -23,7 +23,8 @@ internal constructor(
   // on ActionScopedProviders, which might depend on other ActionScopeds. We break
   // this circular dependency by injecting a map of Provider<ActionScopedProvider>
   // rather than the map of ActionScopedProvider directly
-  private val providers: @JvmSuppressWildcards Map<Key<*>, Provider<ActionScopedProvider<*>>>
+  private val providers: @JvmSuppressWildcards Map<Key<*>, Provider<ActionScopedProvider<*>>>,
+  private val listeners: @JvmSuppressWildcards Provider<Set<ActionScopeListener>>,
 ) : AutoCloseable {
   companion object {
     private val threadLocalInstance = ThreadLocal<Instance>()
@@ -123,7 +124,11 @@ internal constructor(
   }
 
   override fun close() {
-    threadLocalInstance.remove()
+    try {
+      listeners.get().forEach { it.onClose() }
+    } finally {
+      threadLocalInstance.remove()
+    }
 
     // Explicitly NOT removing threadLocalUUID because we want to retain the thread's UUID if
     // the action scope is re-entered on the same thread.

misk-action-scopes/src/main/kotlin/misk/scope/ActionScopeListener.kt

@@ -0,0 +1,23 @@
+package misk.scope
+
+/**
+ * A listener that can be called at various points within an [ActionScope]'s lifecycle. See the comment on each method
+ * to understand when in the lifecycle it is called.
+ */
+interface ActionScopeListener {
+  /**
+   * Called during [ActionScope.close], immediately before the [ThreadLocal] that holds the [ActionScope.Instance] is
+   * removed.
+   *
+   * The [ActionScope] is not closed until all the listeners are executed. That means that:
+   * - Immediately before [onClose] is called, [ActionScope.inScope] returns true
+   * - During [onClose], [ActionScope.inScope] returns true
+   * - Immediately after all listeners have called [onClose], [ActionScope.inScope] returns false
+   *
+   * The [ActionScope] being closed or not is independent of the lifecycle of any action scoped values. A reference to
+   * an object provided by an [ActionScopedProvider] can be held and used even if [ActionScope.inScope] returns false.
+   * For example: an action-scoped HTTP Request Body may still be transmitting to the client despite the scope being
+   * closed.
+   */
+  fun onClose()
+}

misk-action-scopes/src/main/kotlin/misk/scope/ActionScopedProviderModule.kt

@@ -21,6 +21,7 @@ abstract class ActionScopedProviderModule : KAbstractModule() {
   override fun configure() {
     MapBinder.newMapBinder(binder(), KEY_TYPE, ACTION_SCOPED_PROVIDER_TYPE)
     Multibinder.newSetBinder(binder(), KEY_TYPE)
+    Multibinder.newSetBinder(binder(), ACTION_SCOPE_LISTENER_TYPE)
     configureProviders()
   }
 
@@ -146,6 +147,14 @@ abstract class ActionScopedProviderModule : KAbstractModule() {
     bindProvider(typeKey, actionScopedKey, binder().getProvider(providerType.java))
   }
 
+  inline fun <reified T : ActionScopeListener> bindListener() {
+    bindListener(object : TypeLiteral<T>() {})
+  }
+
+  fun bindListener(typeLiteral: TypeLiteral<out ActionScopeListener>) {
+    Multibinder.newSetBinder(binder(), ACTION_SCOPE_LISTENER_TYPE).addBinding().to(typeLiteral)
+  }
+
   private fun <T> bindProvider(
     key: Key<T>,
     actionScopedKey: Key<ActionScoped<T>>,
@@ -190,6 +199,7 @@ abstract class ActionScopedProviderModule : KAbstractModule() {
 
     private val KEY_TYPE = object : TypeLiteral<Key<*>>() {}
     private val ACTION_SCOPED_PROVIDER_TYPE = object : TypeLiteral<ActionScopedProvider<*>>() {}
+    private val ACTION_SCOPE_LISTENER_TYPE = object : TypeLiteral<ActionScopeListener>() {}
 
     private class SeedDataActionScopedProvider<out T>(private val key: Key<T>) : ActionScopedProvider<T> {
       override fun get(): T {

misk-action-scopes/src/test/kotlin/misk/scope/ActionScopedTest.kt

@@ -13,6 +13,7 @@ import kotlinx.coroutines.runBlocking
 import misk.inject.keyOf
 import misk.inject.toKey
 import misk.inject.uninject
+import misk.scope.TestActionScopedProviderModule.TestListener
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
@@ -36,6 +37,8 @@ internal class ActionScopedTest {
 
   @Inject private lateinit var scope: ActionScope
 
+  @Inject private lateinit var testListener: TestListener
+
   @BeforeEach
   fun clearInjections() {
     uninject(this)
@@ -277,4 +280,19 @@ internal class ActionScopedTest {
       assertThat(countingString.get()).isEqualTo("Called CountingProvider 1 time(s)")
     }
   }
+
+  @Test
+  fun `listeners are called before the scope closes`() {
+    val injector = Guice.createInjector(TestActionScopedProviderModule())
+    injector.injectMembers(this)
+
+    assertThat(testListener.result).isNull()
+
+    scope.create(mapOf()).inScope {
+      // We don't have to do anything in the scope, just call inScope() so that close() is called. The listener is
+      // then triggered, setting the result field to an action scoped value, to show we're still in an action scope.
+    }
+
+    assertThat(testListener.result).isEqualTo("constant-value")
+  }
 }

misk-action-scopes/src/test/kotlin/misk/scope/TestActionScopedProviderModule.kt

@@ -4,6 +4,7 @@ import com.google.inject.TypeLiteral
 import com.google.inject.name.Named
 import com.google.inject.name.Names
 import jakarta.inject.Inject
+import jakarta.inject.Singleton
 import java.util.Optional
 
 internal class TestActionScopedProviderModule : ActionScopedProviderModule() {
@@ -19,6 +20,7 @@ internal class TestActionScopedProviderModule : ActionScopedProviderModule() {
     bindProvider(nullableStringTypeLiteral, NullableFooProvider::class, Names.named("nullable-foo"))
     bindProvider(nullableStringTypeLiteral, NullableBasedOnFooProvider::class, Names.named("nullable-based-on-foo"))
     bindProvider(String::class, CountingProvider::class, Names.named(("counting")))
+    bindListener<TestListener>()
   }
 
   class BarProvider @Inject internal constructor(@Named("from-seed") private val seedData: ActionScoped<String>) :
@@ -75,6 +77,16 @@ internal class TestActionScopedProviderModule : ActionScopedProviderModule() {
     override fun get(): String = "Called CountingProvider ${++callCount} time(s)"
   }
 
+  @Singleton
+  class TestListener @Inject constructor(@Named("constant") private val constant: ActionScoped<String>) :
+    ActionScopeListener {
+    var result: String? = null
+
+    override fun onClose() {
+      result = constant.get()
+    }
+  }
+
   companion object {
     val nullableStringTypeLiteral = object : TypeLiteral<String?>() {}
   }

misk-inject/src/main/kotlin/misk/inject/Guice.kt

@@ -110,9 +110,9 @@ inline fun <reified T : Any> keyOf(a: KClass<out Annotation>?): Key<T> = keyOf(a
  */
 inline fun <reified T : Any> keyOf(qualifier: BindingQualifier? = null): Key<T> =
   when (qualifier) {
-    is BindingQualifier.InstanceQualifier -> Key.get(T::class.java, qualifier.annotation)
-    is BindingQualifier.TypeClassifier -> Key.get(T::class.java, qualifier.type.java)
-    null -> Key.get(T::class.java)
+    is BindingQualifier.InstanceQualifier -> Key.get(object : TypeLiteral<T>() {}, qualifier.annotation)
+    is BindingQualifier.TypeClassifier -> Key.get(object : TypeLiteral<T>() {}, qualifier.type.java)
+    null -> Key.get(object : TypeLiteral<T>() {})
   }
 
 /**