Add teardown hook on WorkflowContext.

kotlin/workflow-core/src/main/java/com/squareup/workflow/Workflow.kt

@@ -51,6 +51,14 @@ package com.squareup.workflow
  * or renderings, extend [StatelessWorkflow], or just pass a lambda to the [stateless] function
  * below.
  *
+ * ## Interacting with Events and Other Workflows
+ *
+ * All workflows are passed a [WorkflowContext] in their compose methods. This context allows the
+ * workflow to interact with the outside world by doing things like listening for events,
+ * subscribing to streams of data, rendering child workflows, and performing cleanup when the
+ * workflow is about to be torn down by its parent. See the documentation on [WorkflowContext] for
+ * more information about what it can do.
+ *
  * @param InputT Typically a data class that is used to pass configuration information or bits of
  * state that the workflow can always get from its parent and needn't duplicate in its own state.
  * May be [Unit] if the workflow does not need any input data.

kotlin/workflow-core/src/main/java/com/squareup/workflow/WorkflowContext.kt

@@ -40,6 +40,10 @@ import kotlin.reflect.KType
  * ## Composing Children
  *
  * See [compose].
+ *
+ * ## Handling Workflow Teardown
+ *
+ * See [onTeardown].
  */
 interface WorkflowContext<StateT : Any, in OutputT : Any> {
 
@@ -123,6 +127,18 @@ interface WorkflowContext<StateT : Any, in OutputT : Any> {
     key: String = "",
     handler: (ChildOutputT) -> WorkflowAction<StateT, OutputT>
   ): ChildRenderingT
+
+  /**
+   * Adds an action to be invoked if the workflow is discarded by its parent before the next
+   * compose pass. Multiple hooks can be registered in the same compose pass, they will be invoked
+   * in the same order they're set. Like any other work performed through the context (e.g. calls
+   * to [compose] or [onReceive]), hooks are cleared at the start of each compose pass, so you must
+   * set any hooks you need in each compose pass.
+   *
+   * Teardown handlers should be non-blocking and execute quickly, since they are invoked
+   * synchronously during the compose pass.
+   */
+  fun onTeardown(handler: () -> Unit)
 }
 
 /**

kotlin/workflow-host/src/main/java/com/squareup/workflow/internal/Behavior.kt

@@ -33,7 +33,8 @@ import kotlin.reflect.KType
 internal data class Behavior<StateT : Any, out OutputT : Any>(
   val childCases: List<WorkflowOutputCase<*, *, StateT, OutputT>>,
   val subscriptionCases: List<SubscriptionCase<*, StateT, OutputT>>,
-  val nextActionFromEvent: Deferred<WorkflowAction<StateT, OutputT>>
+  val nextActionFromEvent: Deferred<WorkflowAction<StateT, OutputT>>,
+  val teardownHooks: List<() -> Unit>
 ) {
 
   // @formatter:off

kotlin/workflow-host/src/main/java/com/squareup/workflow/internal/RealWorkflowContext.kt

@@ -46,6 +46,7 @@ internal class RealWorkflowContext<StateT : Any, OutputT : Any>(
   private val nextUpdateFromEvent = CompletableDeferred<WorkflowAction<StateT, OutputT>>()
   private val subscriptionCases = mutableListOf<SubscriptionCase<*, StateT, OutputT>>()
   private val childCases = mutableListOf<WorkflowOutputCase<*, *, StateT, OutputT>>()
+  private val teardownHooks = mutableListOf<() -> Unit>()
 
   override fun <EventT : Any> onEvent(handler: (EventT) -> WorkflowAction<StateT, OutputT>):
       EventHandler<EventT> {
@@ -88,12 +89,17 @@ internal class RealWorkflowContext<StateT : Any, OutputT : Any>(
     return composer.compose(case, child, id, input)
   }
 
+  override fun onTeardown(handler: () -> Unit) {
+    teardownHooks += handler
+  }
+
   /**
    * Constructs an immutable [Behavior] from the context.
    */
   fun buildBehavior(): Behavior<StateT, OutputT> = Behavior(
       childCases = childCases.toList(),
       subscriptionCases = subscriptionCases.toList(),
-      nextActionFromEvent = nextUpdateFromEvent
+      nextActionFromEvent = nextUpdateFromEvent,
+      teardownHooks = teardownHooks.toList()
   )
 }

