feat: Run user logic on virtual threads (#2283)

Bumps Akka 2.10.3 for access to upstream feature.

Author: johanandren

project/Dependencies.scala

@@ -20,7 +20,7 @@ object Dependencies {
 
   val ProtobufVersion = akka.grpc.gen.BuildInfo.googleProtobufVersion
 
-  val AkkaVersion = "2.10.2"
+  val AkkaVersion = "2.10.3"
   val AkkaHttpVersion = "10.7.0" // Note: should at least the Akka HTTP version required by Akka gRPC
   val ScalaTestVersion = "3.2.14"
   // https://github.com/akka/akka/blob/main/project/Dependencies.scala#L31
@@ -163,14 +163,10 @@ object Dependencies {
     addSbtPlugin(sbtProtoc),
     addSbtPlugin(akkaGrpc))
 
-  lazy val excludeTheseDependencies: Seq[ExclusionRule] = Seq(
-    // exclusion rules can be added here
-  )
-
   def akkaDependency(name: String, excludeThese: ExclusionRule*) =
-    ("com.typesafe.akka" %% name % AkkaVersion).excludeAll((excludeTheseDependencies ++ excludeThese): _*)
+    ("com.typesafe.akka" %% name % AkkaVersion).excludeAll(excludeThese: _*)
 
   def akkaHttpDependency(name: String, excludeThese: ExclusionRule*) =
-    ("com.typesafe.akka" %% name % AkkaHttpVersion).excludeAll((excludeTheseDependencies ++ excludeThese): _*)
+    ("com.typesafe.akka" %% name % AkkaHttpVersion).excludeAll(excludeThese: _*)
 
 }

sdk/java-sdk-protobuf/src/main/resources/reference.conf

@@ -100,4 +100,12 @@ kalix {
         collector-endpoint =  ${?COLLECTOR_ENDPOINT}
      }
   }
+
+  sdk-dispatcher {
+    executor = "virtual-thread-executor"
+    virtual-thread-executor {
+      # if not on JDK 21
+      fallback="fork-join-executor"
+    }
+  }
 }

sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/SdkExecutionContext.scala

@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2021-2024 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package kalix.javasdk.impl
+
+import akka.actor.ActorSystem
+import akka.annotation.InternalApi
+import akka.stream.ActorAttributes
+import akka.stream.Attributes
+
+import scala.concurrent.ExecutionContext
+
+/**
+ * INTERNAL API
+ */
+@InternalApi
+object SdkExecutionContext {
+  val DispatcherName: String = "kalix.sdk-dispatcher"
+  def apply(system: ActorSystem): ExecutionContext = system.dispatchers.lookup(DispatcherName)
+
+  val streamDispatcher: Attributes = ActorAttributes.dispatcher(DispatcherName)
+
+}

sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/action/ActionsImpl.scala

@@ -25,6 +25,7 @@ import kalix.protocol.component.{ Failure, MetadataEntry }
 import org.slf4j.{ Logger, LoggerFactory, MDC }
 
 import java.util.Optional
+import scala.concurrent.ExecutionContext
 import scala.concurrent.Future
 import scala.jdk.CollectionConverters.SeqHasAsJava
 import scala.jdk.OptionConverters._
