From cdf88902eb5cca82cd2dbf048177efa6b8c2b505 Mon Sep 17 00:00:00 2001
From: Stephen Edwards <sedwards@squareup.com>
Date: Tue, 10 Jun 2025 11:12:52 -0400
Subject: [PATCH] 1338: Clean up caches after WorkflowNode torn down

---
 workflow-core/api/workflow-core.api           | 11 +++++++++++
 .../squareup/workflow1/StatelessWorkflow.kt   | 19 ++++++++++++++++++-
 .../workflow1/internal/WorkflowNode.kt        | 10 ++++++----
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/workflow-core/api/workflow-core.api b/workflow-core/api/workflow-core.api
index 0230bc9f18..d82fd02f93 100644
--- a/workflow-core/api/workflow-core.api
+++ b/workflow-core/api/workflow-core.api
@@ -284,6 +284,17 @@ public final class com/squareup/workflow1/StatelessWorkflow$RenderContext : com/
 	public fun runningSideEffect (Ljava/lang/String;Lkotlin/jvm/functions/Function2;)V
 }
 
+public final class com/squareup/workflow1/StatelessWorkflow$StatelessAsStatefulWorkflow : com/squareup/workflow1/StatefulWorkflow {
+	public fun <init> (Lcom/squareup/workflow1/StatelessWorkflow;)V
+	public final fun clearCache ()V
+	public synthetic fun initialState (Ljava/lang/Object;Lcom/squareup/workflow1/Snapshot;)Ljava/lang/Object;
+	public fun initialState (Ljava/lang/Object;Lcom/squareup/workflow1/Snapshot;)V
+	public synthetic fun render (Ljava/lang/Object;Ljava/lang/Object;Lcom/squareup/workflow1/StatefulWorkflow$RenderContext;)Ljava/lang/Object;
+	public fun render (Ljava/lang/Object;Lkotlin/Unit;Lcom/squareup/workflow1/StatefulWorkflow$RenderContext;)Ljava/lang/Object;
+	public synthetic fun snapshotState (Ljava/lang/Object;)Lcom/squareup/workflow1/Snapshot;
+	public fun snapshotState (Lkotlin/Unit;)Lcom/squareup/workflow1/Snapshot;
+}
+
 public final class com/squareup/workflow1/TypedWorker : com/squareup/workflow1/Worker {
 	public fun <init> (Lkotlin/reflect/KType;Lkotlinx/coroutines/flow/Flow;)V
 	public fun doesSameWorkAs (Lcom/squareup/workflow1/Worker;)Z
diff --git a/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/StatelessWorkflow.kt b/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/StatelessWorkflow.kt
index c13b47a9c4..c9945b2af4 100644
--- a/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/StatelessWorkflow.kt
+++ b/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/StatelessWorkflow.kt
@@ -227,7 +227,7 @@ public abstract class StatelessWorkflow<PropsT, OutputT, out RenderingT> :
    * Class type returned by [asStatefulWorkflow].
    * See [statefulWorkflow] for the instance.
    */
-  private inner class StatelessAsStatefulWorkflow :
+  inner class StatelessAsStatefulWorkflow :
     StatefulWorkflow<PropsT, Unit, OutputT, RenderingT>() {
 
     /**
@@ -268,6 +268,23 @@ public abstract class StatelessWorkflow<PropsT, OutputT, out RenderingT> :
     }
 
     override fun snapshotState(state: Unit): Snapshot? = null
+
+    /**
+     * When we are finished with at least one node that holds on to this workflow instance,
+     * then we clear the cache. The reason we do that every time is that it *might* be the last
+     * node that is caching this instance, and if so, we do not want to leak these cached
+     * render contexts.
+     *
+     * Yes, that means that it might have to be re-created again when this instance is used
+     * multiple times. The current design for how we get a [StatefulWorkflow] from the
+     * [StatelessWorkflow] is a failed compromise between performance (caching) and type-safe
+     * brevity (erasing the `StateT` type from the concerns of [StatelessWorkflow]). It needs
+     * to be fixed with a bigger re-write (https://github.com/square/workflow-kotlin/issues/1337).
+     */
+    fun clearCache() {
+      cachedStatelessRenderContext = null
+      canonicalStatefulRenderContext = null
+    }
   }
 
   private val statefulWorkflow: StatefulWorkflow<PropsT, Unit, OutputT, RenderingT> =
diff --git a/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt b/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt
index 17dc3ff7d9..1c09af8596 100644
--- a/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt
+++ b/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt
@@ -9,6 +9,7 @@ import com.squareup.workflow1.RuntimeConfig
 import com.squareup.workflow1.RuntimeConfigOptions
 import com.squareup.workflow1.RuntimeConfigOptions.PARTIAL_TREE_RENDERING
 import com.squareup.workflow1.StatefulWorkflow
+import com.squareup.workflow1.StatelessWorkflow
 import com.squareup.workflow1.TreeSnapshot
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
@@ -234,11 +235,12 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
    * after calling this method.
    */
   fun cancel(cause: CancellationException? = null) {
-    // 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(cause)
+    lastRendering = NullableInitBox()
+    (
+      cachedWorkflowInstance as?
+        StatelessWorkflow<PropsT, OutputT, RenderingT>.StatelessAsStatefulWorkflow
+      )?.clearCache()
   }
 
   /**