Better crash reporter error grouping from RenderContext assertions

Normally crash reporters group errors by hashing against the elements of the top stack frame. This causes problems for bottleneck code that is enforcing invariants that should be implemented by client code: we wind up with "kitchen sink" groups that hold every crash of a particular type, rather than distinct groups for each specific way that mistake is made.

We fix this for a few spots in `RenderContext` by inserting a fake top stack frame based on some hashable key provided by client code -- the child workflow, the "key" param for runningSideEffect, etc.

So far this is a JVM-specific hack.

Author: rjrjr

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/RealRenderContext.kt

@@ -72,15 +72,15 @@ internal class RealRenderContext<out PropsT, StateT, OutputT>(
     key: String,
     handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>
   ): ChildRenderingT {
-    checkNotFrozen { "renderChild(${child.identifier})" }
+    checkNotFrozen(child) { "renderChild(${child.identifier})" }
     return renderer.render(child, props, key, handler)
   }
 
   override fun runningSideEffect(
     key: String,
     sideEffect: suspend CoroutineScope.() -> Unit
   ) {
-    checkNotFrozen { "runningSideEffect($key)" }
+    checkNotFrozen(key) { "runningSideEffect($key)" }
     sideEffectRunner.runningSideEffect(key, sideEffect)
   }
 
@@ -90,15 +90,15 @@ internal class RealRenderContext<out PropsT, StateT, OutputT>(
     vararg inputs: Any?,
     calculation: () -> ResultT
   ): ResultT {
-    checkNotFrozen { "remember($key)" }
+    checkNotFrozen(key) { "remember($key)" }
     return rememberStore.remember(key, resultType, inputs = inputs, calculation)
   }
 
   /**
    * Freezes this context so that any further calls to this context will throw.
    */
   fun freeze() {
-    checkNotFrozen { "freeze" }
+    checkNotFrozen("freeze") { "freeze" }
     frozen = true
   }
 
@@ -109,7 +109,13 @@ internal class RealRenderContext<out PropsT, StateT, OutputT>(
     frozen = false
   }
 
-  private fun checkNotFrozen(reason: () -> String) = check(!frozen) {
-    "RenderContext cannot be used after render method returns: ${reason()}"
-  }
+  /**
+   * @param stackTraceKey ensures unique crash reporter error groups.
+   *
+   * @see checkWithKey
+   */
+  private inline fun checkNotFrozen(stackTraceKey: Any, lazyMessage: () -> Any) =
+    checkWithKey(!frozen, stackTraceKey) {
+      "RenderContext cannot be used after render method returns: ${lazyMessage()}"
+    }
 }

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/SubtreeManager.kt