@@ -101,9 +102,10 @@ private[javasdk] final class ActionsImpl(_system: ActorSystem, services: Map[Str
 
   import ActionsImpl._
   import _system.dispatcher
-  implicit val system: ActorSystem = _system
+  private implicit val system: ActorSystem = _system
+  private val sdkEc: ExecutionContext = SdkExecutionContext(system)
   private val telemetry = Telemetry(system)
-  lazy val telemetries: Map[String, Instrumentation] = services.values.map { s =>
+  private lazy val telemetries: Map[String, Instrumentation] = services.values.map { s =>
     (s.serviceName, telemetry.traceInstrumentation(s.serviceName, ActionCategory))
   }.toMap
 
@@ -176,24 +178,29 @@ private[javasdk] final class ActionsImpl(_system: ActorSystem, services: Map[Str
     services.get(in.serviceName) match {
       case Some(service) =>
         val span = telemetries(service.serviceName).buildSpan(service, in)
-        span.foreach(s => MDC.put(Telemetry.TRACE_ID, s.getSpanContext.getTraceId))
+
         val fut =
-          try {
-            val context = createContext(in, service.messageCodec, span.map(_.getSpanContext), service.serviceName)
-            val decodedPayload = service.messageCodec.decodeMessage(
-              in.payload.getOrElse(throw new IllegalArgumentException("No command payload")))
-            val effect = service.factory
-              .create(context)
-              .handleUnary(in.name, MessageEnvelope.of(decodedPayload, context.metadata()), context)
-            effectToResponse(service, in, effect, service.messageCodec)
-          } catch {
-            case NonFatal(ex) =>
-              // command handler threw an "unexpected" error
-              span.foreach(_.end())
-              Future.successful(handleUnexpectedException(service, in, ex))
-          } finally {
-            MDC.remove(Telemetry.TRACE_ID)
-          }
+          // Note: invocation in future to guarantee create and invocation is running on sdk dispatcher with virtual thread support
+          Future {
+            try {
+              span.foreach(s => MDC.put(Telemetry.TRACE_ID, s.getSpanContext.getTraceId))
+              val context = createContext(in, service.messageCodec, span.map(_.getSpanContext), service.serviceName)
+              val decodedPayload = service.messageCodec.decodeMessage(
+                in.payload.getOrElse(throw new IllegalArgumentException("No command payload")))
+              val effect = service.factory
+                .create(context)
+                .handleUnary(in.name, MessageEnvelope.of(decodedPayload, context.metadata()), context)
+              effectToResponse(service, in, effect, service.messageCodec)
+            } catch {
+              case NonFatal(ex) =>
+                // command handler threw an "unexpected" error
+                span.foreach(_.end())
+                Future.successful(handleUnexpectedException(service, in, ex))
+            } finally {
+              MDC.remove(Telemetry.TRACE_ID)
+            }
+          }(sdkEc).flatten
+
         fut.andThen { case _ =>
           span.foreach(_.end())
         }
@@ -246,7 +253,7 @@ private[javasdk] final class ActionsImpl(_system: ActorSystem, services: Map[Str
               Future.successful(
                 ActionResponse(ActionResponse.Response.Failure(Failure(0, "Unknown service: " + call.serviceName))))
           }
-      }
+      }(sdkEc)
 
   /**
    * Handle a streamed out command. The input command will contain the service name, command name, request metadata and
@@ -258,25 +265,32 @@ private[javasdk] final class ActionsImpl(_system: ActorSystem, services: Map[Str
   override def handleStreamedOut(in: ActionCommand): Source[ActionResponse, NotUsed] =
     services.get(in.serviceName) match {
       case Some(service) =>
-        try {
-          val context = createContext(in, service.messageCodec, None, service.serviceName)
-          val decodedPayload = service.messageCodec.decodeMessage(
-            in.payload.getOrElse(throw new IllegalArgumentException("No command payload")))
-          service.factory
-            .create(context)
-            .handleStreamedOut(in.name, MessageEnvelope.of(decodedPayload, context.metadata()), context)
-            .asScala
-            .mapAsync(1)(effect => effectToResponse(service, in, effect, service.messageCodec))
-            .recover { case NonFatal(ex) =>
-              // user stream failed with an "unexpected" error
-              handleUnexpectedException(service, in, ex)
+        // Note: invocation in future to guarantee create and invocation is running on sdk dispatcher with virtual thread support
+        Source
+          .futureSource(Future {
+            try {
+              val context = createContext(in, service.messageCodec, None, service.serviceName)
+              val decodedPayload = service.messageCodec.decodeMessage(
+                in.payload.getOrElse(throw new IllegalArgumentException("No command payload")))
+              service.factory
+                .create(context)
+                .handleStreamedOut(in.name, MessageEnvelope.of(decodedPayload, context.metadata()), context)
+                .asScala
+                .mapAsync(1)(effect => effectToResponse(service, in, effect, service.messageCodec))
+                .recover { case NonFatal(ex) =>
+                  // user stream failed with an "unexpected" error
+                  handleUnexpectedException(service, in, ex)
+                }
+                // run the stream itself on the virtual thread dispatcher in case the user blocks in stream
+                .addAttributes(SdkExecutionContext.streamDispatcher)
+            } catch {
+              case NonFatal(ex) =>
+                // command handler threw an "unexpected" error
+                Source.single(handleUnexpectedException(service, in, ex))
             }
-            .async
-        } catch {
-          case NonFatal(ex) =>
-            // command handler threw an "unexpected" error
-            Source.single(handleUnexpectedException(service, in, ex))
-        }
+          }(sdkEc))
+          .mapMaterializedValue(_ => NotUsed)
+
       case None =>
         Source.single(ActionResponse(ActionResponse.Response.Failure(Failure(0, "Unknown service: " + in.serviceName))))
     }
@@ -303,25 +317,32 @@ private[javasdk] final class ActionsImpl(_system: ActorSystem, services: Map[Str
         case (Seq(call), messages) =>
           services.get(call.serviceName) match {
             case Some(service) =>
+              // Note: invocation in future to guarantee create and invocation is running on sdk dispatcher with virtual thread support
               try {
-                val context = createContext(call, service.messageCodec, None, service.serviceName)
-                service.factory
-                  .create(context)
-                  .handleStreamed(
-                    call.name,
-                    messages.map { message =>
-                      val metadata = MetadataImpl.of(message.metadata.map(_.entries.toVector).getOrElse(Nil))
-                      val decodedPayload = service.messageCodec.decodeMessage(
-                        message.payload.getOrElse(throw new IllegalArgumentException("No command payload")))
-                      MessageEnvelope.of(decodedPayload, metadata)
-                    }.asJava,
-                    context)
-                  .asScala
-                  .mapAsync(1)(effect => effectToResponse(service, call, effect, service.messageCodec))
-                  .recover { case NonFatal(ex) =>
-                    // user stream failed with an "unexpected" error
-                    handleUnexpectedException(service, call, ex)
-                  }
+                Source
+                  .futureSource(Future {
+                    val context = createContext(call, service.messageCodec, None, service.serviceName)
+                    service.factory
+                      .create(context)
+                      .handleStreamed(
+                        call.name,
+                        messages.map { message =>
+                          val metadata = MetadataImpl.of(message.metadata.map(_.entries.toVector).getOrElse(Nil))
+                          val decodedPayload = service.messageCodec.decodeMessage(
+                            message.payload.getOrElse(throw new IllegalArgumentException("No command payload")))
+                          MessageEnvelope.of(decodedPayload, metadata)
+                        }.asJava,
+                        context)
+                      .asScala
+                      .mapAsync(1)(effect => effectToResponse(service, call, effect, service.messageCodec))
+                      .recover { case NonFatal(ex) =>
+                        // user stream failed with an "unexpected" error
+                        handleUnexpectedException(service, call, ex)
+                      }
+                      // run the stream itself on the virtual thread dispatcher in case the user blocks in stream
+                      .addAttributes(SdkExecutionContext.streamDispatcher)
+                  }(sdkEc))
+                  .mapMaterializedValue(_ => NotUsed)
               } catch {
                 case NonFatal(ex) =>
                   // command handler threw an "unexpected" error

sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/eventsourcedentity/EventSourcedEntitiesImpl.scala

@@ -135,6 +135,7 @@ final class EventSourcedEntitiesImpl(
           EventSourcedStreamOut(OutFailure(Failure(description = s"Unexpected failure [$correlationId]")))
         }
       }
+      .withAttributes(SdkExecutionContext.streamDispatcher) // factory instance invoked on this stream
   }
 
   private def runEntity(init: EventSourcedInit): Flow[EventSourcedStreamIn, EventSourcedStreamOut, NotUsed] = {
@@ -260,7 +261,7 @@ final class EventSourcedEntitiesImpl(
           EventSourcedStreamOut(OutFailure(Failure(description = s"Unexpected failure [$correlationId]")))
         }
       }
-      .async
+      .addAttributes(SdkExecutionContext.streamDispatcher)
   }
 
   private class CommandContextImpl(

sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/replicatedentity/ReplicatedEntitiesImpl.scala

@@ -87,7 +87,7 @@ final class ReplicatedEntitiesImpl(system: ActorSystem, services: Map[String, Re
           ReplicatedEntityStreamOut(Out.Failure(Failure(description = s"Unexpected error [$correlationId]")))
         }
       }
-      .async
+      .addAttributes(SdkExecutionContext.streamDispatcher)
 
   private def runEntity(
       init: ReplicatedEntityInit): Flow[ReplicatedEntityStreamIn, ReplicatedEntityStreamOut, NotUsed] = {
@@ -138,6 +138,7 @@ object ReplicatedEntitiesImpl {
     val router: ReplicatedEntityRouter[_ <: Object, _ <: Object] = {
       val context = new ReplicatedEntityCreationContext(entityId, system)
       try {
+        // Note: not run on sdk execution context
         service.factory.create(context)
       } finally {
         context.deactivate()

sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/valueentity/ValueEntitiesImpl.scala

@@ -107,7 +107,7 @@ final class ValueEntitiesImpl(
           ValueEntityStreamOut(OutFailure(Failure(description = s"Unexpected error [$correlationId]")))
         }
       }
-      .async
+      .addAttributes(SdkExecutionContext.streamDispatcher)
 
   private def runEntity(init: ValueEntityInit): Flow[ValueEntityStreamIn, ValueEntityStreamOut, NotUsed] = {
     val service =

sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/view/ViewsImpl.scala

@@ -156,7 +156,7 @@ final class ViewsImpl(system: ActorSystem, _services: Map[String, ViewService])
             s"Kalix protocol failure: expected ReceiveEvent message, but got ${other.getClass.getName}"
           Source.failed(new RuntimeException(errMsg))
       }
-      .async
+      .addAttributes(SdkExecutionContext.streamDispatcher)
 
   private final class UpdateContextImpl(
       override val viewId: String,

sdk/java-sdk-protobuf/src/main/scala/kalix/javasdk/impl/workflow/WorkflowImpl.scala

@@ -125,7 +125,7 @@ final class WorkflowImpl(system: ActorSystem, val services: Map[String, Workflow
           WorkflowStreamOut(OutFailure(component.Failure(description = s"Unexpected error [$correlationId]")))
         }
       }
-      .async
+      .withAttributes(SdkExecutionContext.streamDispatcher)
 
   private def toRecoverStrategy(messageCodec: MessageCodec)(
       recoverStrategy: AbstractWorkflow.RecoverStrategy[_]): RecoverStrategy = {
@@ -169,6 +169,7 @@ final class WorkflowImpl(system: ActorSystem, val services: Map[String, Workflow
     WorkflowConfig(workflowTimeout, failoverTo, failoverRecovery, Some(stepConfig), stepConfigs)
   }
 
+  // Note: called from stream, already on sdk dispatcher
   private def runWorkflow(
       init: WorkflowEntityInit): (Flow[WorkflowStreamIn, WorkflowStreamOut, NotUsed], WorkflowStreamOut) = {
     val service =
@@ -345,6 +346,7 @@ final class WorkflowImpl(system: ActorSystem, val services: Map[String, Workflow
           // currently added to satisfy the compiler
           Future.successful(WorkflowStreamOut(WorkflowStreamOut.Message.Empty))
       }
+      .addAttributes(SdkExecutionContext.streamDispatcher)
 
     (flow, workflowConfig)
   }