chore: Move some reflect out of entity instance

* can be done once per component (class)
akka · Dec 18, 2024 · 5485adf · 5485adf
1 parent a42b7fc
commit 5485adf
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 16 deletions.
diff --git a/akka-javasdk/src/main/scala/akka/javasdk/impl/SdkRunner.scala b/akka-javasdk/src/main/scala/akka/javasdk/impl/SdkRunner.scala
@@ -424,6 +424,11 @@ private final class Sdk(
               method.getName
           }.toSet
 
+        // we preemptively register the events type to the serializer
+        Reflect.allKnownEventSourcedEntityEventType(clz).foreach(serializer.registerTypeHints)
+
+        val entityStateType: Class[AnyRef] = Reflect.eventSourcedEntityStateType(clz).asInstanceOf[Class[AnyRef]]
+
         val instanceFactory: SpiEventSourcedEntity.FactoryContext => SpiEventSourcedEntity = { factoryContext =>
           new EventSourcedEntityImpl[AnyRef, AnyRef, EventSourcedEntity[AnyRef, AnyRef]](
             sdkSettings,
@@ -432,6 +437,7 @@ private final class Sdk(
             factoryContext.entityId,
             serializer,
             ComponentDescriptor.descriptorFor(clz, serializer),
+            entityStateType,
             context =>
               wiredInstance(clz.asInstanceOf[Class[EventSourcedEntity[AnyRef, AnyRef]]]) {
                 // remember to update component type API doc and docs if changing the set of injectables
@@ -446,6 +452,8 @@ private final class Sdk(
 
         val readOnlyCommandNames = Set.empty[String]
 
+        val entityStateType: Class[AnyRef] = Reflect.keyValueEntityStateType(clz).asInstanceOf[Class[AnyRef]]
+
         val instanceFactory: SpiEventSourcedEntity.FactoryContext => SpiEventSourcedEntity = { factoryContext =>
           new KeyValueEntityImpl[AnyRef, KeyValueEntity[AnyRef]](
             sdkSettings,
@@ -454,6 +462,7 @@ private final class Sdk(
             factoryContext.entityId,
             serializer,
             ComponentDescriptor.descriptorFor(clz, serializer),
+            entityStateType,
             context =>
               wiredInstance(clz.asInstanceOf[Class[KeyValueEntity[AnyRef]]]) {
                 // remember to update component type API doc and docs if changing the set of injectables

diff --git a/...-javasdk/src/main/scala/akka/javasdk/impl/eventsourcedentity/EventSourcedEntityImpl.scala b/...-javasdk/src/main/scala/akka/javasdk/impl/eventsourcedentity/EventSourcedEntityImpl.scala
@@ -84,6 +84,7 @@ private[impl] final class EventSourcedEntityImpl[S, E, ES <: EventSourcedEntity[
     entityId: String,
     serializer: JsonSerializer,
     componentDescriptor: ComponentDescriptor,
+    entityStateType: Class[S],
     factory: EventSourcedEntityContext => ES)
     extends SpiEventSourcedEntity {
   import EventSourcedEntityImpl._
@@ -229,5 +230,5 @@ private[impl] final class EventSourcedEntityImpl[S, E, ES <: EventSourcedEntity[
     serializer.toBytes(obj)
 
   override def stateFromBytes(pb: BytesPayload): SpiEventSourcedEntity.State =
-    serializer.fromBytes(router.entityStateType, pb)
+    serializer.fromBytes(entityStateType, pb).asInstanceOf[SpiEventSourcedEntity.State]
 }
diff --git a/.../main/scala/akka/javasdk/impl/eventsourcedentity/ReflectiveEventSourcedEntityRouter.scala b/.../main/scala/akka/javasdk/impl/eventsourcedentity/ReflectiveEventSourcedEntityRouter.scala
@@ -9,7 +9,6 @@ import akka.javasdk.eventsourcedentity.EventSourcedEntity
 import akka.javasdk.impl.CommandHandler
 import akka.javasdk.impl.CommandSerialization
 import akka.javasdk.impl.HandlerNotFoundException
-import akka.javasdk.impl.reflection.Reflect
 import akka.javasdk.impl.serialization.JsonSerializer
 import akka.runtime.sdk.spi.BytesPayload
 
@@ -22,11 +21,6 @@ private[impl] class ReflectiveEventSourcedEntityRouter[S, E, ES <: EventSourcedE
     commandHandlers: Map[String, CommandHandler],
     serializer: JsonSerializer) {
 
-  // we preemptively register the events type to the serializer
-  Reflect.allKnownEventTypes[S, E, ES](entity).foreach(serializer.registerTypeHints)
-
-  val entityStateType: Class[S] = Reflect.eventSourcedEntityStateType(entity.getClass).asInstanceOf[Class[S]]
-
   private def commandHandlerLookup(commandName: String): CommandHandler =
     commandHandlers.get(commandName) match {
       case Some(handler) => handler

diff --git a/akka-javasdk/src/main/scala/akka/javasdk/impl/keyvalueentity/KeyValueEntityImpl.scala b/akka-javasdk/src/main/scala/akka/javasdk/impl/keyvalueentity/KeyValueEntityImpl.scala
@@ -77,6 +77,7 @@ private[impl] final class KeyValueEntityImpl[S, KV <: KeyValueEntity[S]](
     entityId: String,
     serializer: JsonSerializer,
     componentDescriptor: ComponentDescriptor,
+    entityStateType: Class[S],
     factory: KeyValueEntityContext => KV)
     extends SpiEventSourcedEntity {
   import KeyValueEntityEffectImpl._
@@ -205,5 +206,5 @@ private[impl] final class KeyValueEntityImpl[S, KV <: KeyValueEntity[S]](
     serializer.toBytes(obj)
 
   override def stateFromBytes(pb: BytesPayload): SpiEventSourcedEntity.State =
-    serializer.fromBytes(router.entityStateType, pb)
+    serializer.fromBytes(entityStateType, pb).asInstanceOf[SpiEventSourcedEntity.State]
 }
diff --git a/...asdk/src/main/scala/akka/javasdk/impl/keyvalueentity/ReflectiveKeyValueEntityRouter.scala b/...asdk/src/main/scala/akka/javasdk/impl/keyvalueentity/ReflectiveKeyValueEntityRouter.scala
@@ -8,7 +8,6 @@ import akka.annotation.InternalApi
 import akka.javasdk.impl.CommandHandler
 import akka.javasdk.impl.CommandSerialization
 import akka.javasdk.impl.HandlerNotFoundException
-import akka.javasdk.impl.reflection.Reflect
 import akka.javasdk.impl.serialization.JsonSerializer
 import akka.javasdk.keyvalueentity.KeyValueEntity
 import akka.runtime.sdk.spi.BytesPayload
@@ -22,8 +21,6 @@ private[impl] class ReflectiveKeyValueEntityRouter[S, KV <: KeyValueEntity[S]](
     commandHandlers: Map[String, CommandHandler],
     serializer: JsonSerializer) {
 
-  val entityStateType: Class[S] = Reflect.keyValueEntityStateType(entity.getClass).asInstanceOf[Class[S]]
-
   private def commandHandlerLookup(commandName: String): CommandHandler =
     commandHandlers.get(commandName) match {
       case Some(handler) => handler

diff --git a/akka-javasdk/src/main/scala/akka/javasdk/impl/reflection/Reflect.scala b/akka-javasdk/src/main/scala/akka/javasdk/impl/reflection/Reflect.scala
@@ -130,11 +130,6 @@ private[impl] object Reflect {
     Modifier.isStatic(component.getModifiers) &&
     Modifier.isPublic(component.getModifiers)
 
-  def allKnownEventTypes[S, E, ES <: EventSourcedEntity[S, E]](entity: ES): Seq[Class[_]] = {
-    val eventType = eventSourcedEntityEventType(entity.getClass)
-    eventType.getPermittedSubclasses.toSeq
-  }
-
   def workflowStateType[S, W <: Workflow[S]](workflow: W): Class[S] = {
     @tailrec
     def loop(current: Class[_]): Class[_] =
@@ -156,6 +151,11 @@ private[impl] object Reflect {
     loop(workflow.getClass).asInstanceOf[Class[S]]
   }
 
+  def allKnownEventSourcedEntityEventType(component: Class[_]): Seq[Class[_]] = {
+    val eventType = eventSourcedEntityEventType(component)
+    eventType.getPermittedSubclasses.toSeq
+  }
+
   def eventSourcedEntityEventType(component: Class[_]): Class[_] =
     concreteEsApplyEventMethod(component).getParameterTypes.head