@@ -126,7 +126,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     // Prevent duplicate workflows with the same key.
     workflowTracer.trace("CheckingUniqueMatches") {
       children.forEachStaging {
-        require(!(it.matches(child, key, workflowTracer))) {
+        requireWithKey(!(it.matches(child, key, workflowTracer)), stackTraceKey = child) {
           "Expected keys to be unique for ${child.identifier}: key=\"$key\""
         }
       }

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/Throwables.kt

@@ -1,8 +1,66 @@
 package com.squareup.workflow1.internal
 
-import kotlinx.coroutines.CancellationException
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 
-internal tailrec fun Throwable.unwrapCancellationCause(): Throwable? {
-  if (this !is CancellationException) return this
-  return cause?.unwrapCancellationCause()
+/**
+ * Like Kotlin's [require], but uses [stackTraceKey] to create a fake top element
+ * on the stack trace, ensuring that crash reporter's default grouping will create unique
+ * groups for unique keys.
+ *
+ * So far [stackTraceKey] is only effective on JVM, it has no effect in other languages.
+ *
+ * @see [withKey]
+ *
+ * @throws IllegalArgumentException if the [value] is false.
+ */
+@OptIn(ExperimentalContracts::class)
+internal inline fun requireWithKey(
+  value: Boolean,
+  stackTraceKey: Any,
+  lazyMessage: () -> Any = { "Failed requirement." }
+) {
+  contract {
+    returns() implies value
+  }
+  if (!value) {
+    val message = lazyMessage()
+    val exception: Throwable = IllegalArgumentException(message.toString())
+    throw exception.withKey(stackTraceKey)
+  }
 }
+
+/**
+ * Like Kotlin's [check], but uses [stackTraceKey] to create a fake top element
+ * on the stack trace, ensuring that crash reporter's default grouping will create unique
+ * groups for unique keys.
+ *
+ * So far [stackTraceKey] is only effective on JVM, it has no effect in other languages.
+ *
+ * @see [withKey]
+ *
+ * @throws IllegalStateException if the [value] is false.
+ */
+@OptIn(ExperimentalContracts::class)
+internal inline fun checkWithKey(
+  value: Boolean,
+  stackTraceKey: Any,
+  lazyMessage: () -> Any = { "Check failed." }
+) {
+  contract {
+    returns() implies value
+  }
+  if (!value) {
+    val message = lazyMessage()
+    val exception: Throwable = IllegalStateException(message.toString())
+    throw exception.withKey(stackTraceKey)
+  }
+}
+
+/**
+ * Uses [stackTraceKey] to create a fake top element on the stack trace, ensuring
+ * that crash reporter's default grouping will create unique groups for unique keys.
+ *
+ * So far only effective on JVM, this is a pass through in other languages.
+ */
+internal expect fun <T : Throwable> T.withKey(stackTraceKey: Any): T

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt

@@ -163,7 +163,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   ) {
     // Prevent duplicate side effects with the same key.
     sideEffects.forEachStaging {
-      require(key != it.key) { "Expected side effect keys to be unique: \"$key\"" }
+      requireWithKey(key != it.key, key) { "Expected side effect keys to be unique: \"$key\"" }
     }
 
     sideEffects.retainOrCreate(
@@ -179,7 +179,10 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     calculation: () -> ResultT
   ): ResultT {
     remembered.forEachStaging {
-      require(key != it.key || resultType != it.resultType || !inputs.contentEquals(it.inputs)) {
+      requireWithKey(
+        key != it.key || resultType != it.resultType || !inputs.contentEquals(it.inputs),
+        stackTraceKey = key
+      ) {
         "Expected combination of key, inputs and result type to be unique: \"$key\""
       }
     }

workflow-runtime/src/iosMain/kotlin/com/squareup/workflow1/internal/Throwables.ios.kt

@@ -0,0 +1,3 @@
+package com.squareup.workflow1.internal
+
+actual fun <T : Throwable> T.withKey(stackTraceKey: Any): T = this

workflow-runtime/src/jsMain/kotlin/com/squareup/workflow1/internal/Throwables.js.kt

@@ -0,0 +1,3 @@
+package com.squareup.workflow1.internal
+
+actual fun <T : Throwable> T.withKey(stackTraceKey: Any): T = this

workflow-runtime/src/jvmMain/kotlin/com/squareup/workflow1/internal/Throwables.jvm.kt

@@ -0,0 +1,13 @@
+package com.squareup.workflow1.internal
+
+internal actual fun <T : Throwable> T.withKey(stackTraceKey: Any): T = apply {
+  val realTop = stackTrace[0]
+  val fakeTop = StackTraceElement(
+    // Real class name to ensure that we are still "in project".
+    realTop.className,
+    "fakeMethodForCrashGrouping",
+    /* fileName = */ stackTraceKey.toString(),
+    /* lineNumber = */ stackTraceKey.hashCode()
+  )
+  stackTrace = stackTrace.toMutableList().apply { add(0, fakeTop) }.toTypedArray()
+}

workflow-runtime/src/jvmTest/kotlin/com/squareup/workflow1/internal/ThrowablesTest.kt

@@ -0,0 +1,39 @@
+package com.squareup.workflow1.internal
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class ThrowablesTest {
+
+  @Test fun `requireWithKey throws IllegalArgumentException`() {
+    try {
+      requireWithKey(false, "key") { "message" }
+    } catch (e: IllegalArgumentException) {
+      assertEquals("message", e.message)
+      e.assertIsKeyedException()
+      return
+    }
+  }
+
+  @Test fun `checkWithKey throws IllegalStateException`() {
+    try {
+      checkWithKey(false, "key") { "message" }
+    } catch (e: IllegalStateException) {
+      assertEquals("message", e.message)
+      e.assertIsKeyedException()
+      return
+    }
+  }
+
+  @Test fun `Throwable withKey adds frame based on key`() {
+    RuntimeException("cause").withKey("key").assertIsKeyedException()
+  }
+
+  private fun RuntimeException.assertIsKeyedException() {
+    val top = stackTrace[0]
+    val topPlusOne = stackTrace[1]
+    assertEquals(topPlusOne.className, top.className, "Uses real class name")
+    assertEquals("key", top.fileName)
+    assertEquals("key".hashCode(), top.lineNumber)
+  }
+}