Allow exception types to be preserved when serializing ValueOrCompletion objects

Author: rossdanderson

util/src/main/kotlin/com/caplin/integration/datasourcex/util/serialization/fory/DataSourceModule.kt

@@ -8,7 +8,13 @@ import com.caplin.integration.datasourcex.util.flow.ValueOrCompletion
 import org.apache.fory.Fory
 
 /** Registers serializers for internal types with the provided [Fory] instance. */
-fun Fory.registerDataSourceSerializers(): Fory = apply {
+@Suppress("UNCHECKED_CAST")
+fun Fory.registerDataSourceSerializers(preserveExceptionTypes: Boolean = false): Fory = apply {
+  if (preserveExceptionTypes)
+      check(config.trackingRef()) {
+        "Tracking references must be enabled for exception types preservation"
+      }
+
   // Register serializers for FlowMapStreamEvent value classes
   registerSerializer(FlowMapStreamEvent::class.java, FlowMapStreamEventSerializer::class.java)
   registerSerializer(
@@ -24,7 +30,15 @@ fun Fory.registerDataSourceSerializers(): Fory = apply {
   registerSerializer(MapEvent::class.java, MapEventSerializer::class.java)
   registerSerializer(SimpleMapEvent::class.java, SimpleMapEventSerializer::class.java)
   registerSerializer(SetEvent::class.java, SetEventSerializer::class.java)
-  registerSerializer(ValueOrCompletion::class.java, ValueOrCompletionSerializer::class.java)
+
+  registerSerializer(
+      ValueOrCompletion::class.java,
+      ValueOrCompletionSerializer(
+          this,
+          ValueOrCompletion::class.java,
+          preserveExceptionTypes,
+      ),
+  )
 }
 
 fun Fory.registerPersistentCollectionSerializers(): Fory = apply {

util/src/main/kotlin/com/caplin/integration/datasourcex/util/serialization/fory/ValueOrCompletionSerializer.kt

@@ -5,8 +5,11 @@ import org.apache.fory.Fory
 import org.apache.fory.memory.MemoryBuffer
 import org.apache.fory.serializer.Serializer
 
-internal class ValueOrCompletionSerializer(fory: Fory, type: Class<ValueOrCompletion<*>>) :
-    Serializer<ValueOrCompletion<*>>(fory, type) {
+internal class ValueOrCompletionSerializer(
+    fory: Fory,
+    type: Class<ValueOrCompletion<*>>,
+    private val preserveExceptionTypes: Boolean = true,
+) : Serializer<ValueOrCompletion<*>>(fory, type) {
 
   private enum class Type {
     VALUE,
@@ -21,8 +24,12 @@ internal class ValueOrCompletionSerializer(fory: Fory, type: Class<ValueOrComple
       }
       is ValueOrCompletion.Completion -> {
         buffer.writeByte(Type.COMPLETION.ordinal.toByte())
-        val message = value.throwable?.message ?: value.throwable?.toString()
-        fory.writeRef(buffer, message)
+        if (preserveExceptionTypes) {
+          fory.writeRef(buffer, value.throwable)
+        } else {
+          val message = value.throwable?.message ?: value.throwable?.toString()
+          fory.writeRef(buffer, message)
+        }
       }
     }
   }
@@ -34,8 +41,13 @@ internal class ValueOrCompletionSerializer(fory: Fory, type: Class<ValueOrComple
         ValueOrCompletion.Value(value)
       }
       Type.COMPLETION -> {
-        val message = fory.readRef(buffer) as String?
-        ValueOrCompletion.Completion(message?.let { RuntimeException(it) })
+        if (preserveExceptionTypes) {
+          val throwable = fory.readRef(buffer) as Throwable?
+          ValueOrCompletion.Completion(throwable)
+        } else {
+          val message = fory.readRef(buffer) as String?
+          ValueOrCompletion.Completion(message?.let { RuntimeException(it) })
+        }
       }
     }
   }

util/src/test/kotlin/com/caplin/integration/datasourcex/util/serialization/fory/ValueOrCompletionSerializationTest.kt

@@ -3,15 +3,27 @@ package com.caplin.integration.datasourcex.util.serialization.fory
 import com.caplin.integration.datasourcex.util.flow.ValueOrCompletion
 import io.kotest.core.spec.style.FunSpec
 import io.kotest.matchers.shouldBe
+import io.kotest.matchers.types.shouldBeInstanceOf
 import org.apache.fory.Fory
 import org.apache.fory.config.Language
 
+class CustomException(val customMessage: String) : Exception(customMessage)
+
 class ValueOrCompletionSerializationTest :
     FunSpec({
       val fory =
           Fory.builder()
               .withLanguage(Language.JAVA)
               .requireClassRegistration(false)
+              .withRefTracking(true)
+              .build()
+              .registerDataSourceSerializers(preserveExceptionTypes = true)
+
+      val foryNoType =
+          Fory.builder()
+              .withLanguage(Language.JAVA)
+              .requireClassRegistration(false)
+              .withRefTracking(false)
               .build()
               .registerDataSourceSerializers()
 
@@ -37,11 +49,22 @@ class ValueOrCompletionSerializationTest :
           deserialized shouldBe event
         }
 
-        test("Completion") {
-          val event: ValueOrCompletion.Completion = ValueOrCompletion.Completion(null)
+        test("Completion with custom exception (preserved type)") {
+          val event: ValueOrCompletion.Completion =
+              ValueOrCompletion.Completion(CustomException("aah"))
           val bytes = fory.serialize(event)
-          val deserialized = fory.deserialize(bytes)
-          deserialized shouldBe event
+          val deserialized = fory.deserialize(bytes) as ValueOrCompletion.Completion
+          val throwable = deserialized.throwable.shouldBeInstanceOf<CustomException>()
+          throwable.customMessage shouldBe "aah"
+        }
+
+        test("Completion with exception (no type)") {
+          val event: ValueOrCompletion.Completion =
+              ValueOrCompletion.Completion(IllegalStateException("aah"))
+          val bytes = foryNoType.serialize(event)
+          val deserialized = foryNoType.deserialize(bytes) as ValueOrCompletion.Completion
+          deserialized.throwable.shouldBeInstanceOf<RuntimeException>()
+          deserialized.throwable?.message shouldBe "aah"
         }
       }
     })