kotlin/workflow-host/src/main/java/com/squareup/workflow/internal/WorkflowNode.kt

@@ -71,6 +71,15 @@ internal class WorkflowNode<InputT : Any, StateT : Any, OutputT : Any, Rendering
    */
   override val coroutineContext = baseContext + Job(baseContext[Job]) + CoroutineName(id.toString())
 
+  init {
+    // This will get invoked whenever we are cancelled (via our cancel() method), or whenever
+    // any of our ancestor workflows are cancelled, anywhere up the tree, including whenever an
+    // exception is thrown that cancels a workflow.
+    coroutineContext[Job]!!.invokeOnCompletion {
+      behavior?.teardownHooks?.forEach { it.invoke() }
+    }
+  }
+
   private val subtreeManager = SubtreeManager<StateT, OutputT>(coroutineContext)
   private val subscriptionTracker =
     LifetimeTracker<SubscriptionCase<*, StateT, OutputT>, Any, Subscription>(
@@ -85,7 +94,7 @@ internal class WorkflowNode<InputT : Any, StateT : Any, OutputT : Any, Rendering
 
   private var lastInput: InputT = initialInput
 
-  private lateinit var behavior: Behavior<StateT, OutputT>
+  private var behavior: Behavior<StateT, OutputT>? = null
 
   /**
    * Walk the tree of workflows, rendering each one and using
@@ -149,7 +158,7 @@ internal class WorkflowNode<InputT : Any, StateT : Any, OutputT : Any, Rendering
 
     // Listen for any events.
     with(selector) {
-      behavior.nextActionFromEvent.onAwait { update ->
+      behavior!!.nextActionFromEvent.onAwait { update ->
         acceptUpdate(update)
       }
     }
@@ -162,6 +171,10 @@ internal class WorkflowNode<InputT : Any, StateT : Any, OutputT : Any, Rendering
    * after calling this method.
    */
   fun cancel() {
+    // No other cleanup work should be done in this function, since it will only be invoked when
+    // this workflow is *directly* discarded by its parent (or the host).
+    // If you need to do something whenever this workflow is torn down, add it to the
+    // invokeOnCompletion handler for the Job above.
     coroutineContext.cancel()
   }
 
@@ -179,9 +192,11 @@ internal class WorkflowNode<InputT : Any, StateT : Any, OutputT : Any, Rendering
     val rendering = workflow.compose(input, state, context)
 
     behavior = context.buildBehavior()
-    // Start new children/subscriptions, and drop old ones.
-    subtreeManager.track(behavior.childCases)
-    subscriptionTracker.track(behavior.subscriptionCases)
+        .apply {
+          // Start new children/subscriptions, and drop old ones.
+          subtreeManager.track(childCases)
+          subscriptionTracker.track(subscriptionCases)
+        }
 
     return rendering
   }

kotlin/workflow-host/src/test/java/com/squareup/workflow/internal/WorkflowNodeTest.kt

@@ -18,6 +18,7 @@ package com.squareup.workflow.internal
 import com.squareup.workflow.EventHandler
 import com.squareup.workflow.Snapshot
 import com.squareup.workflow.StatefulWorkflow
+import com.squareup.workflow.Workflow
 import com.squareup.workflow.WorkflowAction.Companion.emitOutput
 import com.squareup.workflow.WorkflowAction.Companion.enterState
 import com.squareup.workflow.WorkflowContext
@@ -26,6 +27,7 @@ import com.squareup.workflow.invoke
 import com.squareup.workflow.onReceive
 import com.squareup.workflow.parse
 import com.squareup.workflow.readUtf8WithLength
+import com.squareup.workflow.stateless
 import com.squareup.workflow.util.ChannelUpdate
 import com.squareup.workflow.util.ChannelUpdate.Closed
 import com.squareup.workflow.util.ChannelUpdate.Value
@@ -604,4 +606,20 @@ class WorkflowNodeTest {
     )
     assertEquals("input:new input|state:initial input", restoredNode.compose(workflow, "foo"))
   }
