fix: SDK gRPC endpoint dispatchers (#2097)

Author: johanandren

build.sbt

@@ -270,6 +270,21 @@ lazy val pluginTesterJava = Project(id = "akka-grpc-plugin-tester-java", base =
     ReflectiveCodeGen.codeGeneratorSettings ++= Seq("server_power_apis"))
   .pluginTestingSettings
 
+lazy val pluginTesterSdkHandler =
+  Project(id = "akka-grpc-plugin-tester-sdk-handler", base = file("plugin-tester-sdk-handler"))
+    .disablePlugins(MimaPlugin, CiReleasePlugin)
+    .settings(Dependencies.pluginTester)
+    .settings(
+      (publish / skip) := true,
+      fork := true,
+      PB.protocVersion := Dependencies.Versions.googleProtobuf,
+      ReflectiveCodeGen.generatedLanguages := Seq("Java"),
+      crossScalaVersions := Dependencies.Versions.CrossScalaForLib,
+      scalaVersion := Dependencies.Versions.CrossScalaForLib.head,
+      Test / parallelExecution := false,
+      ReflectiveCodeGen.codeGeneratorSettings ++= Seq("generate_scala_handler_factory"))
+    .pluginTestingSettings
+
 lazy val root = Project(id = "akka-grpc", base = file("."))
   .disablePlugins(SitePlugin, MimaPlugin, CiReleasePlugin)
   .aggregate(
@@ -281,6 +296,7 @@ lazy val root = Project(id = "akka-grpc", base = file("."))
     interopTests,
     pluginTesterScala,
     pluginTesterJava,
+    pluginTesterSdkHandler,
     benchmarks,
     docs)
   .settings(

codegen/src/main/twirl/templates/JavaServer/ScalaHandler.scala.txt

@@ -58,36 +58,36 @@ public final class @{serviceName}ScalaHandlerFactory implements InstancePerReque
     private static final CompletionStage<akka.http.javadsl.model.HttpResponse> unsupportedMediaType = CompletableFuture.completedFuture(
       akka.http.javadsl.model.HttpResponse.create().withStatus(akka.http.javadsl.model.StatusCodes.UNSUPPORTED_MEDIA_TYPE));
 
-
-  /**
+    /**
      * Creates a `HttpRequest` to `HttpResponse` handler that can be used in for example
      * `Http.get(system).bindAndHandleAsync`. It ends with `StatusCodes.NotFound` if the request is not matching.
      *
      * Use {@@link akka.grpc.javadsl.ServiceHandler#concatOrNotFound} when combining several services.
      */
-  @@SuppressWarnings("All")
-  @@Override
-  public PartialFunction<akka.http.scaladsl.model.HttpRequest, Future<akka.http.scaladsl.model.HttpResponse>> partialInstancePerRequest(Function1<akka.http.scaladsl.model.HttpRequest, @{serviceName}> serviceFactory, String prefix, PartialFunction<Throwable, Trailers> eHandler, ClassicActorSystemProvider systemProvider) {
-      InstancePerRequestPF.GrpcMethod[] methods = {
-        new InstancePerRequestPF.GrpcMethod<@{serviceName}>(
-          "@service.methods.head.grpcName",
-          this::@{service.methods.head.name})
-          @for(method <- service.methods.tail) {
-            ,
-            new InstancePerRequestPF.GrpcMethod<@{serviceName}>(
-              "@{method.grpcName}",
-              this::@{method.name})
-          }
-      };
-
-      return new InstancePerRequestPF<@{serviceName}>(
-          serviceFactory,
-          prefix,
-          methods,
-          eHandler,
-          systemProvider
-      );
-  }
+    @@Override
+    @@SuppressWarnings({"all","unchecked"})
+    public PartialFunction<akka.http.scaladsl.model.HttpRequest, Future<akka.http.scaladsl.model.HttpResponse>> partialInstancePerRequest(Function1<akka.http.scaladsl.model.HttpRequest, @{serviceName}> serviceFactory, String prefix, PartialFunction<Throwable, Trailers> eHandler, ClassicActorSystemProvider systemProvider, Materializer materializer) {
+        InstancePerRequestPF.GrpcMethod[] methods = {
+          new InstancePerRequestPF.GrpcMethod<@{serviceName}>(
+            "@service.methods.head.grpcName",
+            this::@{service.methods.head.name})
+            @for(method <- service.methods.tail) {
+              ,
+              new InstancePerRequestPF.GrpcMethod<@{serviceName}>(
+                "@{method.grpcName}",
+                this::@{method.name})
+            }
+        };
+
+        return new InstancePerRequestPF<@{serviceName}>(
+            serviceFactory,
+            prefix,
+            methods,
+            eHandler,
+            systemProvider,
+            materializer
+        );
+    }
 
     public String getServiceName() {
       return @{service.name}.name;

plugin-tester-sdk-handler/LICENSE

@@ -0,0 +1,3 @@
+# Test module for testing the SDK-specific server handler
+
+Tests run through the main sbt project

plugin-tester-sdk-handler/README.md

@@ -0,0 +1,20 @@
+
+syntax = "proto3";
+
+option java_multiple_files = true;
+option java_package = "example.myapp.helloworld.grpc";
+option java_outer_classname = "HelloWorldProto";
+
+package helloworld;
+
+service GreeterService {
+    rpc SayHello (HelloRequest) returns (HelloReply) {}
+    rpc ItKeepsReplying (HelloRequest) returns (stream HelloReply) {}
+}
+message HelloRequest {
+    string name = 1;
+}
+
+message HelloReply {
+    string message = 1;
+}

plugin-tester-sdk-handler/src/main/proto/helloworld.proto

@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2025 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package com.example
+
+import akka.NotUsed
+import akka.actor.typed.ActorSystem
+import akka.actor.testkit.typed.scaladsl.ActorTestKit
+import akka.actor.typed.scaladsl.Behaviors
+import akka.grpc.GrpcClientSettings
+import akka.grpc.scaladsl.GrpcExceptionHandler
+import akka.http.scaladsl.Http
+import akka.stream.ActorAttributes
+import akka.stream.Materializer
+import akka.stream.javadsl.Source
+import akka.stream.scaladsl.Sink
+import com.typesafe.config.ConfigFactory
+import example.myapp.helloworld.grpc.GreeterService
+import example.myapp.helloworld.grpc.GreeterServiceClient
+import example.myapp.helloworld.grpc.GreeterServiceScalaHandlerFactory
+import example.myapp.helloworld.grpc.HelloReply
+import example.myapp.helloworld.grpc.HelloRequest
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.concurrent.ScalaFutures
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.time.Span
+import org.scalatest.wordspec.AnyWordSpec
+
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionStage
+import scala.concurrent.duration.DurationInt
+import scala.jdk.FutureConverters.CompletionStageOps
+
+class ServiceExecutionContextSpec extends AnyWordSpec with Matchers with ScalaFutures with BeforeAndAfterAll {
+
+  implicit val system: ActorSystem[Nothing] = ActorSystem[Any](
+    Behaviors.empty,
+    "ServiceExecutionContextSpec",
+    ConfigFactory.parseString("""
+      akka.http.server.enable-http2 = on
+      custom-dispatcher = {
+        type = "Dispatcher"
+        executor = "fork-join-executor"
+      }
+      """))
+  override implicit val patienceConfig: PatienceConfig = PatienceConfig(5.seconds, Span(100, org.scalatest.time.Millis))
+
+  private final class GreeterServiceImpl extends GreeterService {
+
+    private def replyWithCurrentThreadName() =
+      HelloReply.newBuilder().setMessage(Thread.currentThread().getName).build()
+
+    override def sayHello(in: HelloRequest): CompletionStage[HelloReply] = {
+      CompletableFuture.completedFuture(replyWithCurrentThreadName())
+    }
+
+    override def itKeepsReplying(in: HelloRequest): Source[HelloReply, NotUsed] = {
+      // invokin thread
+      Source
+        .single(replyWithCurrentThreadName())
+        .concat(
+          Source
+            .single(in)
+            // thread name in the running stream (materialized by akka grpc)
+            .map(_ => replyWithCurrentThreadName()))
+    }
+  }
+
+  "The SDK instance per request handler" should {
+
+    val customMaterializer = Materializer(system, ActorAttributes.dispatcher("custom-dispatcher"))
+
+    val handler = new GreeterServiceScalaHandlerFactory()
+
+    val handlerPf = handler.partialInstancePerRequest(
+      _ => new GreeterServiceImpl,
+      "",
+      GrpcExceptionHandler.defaultMapper(system.classicSystem),
+      system,
+      customMaterializer)
+
+    val bound = Http().newServerAt("127.0.0.1", 0).withMaterializer(customMaterializer).bind(handlerPf).futureValue
+    val port = bound.localAddress.getPort
+
+    val client =
+      GreeterServiceClient.create(GrpcClientSettings.connectToServiceAt("127.0.0.1", port).withTls(false), system)
+
+    "run unary requests on the specified execution context" in {
+      val reply = client.sayHello(HelloRequest.getDefaultInstance).asScala.futureValue
+
+      reply.getMessage should include("custom-dispatcher")
+    }
+
+    "run streamed responses and their streams on specified execution context" in {
+      val replies = client.itKeepsReplying(HelloRequest.getDefaultInstance).asScala.runWith(Sink.seq).futureValue
+
+      replies should have size 2
+      replies(0).getMessage should include("custom-dispatcher")
+      // Note: this is only true because we started http server with the materializer as well
+      replies(1).getMessage should include("custom-dispatcher")
+    }
+
+  }
+
+  override def afterAll(): Unit = {
+    ActorTestKit.shutdown(system)
+  }
+
+}

plugin-tester-sdk-handler/src/test/scala/com/example/ServiceExecutionContextSpec.scala

@@ -157,4 +157,13 @@ object Dependencies {
     Test.scalaTestPlusJunit,
     Test.akkaTestkitTyped,
     GrpcApi.googleApiProtos)
+
+  val sdkPluginTester = l ++= Seq(
+    // usually automatically added by `suggestedDependencies`, which doesn't work with ReflectiveCodeGen
+    Compile.grpcStub,
+    Compile.akkaPki,
+    Runtime.logback,
+    Test.scalaTest,
+    Test.scalaTestPlusJunit,
+    Test.akkaTestkitTyped)
 }

project/Dependencies.scala

@@ -16,7 +16,7 @@ import akka.stream.Materializer
 import akka.stream.SystemMaterializer
 
 import java.util.concurrent.CompletionStage
-import scala.concurrent.ExecutionContext
+
 import scala.concurrent.Future
 import scala.jdk.FutureConverters.CompletionStageOps
 
@@ -51,12 +51,20 @@ private[akka] final class InstancePerRequestPF[S](
     prefix: String,
     methods: Array[InstancePerRequestPF.GrpcMethod[S]],
     eHandler: PartialFunction[Throwable, Trailers],
-    system: ClassicActorSystemProvider)
+    system: ClassicActorSystemProvider,
+    materializer: Materializer)
     extends PartialFunction[HttpRequest, Future[HttpResponse]] {
-  // Note how the factory and partial function is Akka HTTP scaladsl. This is intentional.
 
-  private val materializer: Materializer = SystemMaterializer(system).materializer
-  private implicit val ec: ExecutionContext = materializer.executionContext
+  // for bin/source comp with existing generated code
+  def this(
+      factory: HttpRequest => S,
+      prefix: String,
+      methods: Array[InstancePerRequestPF.GrpcMethod[S]],
+      eHandler: PartialFunction[Throwable, Trailers],
+      system: ClassicActorSystemProvider) =
+    this(factory, prefix, methods, eHandler, system, SystemMaterializer(system).materializer)
+
+  // Note how the factory and partial function is Akka HTTP scaladsl. This is intentional.
   private val spi = TelemetryExtension(system).spi
   private val javaEHandler: akka.japi.Function[ActorSystem, akka.japi.Function[Throwable, Trailers]] =
     (_: ActorSystem) => { eHandler.apply _ }

runtime/src/main/scala/akka/grpc/internal/InstancePerRequestPF.scala

@@ -10,7 +10,10 @@ import akka.annotation.DoNotInherit
 import akka.grpc.Trailers
 import akka.http.scaladsl.model.HttpRequest
 import akka.http.scaladsl.model.HttpResponse
+import akka.stream.Materializer
+import akka.stream.SystemMaterializer
 
+import scala.annotation.nowarn
 import scala.concurrent.Future
 
 /**
@@ -24,10 +27,32 @@ import scala.concurrent.Future
 @DoNotInherit
 @ApiMayChange
 trait InstancePerRequestFactory[S] {
+
+  // Note: this may be strange, one of these two methods is always overridden, either the old one, or the new
+  // having this cycle avoids having to implement the old method in new generated code just for compatibility
+
+  @nowarn("msg=deprecated")
+  def partialInstancePerRequest(
+      serviceFactory: HttpRequest => S,
+      prefix: String,
+      eHandler: PartialFunction[Throwable, Trailers],
+      systemProvider: ClassicActorSystemProvider,
+      materializer: Materializer): PartialFunction[HttpRequest, Future[HttpResponse]] = {
+    partialInstancePerRequest(serviceFactory, prefix, eHandler, systemProvider)
+  }
+
+  // for compatibility with existing generated sources and SDK
+  @deprecated(since = "2.5.9")
   def partialInstancePerRequest(
       serviceFactory: HttpRequest => S,
       prefix: String,
       eHandler: PartialFunction[Throwable, Trailers],
-      systemProvider: ClassicActorSystemProvider): PartialFunction[HttpRequest, Future[HttpResponse]]
+      systemProvider: ClassicActorSystemProvider): PartialFunction[HttpRequest, Future[HttpResponse]] =
+    partialInstancePerRequest(
+      serviceFactory,
+      prefix,
+      eHandler,
+      systemProvider,
+      SystemMaterializer(systemProvider).materializer)
 
 }

runtime/src/main/scala/akka/grpc/scaladsl/InstancePerRequestFactory.scala

@@ -0,0 +1,25 @@
+/*
+ * Copyright (C) 2024-2025 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.grpc.internal;
+
+import akka.actor.ClassicActorSystemProvider;
+import akka.grpc.Trailers;
+import akka.grpc.scaladsl.InstancePerRequestFactory;
+import akka.http.scaladsl.model.HttpRequest;
+import akka.http.scaladsl.model.HttpResponse;
+import akka.stream.Materializer;
+import scala.Function1;
+import scala.PartialFunction;
+import scala.concurrent.Future;
+
+// in previous versions of the plugin this would be generated
+public class OldTestServiceScalaHandlerFactory implements InstancePerRequestFactory<TestService> {
+
+    @Override
+    public PartialFunction<HttpRequest, Future<HttpResponse>> partialInstancePerRequest(Function1<HttpRequest, TestService> serviceFactory, String prefix, PartialFunction<Throwable, Trailers> eHandler, ClassicActorSystemProvider systemProvider) {
+        return null;
+    }
+    
+}