Introduce unsnapshottableIdentifier() to allow transient WorkflowIdentifiers.

workflow-core/api/workflow-core.api

@@ -172,12 +172,12 @@ public final class com/squareup/workflow/WorkflowIdentifier {
 	public static final field Companion Lcom/squareup/workflow/WorkflowIdentifier$Companion;
 	public fun equals (Ljava/lang/Object;)Z
 	public fun hashCode ()I
+	public final fun toByteStringOrNull ()Lokio/ByteString;
 	public fun toString ()Ljava/lang/String;
-	public final fun write (Lokio/BufferedSink;)V
 }
 
 public final class com/squareup/workflow/WorkflowIdentifier$Companion {
-	public final fun read (Lokio/BufferedSource;)Lcom/squareup/workflow/WorkflowIdentifier;
+	public final fun parse (Lokio/ByteString;)Lcom/squareup/workflow/WorkflowIdentifier;
 }
 
 public final class com/squareup/workflow/WorkflowOutput {
@@ -225,6 +225,7 @@ public final class com/squareup/workflow/Workflows {
 	public static synthetic fun stateful$default (Lcom/squareup/workflow/Workflow$Companion;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function3;ILjava/lang/Object;)Lcom/squareup/workflow/StatefulWorkflow;
 	public static final fun stateless (Lcom/squareup/workflow/Workflow$Companion;Lkotlin/jvm/functions/Function2;)Lcom/squareup/workflow/Workflow;
 	public static final fun transform (Lcom/squareup/workflow/Worker;Lkotlin/jvm/functions/Function1;)Lcom/squareup/workflow/Worker;
+	public static final fun unsnapshottableIdentifier (Lkotlin/reflect/KType;)Lcom/squareup/workflow/WorkflowIdentifier;
 	public static final fun workflowAction (Lcom/squareup/workflow/StatefulWorkflow;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lcom/squareup/workflow/WorkflowAction;
 	public static final fun workflowAction (Lcom/squareup/workflow/StatefulWorkflow;Lkotlin/jvm/functions/Function0;Lkotlin/jvm/functions/Function1;)Lcom/squareup/workflow/WorkflowAction;
 	public static synthetic fun workflowAction$default (Lcom/squareup/workflow/StatefulWorkflow;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lcom/squareup/workflow/WorkflowAction;

workflow-core/src/main/java/com/squareup/workflow/WorkflowIdentifier.kt

@@ -18,11 +18,13 @@
 
 package com.squareup.workflow
 
-import okio.BufferedSink
-import okio.BufferedSource
+import okio.Buffer
+import okio.ByteString
 import okio.EOFException
 import kotlin.LazyThreadSafetyMode.PUBLICATION
+import kotlin.reflect.KAnnotatedElement
 import kotlin.reflect.KClass
+import kotlin.reflect.KType
 
 /**
  * Represents a [Workflow]'s "identity" and is used by the runtime to determine whether a workflow
@@ -31,61 +33,82 @@ import kotlin.reflect.KClass
  *
  * A workflow's identity consists primarily of its concrete type (i.e. the class that implements
  * the [Workflow] interface). Two workflows of the same concrete type are considered identical.
- *
  * However, if a workflow class implements [ImpostorWorkflow], the identifier will also include
  * that workflow's [ImpostorWorkflow.realIdentifier].
  *
+ * Instances of this class are [equatable][equals] and [hashable][hashCode].
+ *
+ * ## Identifiers and snapshots
+ *
+ * Since workflows can be [serialized][StatefulWorkflow.snapshotState], workflows' identifiers must
+ * also be serializable in order to match workflows back up with their snapshots when restoring.
+ * However, some [WorkflowIdentifier]s may represent workflows that cannot be snapshotted. When an
+ * identifier is not snapshottable, [toByteStringOrNull] will return null, and any identifiers that
+ * reference [ImpostorWorkflow]s whose [ImpostorWorkflow.realIdentifier] is not snapshottable will
+ * also not be snapshottable. Such identifiers are created with [unsnapshottableIdentifier], but
+ * should not be used to wrap arbitrary workflows since those workflows may expect to be
+ * snapshotted.
+ *
  * @constructor
- * @param type The [KClass] of the [Workflow] this identifier identifies.
+ * @param type The [KClass] of the [Workflow] this identifier identifies, or the [KType] of an
+ * [unsnapshottableIdentifier].
  * @param proxiedIdentifier An optional identifier from [ImpostorWorkflow.realIdentifier] that will
  * be used to further narrow the scope of this identifier.
  */
 @ExperimentalWorkflowApi
 class WorkflowIdentifier internal constructor(
-  private val type: KClass<out Workflow<*, *, *>>,
-  private val proxiedIdentifier: WorkflowIdentifier?
+  private val type: KAnnotatedElement,
+  private val proxiedIdentifier: WorkflowIdentifier? = null
 ) {
 
   /**
-   * The fully-qualified name of [type]. Computed lazily.
+   * The fully-qualified name of the type of workflow this identifier identifies. Computed lazily
+   * and cached.
    */
-  private val typeString: String by lazy(PUBLICATION) { type.java.name }
+  private val typeName: String by lazy(PUBLICATION) {
+    if (type is KClass<*>) type.java.name else type.toString()
+  }
 
   /**
-   * Returns a description of this identifier including the name of its workflow type and any
-   * [proxiedIdentifier]s. Computes [typeString] if it has not already been computed.
+   * If this identifier is snapshottable, returns the serialized form of the identifier.
+   * If it is not snapshottable, returns null.
    */
-  override fun toString(): String =
-    generateSequence(this) { it.proxiedIdentifier }
-        .joinToString { it.typeString }
-        .let { "WorkflowIdentifier($it)" }
+  fun toByteStringOrNull(): ByteString? {
+    if (type !is KClass<*>) return null
 
-  /**
-   * Serializes this identifier to the sink. It can be read back with [WorkflowIdentifier.read].
-   */
-  fun write(sink: BufferedSink) {
-    sink.writeUtf8WithLength(typeString)
-    if (proxiedIdentifier != null) {
-      sink.writeByte(PROXY_IDENTIFIER_TAG.toInt())
-      proxiedIdentifier.write(sink)
-    } else {
-      sink.writeByte(NO_PROXY_IDENTIFIER_TAG.toInt())
+    val proxiedBytes = proxiedIdentifier?.let {
+      // If we have a proxied identifier but it's not serializable, then we can't be serializable
+      // either.
+      it.toByteStringOrNull() ?: return null
+    }
+
+    return Buffer().let { sink ->
+      sink.writeUtf8WithLength(typeName)
+      if (proxiedBytes != null) {
+        sink.writeByte(PROXY_IDENTIFIER_TAG.toInt())
+        sink.write(proxiedBytes)
+      } else {
+        sink.writeByte(NO_PROXY_IDENTIFIER_TAG.toInt())
+      }
+      sink.readByteString()
     }
   }
 
   /**
-   * Determines equality to another [WorkflowIdentifier] by comparing their [type]s and their
-   * [proxiedIdentifier]s.
+   * Returns a description of this identifier including the name of its workflow type and any
+   * [ImpostorWorkflow.realIdentifier]s.
    */
+  override fun toString(): String =
+    generateSequence(this) { it.proxiedIdentifier }
+        .joinToString { it.typeName }
+        .let { "WorkflowIdentifier($it)" }
+
   override fun equals(other: Any?): Boolean = when {
     this === other -> true
     other !is WorkflowIdentifier -> false
     else -> type == other.type && proxiedIdentifier == other.proxiedIdentifier
   }
 
-  /**
-   * Derives a hashcode from [type] and [proxiedIdentifier].
-   */
   override fun hashCode(): Int {
     var result = type.hashCode()
     result = 31 * result + (proxiedIdentifier?.hashCode() ?: 0)
@@ -97,18 +120,20 @@ class WorkflowIdentifier internal constructor(
     private const val PROXY_IDENTIFIER_TAG = 1.toByte()
 
     /**
-     * Reads a [WorkflowIdentifier] from [source].
+     * Reads a [WorkflowIdentifier] from a [ByteString] as written by [toByteStringOrNull].
      *
      * @throws IllegalArgumentException if the source does not contain a valid [WorkflowIdentifier]
      * @throws ClassNotFoundException if one of the workflow types can't be found in the class
      * loader
      */
-    fun read(source: BufferedSource): WorkflowIdentifier? {
+    fun parse(bytes: ByteString): WorkflowIdentifier? = Buffer().let { source ->
+      source.write(bytes)
+
       try {
         val typeString = source.readUtf8WithLength()
         val proxiedIdentifier = when (source.readByte()) {
           NO_PROXY_IDENTIFIER_TAG -> null
-          PROXY_IDENTIFIER_TAG -> read(source)
+          PROXY_IDENTIFIER_TAG -> parse(source.readByteString())
           else -> throw IllegalArgumentException("Invalid WorkflowIdentifier")
         }
 
@@ -122,6 +147,20 @@ class WorkflowIdentifier internal constructor(
   }
 }
 
+/**
+ * Creates a [WorkflowIdentifier] that is not capable of being snapshotted and will cause any
+ * [ImpostorWorkflow] workflow identified by it to also not be snapshotted.
+ *
+ * **This function should not be used for [ImpostorWorkflow]s that wrap arbitrary workflows**, since
+ * those workflows may expect to be on snapshotted. Using such identifiers _anywhere in the
+ * [ImpostorWorkflow.realIdentifier] chain_ will disable snapshotting for that workflow. **This
+ * function should only be used for [ImpostorWorkflow]s that wrap a closed set of known workflow
+ * types.**
+ */
+@ExperimentalWorkflowApi
+@Suppress("unused")
+fun unsnapshottableIdentifier(type: KType): WorkflowIdentifier = WorkflowIdentifier(type)
+
 /**
  * The [WorkflowIdentifier] that identifies this [Workflow].
  */

workflow-core/src/test/java/com/squareup/workflow/WorkflowIdentifierTest.kt

@@ -16,12 +16,16 @@
 package com.squareup.workflow
 
 import okio.Buffer
+import okio.ByteString
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 import kotlin.test.assertNotEquals
+import kotlin.test.assertNull
 
-@OptIn(ExperimentalWorkflowApi::class)
+@OptIn(ExperimentalWorkflowApi::class, ExperimentalStdlibApi::class)
 class WorkflowIdentifierTest {
 
   @Test fun `flat identifier toString`() {
@@ -43,8 +47,8 @@ class WorkflowIdentifierTest {
 
   @Test fun `restored identifier toString`() {
     val id = TestWorkflow1.identifier
-    val serializedId = Buffer().also(id::write)
-    val restoredId = WorkflowIdentifier.read(serializedId)
+    val serializedId = id.toByteStringOrNull()!!
+    val restoredId = WorkflowIdentifier.parse(serializedId)
     assertEquals(id.toString(), restoredId.toString())
   }
 
@@ -82,24 +86,24 @@ class WorkflowIdentifierTest {
 
   @Test fun `identifier restored from source is equal to itself`() {
     val id = TestWorkflow1.identifier
-    val serializedId = Buffer().also(id::write)
-    val restoredId = WorkflowIdentifier.read(serializedId)
+    val serializedId = id.toByteStringOrNull()!!
+    val restoredId = WorkflowIdentifier.parse(serializedId)
     assertEquals(id, restoredId)
     assertEquals(id.hashCode(), restoredId.hashCode())
   }
 
   @Test fun `identifier restored from source is not equal to different identifier`() {
     val id1 = TestWorkflow1.identifier
     val id2 = TestWorkflow2.identifier
-    val serializedId = Buffer().also(id1::write)
-    val restoredId = WorkflowIdentifier.read(serializedId)
+    val serializedId = id1.toByteStringOrNull()!!
+    val restoredId = WorkflowIdentifier.parse(serializedId)
     assertNotEquals(id2, restoredId)
   }
 
   @Test fun `impostor identifier restored from source is equal to itself`() {
     val id = TestImpostor1(TestWorkflow1).identifier
-    val serializedId = Buffer().also(id::write)
-    val restoredId = WorkflowIdentifier.read(serializedId)
+    val serializedId = id.toByteStringOrNull()!!
+    val restoredId = WorkflowIdentifier.parse(serializedId)
     assertEquals(id, restoredId)
     assertEquals(id.hashCode(), restoredId.hashCode())
   }
@@ -108,48 +112,92 @@ class WorkflowIdentifierTest {
   fun `impostor identifier restored from source is not equal to impostor with different proxied class`() {
     val id1 = TestImpostor1(TestWorkflow1).identifier
     val id2 = TestImpostor1(TestWorkflow2).identifier
-    val serializedId = Buffer().also(id1::write)
-    val restoredId = WorkflowIdentifier.read(serializedId)
+    val serializedId = id1.toByteStringOrNull()!!
+    val restoredId = WorkflowIdentifier.parse(serializedId)
     assertNotEquals(id2, restoredId)
   }
 
   @Test
   fun `impostor identifier restored from source is not equal to different impostor with same proxied class`() {
     val id1 = TestImpostor1(TestWorkflow1).identifier
     val id2 = TestImpostor2(TestWorkflow1).identifier
-    val serializedId = Buffer().also(id1::write)
-    val restoredId = WorkflowIdentifier.read(serializedId)
+    val serializedId = id1.toByteStringOrNull()!!
+    val restoredId = WorkflowIdentifier.parse(serializedId)
     assertNotEquals(id2, restoredId)
   }
 
   @Test fun `read from empty source throws`() {
-    val source = Buffer()
     assertFailsWith<IllegalArgumentException> {
-      WorkflowIdentifier.read(source)
+      WorkflowIdentifier.parse(ByteString.EMPTY)
     }
   }
 
   @Test fun `read from invalid source throws`() {
     val source = Buffer().apply { writeUtf8("invalid data") }
+        .readByteString()
     assertFailsWith<IllegalArgumentException> {
-      WorkflowIdentifier.read(source)
+      WorkflowIdentifier.parse(source)
     }
   }
 
   @Test fun `read from corrupted source throws`() {
-    val source = Buffer().also(TestWorkflow1.identifier::write)
-        .readByteArray()
+    val source = TestWorkflow1.identifier.toByteStringOrNull()!!
+        .toByteArray()
     source.indices.reversed()
         .take(10)
         .forEach { i ->
           source[i] = 0
         }
     val corruptedSource = Buffer().apply { write(source) }
+        .readByteString()
     assertFailsWith<ClassNotFoundException> {
-      WorkflowIdentifier.read(corruptedSource)
+      WorkflowIdentifier.parse(corruptedSource)
     }
   }
 
+  @Test fun `unsnapshottable identifier returns null ByteString`() {
+    val id = unsnapshottableIdentifier(typeOf<TestWorkflow1>())
+    assertNull(id.toByteStringOrNull())
+  }
+
+  @Test fun `unsnapshottable identifier toString()`() {
+    val id = unsnapshottableIdentifier(typeOf<String>())
+    assertEquals(
+        "WorkflowIdentifier(${String::class.java.name} (Kotlin reflection is not available))",
+        id.toString()
+    )
+  }
+
+  @Test fun `unsnapshottable identifiers for same class are equal`() {
+    val id1 = unsnapshottableIdentifier(typeOf<String>())
+    val id2 = unsnapshottableIdentifier(typeOf<String>())
+    assertEquals(id1, id2)
+  }
+
+  @Test fun `unsnapshottable identifiers for different class are not equal`() {
+    val id1 = unsnapshottableIdentifier(typeOf<String>())
+    val id2 = unsnapshottableIdentifier(typeOf<Int>())
+    assertNotEquals(id1, id2)
+  }
+
+  @Test fun `unsnapshottable impostor identifier returns null ByteString`() {
+    val id = TestUnsnapshottableImpostor(typeOf<String>()).identifier
+    assertNull(id.toByteStringOrNull())
+  }
+
+  @Test fun `impostor of unsnapshottable impostor identifier returns null ByteString`() {
+    val id = TestImpostor1(TestUnsnapshottableImpostor(typeOf<String>())).identifier
+    assertNull(id.toByteStringOrNull())
+  }
+
+  @Test fun `unsnapshottable impostor identifier toString()`() {
+    val id = TestUnsnapshottableImpostor(typeOf<String>()).identifier
+    assertEquals(
+        "WorkflowIdentifier(${TestUnsnapshottableImpostor::class.java.name}, " +
+            "${String::class.java.name} (Kotlin reflection is not available))", id.toString()
+    )
+  }
+
   private object TestWorkflow1 : Workflow<Nothing, Nothing, Nothing> {
     override fun asStatefulWorkflow(): StatefulWorkflow<Nothing, *, Nothing, Nothing> =
       throw NotImplementedError()
@@ -175,4 +223,12 @@ class WorkflowIdentifierTest {
     override fun asStatefulWorkflow(): StatefulWorkflow<Nothing, *, Nothing, Nothing> =
       throw NotImplementedError()
   }
+
+  private class TestUnsnapshottableImpostor(
+    type: KType
+  ) : Workflow<Nothing, Nothing, Nothing>, ImpostorWorkflow {
+    override val realIdentifier: WorkflowIdentifier = unsnapshottableIdentifier(type)
+    override fun asStatefulWorkflow(): StatefulWorkflow<Nothing, *, Nothing, Nothing> =
+      throw NotImplementedError()
+  }
 }