+
+  @Test fun `invokes teardown hooks in order on cancel`() {
+    val teardowns = mutableListOf<Int>()
+    val workflow = Workflow.stateless<Nothing, Unit> { context ->
+      context.onTeardown { teardowns += 1 }
+      context.onTeardown { teardowns += 2 }
+    }
+    val node = WorkflowNode(workflow.id(), workflow.asStatefulWorkflow(), Unit, null, Unconfined)
+    node.compose(workflow.asStatefulWorkflow(), Unit)
+
+    assertTrue(teardowns.isEmpty())
+
+    node.cancel()
+
+    assertEquals(listOf(1, 2), teardowns)
+  }
 }

kotlin/workflow-testing/src/test/java/com/squareup/workflow/CompositionIntegrationTest.kt

@@ -15,10 +15,12 @@
  */
 package com.squareup.workflow
 
+import com.squareup.workflow.WorkflowAction.Companion.enterState
 import com.squareup.workflow.testing.testFromStart
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
+import kotlin.test.assertTrue
 
 class CompositionIntegrationTest {
 
@@ -91,4 +93,88 @@ class CompositionIntegrationTest {
       }
     }
   }
+
+  @Test fun `all childrens teardown hooks invoked when parent discards it`() {
+    val teardowns = mutableListOf<String>()
+    val child1 = Workflow.stateless<Nothing, Unit> { context ->
+      context.onTeardown { teardowns += "child1" }
+    }
+    val child2 = Workflow.stateless<Nothing, Unit> { context ->
+      context.onTeardown { teardowns += "child2" }
+    }
+    // A workflow that will render child1 and child2 until its rendering is invoked, at which point
+    // it will compose neither of them, which should trigger the teardown callbacks.
+    val root = object : StatefulWorkflow<Unit, Boolean, Nothing, () -> Unit>() {
+      override fun initialState(
+        input: Unit,
+        snapshot: Snapshot?
+      ): Boolean = true
+
+      override fun compose(
+        input: Unit,
+        state: Boolean,
+        context: WorkflowContext<Boolean, Nothing>
+      ): () -> Unit {
+        if (state) {
+          context.compose(child1, key = "child1")
+          context.compose(child2, key = "child2")
+        }
+        return context.onEvent<Unit> { enterState(false) }::invoke
+      }
+
+      override fun snapshotState(state: Boolean): Snapshot = Snapshot.EMPTY
+    }
+
+    root.testFromStart { tester ->
+      tester.withNextRendering { teardownChildren ->
+        assertTrue(teardowns.isEmpty())
+
+        teardownChildren()
+
+        assertEquals(listOf("child1", "child2"), teardowns)
+      }
+    }
+  }
+
+  @Test fun `nested childrens teardown hooks invoked when parent discards it`() {
+    val teardowns = mutableListOf<String>()
+    val grandchild = Workflow.stateless<Nothing, Unit> { context ->
+      context.onTeardown { teardowns += "grandchild" }
+    }
+    val child = Workflow.stateless<Nothing, Unit> { context ->
+      context.compose(grandchild)
+      context.onTeardown { teardowns += "child" }
+    }
+    // A workflow that will render child1 and child2 until its rendering is invoked, at which point
+    // it will compose neither of them, which should trigger the teardown callbacks.
+    val root = object : StatefulWorkflow<Unit, Boolean, Nothing, () -> Unit>() {
+      override fun initialState(
+        input: Unit,
+        snapshot: Snapshot?
+      ): Boolean = true
+
+      override fun compose(
+        input: Unit,
+        state: Boolean,
+        context: WorkflowContext<Boolean, Nothing>
+      ): () -> Unit {
+        if (state) {
+          context.compose(child)
+        }
+        return context.onEvent<Unit> { enterState(false) }::invoke
+      }
+
+      override fun snapshotState(state: Boolean): Snapshot = Snapshot.EMPTY
+    }
+
+    root.testFromStart { tester ->
+      tester.withNextRendering { teardownChildren ->
+        assertTrue(teardowns.isEmpty())
+
+        teardownChildren()
+
+        assertEquals(listOf("grandchild", "child"), teardowns)
+      }
+    }
+  }
 }