Adds support for wire.kotlin rpcCallStyle =

"suspending"

GitOrigin-RevId: 05fa1f17678cdc51f2f2e57acffd638038fc81d4

Author: jclyne

docs/actions.md

@@ -187,6 +187,41 @@ class GreeterActionModule : KAbstractModule() {
 }
 ```
 
+Misk also supports `rpcCallStyle = "suspending"` for suspending gRPC actions.This is the preferred way to generate server 
+actions if you intend on using coroutines to implement the business logic of your action. See [coroutines](coroutines.md) for more information.
+
+```kotlin
+wire {
+  sourcePath {
+    srcDir("src/main/proto")
+  }
+
+  kotlin {
+    include("squareup.cash.hello.GreeterService")
+    rpcCallStyle = "suspending"
+    rpcRole = "server"
+    singleMethodServices = true
+  }
+
+  java {
+  }
+}
+```
+
+The above will generate a similar action class, but with a suspending action function
+
+```kotlin
+@Singleton
+class HelloGrpcAction @Inject internal constructor()
+  : GreeterServiceHelloBlockingServer, WebAction {
+
+  @Unauthorized
+  override suspend fun Hello(request: HelloRequest): HelloResponse {
+    return HelloResponse("message")
+  }
+}
+```
+
 Creating a gRPC action automatically creates a JSON endpoint with all of the same annotations in the 
 path defined by the `...BlockingServer`, typically `/<package>.<service name>/<rpc name>`.
 

docs/coroutines.md

@@ -0,0 +1,49 @@
+Misk Coroutines Rules
+
+Coroutines are cooperative concurrency. If you don't cooperate (ie. suspend regularly), things will not be as efficient 
+and may potentially deadlock.  
+
+Misk allows for per request threading for which coroutines can be used to model concurrent operations. If the coroutines suspend,
+then multiple coroutines can be run in parallel on the same thread. If coroutines need to block, they will run serially unless
+a `Dispatcher` is provided that allows for running on multiple threads. Both models are supported.
+
+Mixing thread based concurrency primitives, to synchronize coroutines, can result in deadlocks if threads are blocked waiting
+Example
+```class DangerTest {
+  @Test
+  fun threads() {
+    val latch = CountDownLatch(3)
+    thread { latch.countDown() }
+    thread { latch.countDown() }
+    thread { latch.countDown() }
+    latch.await()
+  }
+
+  @Test
+  fun coroutineDeadlock(){
+    runBlocking {
+      val latch = CountDownLatch(3)
+      async { latch.countDown() }
+      async { latch.countDown() }
+      async { latch.countDown() }
+      latch.await()
+    }
+  }
+}```
+
+
+When an action is declared with the `suspend` modifier, it'll be called with a `Dispatcher` that has a single backing
+thread (`runBlocking`). This thread is part of the Jetty Thread Pool and allocated to this specific request, therefore
+it is safe to make blocking calls on. This will also take care of request scoped features that are thread local, 
+such as ActionScoped values, MDC, tracing, etc.
+
+When that function returns, the request and response body streams will be flushed and closed immediately.
+The framework will release these resources for you.
+
+Follow structured concurrency best practices, including:
+ - All coroutines should be in child scopes for the incoming request scope.
+ - Don't use `GlobalScope` or create a new `CoroutineScope()` object.
+
+
+
+

gradle/libs.versions.toml

@@ -160,6 +160,9 @@ kotlinStdLibJdk8 = { module = "org.jetbrains.kotlin:kotlin-stdlib-jdk8", version
 kotlinTest = { module = "org.jetbrains.kotlin:kotlin-test", version.ref = "kotlin" }
 kotlinXHtml = { module = "org.jetbrains.kotlinx:kotlinx-html-jvm", version = "0.11.0" }
 kotlinxCoroutinesCore = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-core", version = "1.9.0" }
+kotlinxCoroutinesDebug = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-debug", version = "1.9.0" }
+kotlinxCoroutinesJdk8 = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-jdk8", version = "1.9.0" }
+kotlinxCoroutinesTest = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-test", version = "1.9.0" }
 kubernetesClient = { module = "io.kubernetes:client-java", version = "18.0.1" }
 kubernetesClientApi = { module = "io.kubernetes:client-java-api", version = "18.0.1" }
 launchDarkly = { module = "com.launchdarkly:launchdarkly-java-server-sdk", version = "6.3.0" }

misk-action-scopes/src/main/kotlin/misk/scope/ActionScope.kt

@@ -31,6 +31,10 @@ class ActionScope @Inject internal constructor(
   /**
    * Wraps a [kotlinx.coroutines.runBlocking] to propagate the current action scope.
    */
+  @Deprecated(
+    message = "don't use runBlocking explicitly",
+    replaceWith = ReplaceWith("use suspending invocation")
+  )
   fun <T> runBlocking(block: suspend CoroutineScope.() -> T): T {
     return if (inScope()) {
       kotlinx.coroutines.runBlocking(asContextElement(), block)
@@ -44,6 +48,10 @@ class ActionScope @Inject internal constructor(
   /**
    * Wraps a [kotlinx.coroutines.runBlocking] to propagate the current action scope.
    */
+  @Deprecated(
+    message = "don't use runBlocking explicitly",
+    replaceWith = ReplaceWith("use suspending invocation")
+  )
   fun <T> runBlocking(context: CoroutineContext, block: suspend CoroutineScope.() -> T): T {
     return if (inScope()) {
       kotlinx.coroutines.runBlocking(context + asContextElement(), block)

misk/api/misk.api

@@ -1683,10 +1683,6 @@ public final class misk/web/actions/StatusAction$ServerStatus {
 	public fun toString ()Ljava/lang/String;
 }
 
-public final class misk/web/actions/WebActionsKt {
-	public static final fun asChain (Lmisk/web/actions/WebAction;Lkotlin/reflect/KFunction;Ljava/util/List;Ljava/util/List;Lmisk/web/HttpCall;)Lmisk/Chain;
-}
-
 public abstract interface class misk/web/concurrencylimits/ConcurrencyLimiterFactory {
 	public abstract fun create (Lmisk/Action;)Lcom/netflix/concurrency/limits/Limiter;
 }

misk/build.gradle.kts

@@ -53,6 +53,7 @@ dependencies {
   implementation(libs.loggingApi)
   implementation(libs.kotlinReflect)
   implementation(libs.kotlinStdLibJdk8)
+  implementation(libs.kotlinxCoroutinesCore)
   implementation(libs.moshiAdapters)
   implementation(libs.okio)
   implementation(libs.openTracingConcurrent)
@@ -78,7 +79,9 @@ dependencies {
   testImplementation(libs.guavaTestLib)
   testImplementation(libs.junitApi)
   testImplementation(libs.junitParams)
+  testImplementation(libs.kotestAssertions)
   testImplementation(libs.kotlinTest)
+  testImplementation(libs.kotlinxCoroutinesTest)
   testImplementation(libs.logbackClassic)
   testImplementation(libs.okHttpMockWebServer)
   testImplementation(libs.openTracingMock)

misk/src/main/kotlin/misk/grpc/Grpc.kt

@@ -2,6 +2,8 @@ package misk.grpc
 
 import com.squareup.wire.MessageSink
 import com.squareup.wire.MessageSource
+import kotlinx.coroutines.channels.ReceiveChannel
+import kotlinx.coroutines.channels.SendChannel
 import java.lang.reflect.ParameterizedType
 import java.lang.reflect.Type
 import java.lang.reflect.WildcardType
@@ -10,13 +12,15 @@ import kotlin.reflect.jvm.javaType
 
 /**
  * Returns the stream element type, like `MyRequest` if this is `MessageSource<MyRequest>`.
- * Returns null if this is not a [MessageSource] or [MessageSink].
+ * Returns null if this is not a [MessageSource], [MessageSink], [SendChannel] or [ReceiveChannel].
  */
 internal fun KType.streamElementType(): Type? {
   // Unbox the type parameter.
   val parameterizedType = javaType as? ParameterizedType ?: return null
   if (parameterizedType.rawType != MessageSource::class.java &&
-    parameterizedType.rawType != MessageSink::class.java
+    parameterizedType.rawType != MessageSink::class.java &&
+    parameterizedType.rawType != SendChannel::class.java &&
+    parameterizedType.rawType != ReceiveChannel::class.java
   ) return null
   // Remove the wildcard, like 'out MessageSource' (Kotlin) or '? super MessageSource' (Java).
   return when (val typeArgument = parameterizedType.actualTypeArguments[0]) {

misk/src/main/kotlin/misk/grpc/GrpcFeatureBinding.kt

@@ -2,25 +2,33 @@ package misk.grpc
 
 import com.squareup.wire.ProtoAdapter
 import com.squareup.wire.WireRpc
+import jakarta.inject.Inject
+import jakarta.inject.Singleton
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.Channel.Factory.RENDEZVOUS
 import misk.Action
 import misk.web.DispatchMechanism
 import misk.web.FeatureBinding
 import misk.web.FeatureBinding.Claimer
 import misk.web.FeatureBinding.Subject
-import misk.web.Grpc
 import misk.web.PathPattern
+import misk.web.WebConfig
 import misk.web.actions.findAnnotationWithOverrides
 import misk.web.mediatype.MediaTypes
 import java.lang.reflect.Type
-import jakarta.inject.Inject
-import jakarta.inject.Singleton
+import kotlin.coroutines.CoroutineContext
 
 internal class GrpcFeatureBinding(
   private val requestAdapter: ProtoAdapter<Any>,
   private val responseAdapter: ProtoAdapter<Any>,
   private val streamingRequest: Boolean,
-  private val streamingResponse: Boolean
+  private val streamingResponse: Boolean,
+  private val isSuspend: Boolean,
+  private val grpcMessageSourceChannelContext: CoroutineContext
 ) : FeatureBinding {
+
   override fun beforeCall(subject: Subject) {
     val requestBody = subject.takeRequestBody()
     val messageSource = GrpcMessageSource(
@@ -29,7 +37,19 @@ internal class GrpcFeatureBinding(
     )
 
     if (streamingRequest) {
-      subject.setParameter(0, messageSource)
+      val param: Any = if (isSuspend) {
+        GrpcMessageSourceChannel(
+          channel = Channel(
+            capacity = RENDEZVOUS,
+            onBufferOverflow = BufferOverflow.SUSPEND,
+          ),
+          source = messageSource,
+          coroutineContext = grpcMessageSourceChannelContext,
+        )
+      } else {
+        messageSource
+      }
+      subject.setParameter(0, param)
     } else {
       val request = messageSource.read()!!
       subject.setParameter(0, request)
@@ -38,9 +58,20 @@ internal class GrpcFeatureBinding(
     if (streamingResponse) {
       val responseBody = subject.takeResponseBody()
       val messageSink = GrpcMessageSink(responseBody, responseAdapter, grpcEncoding = "identity")
+      val param: Any = if (isSuspend) {
+        GrpcMessageSinkChannel(
+          channel = Channel(
+            capacity = RENDEZVOUS,
+            onBufferOverflow = BufferOverflow.SUSPEND,
+          ),
+          sink = messageSink,
+        )
+      } else {
+        messageSink
+      }
 
       // It's a streaming response, give the call a SendChannel to write to.
-      subject.setParameter(1, messageSink)
+      subject.setParameter(1, param)
       setResponseHeaders(subject)
     }
   }
@@ -72,7 +103,19 @@ internal class GrpcFeatureBinding(
   }
 
   @Singleton
-  class Factory @Inject internal constructor() : FeatureBinding.Factory {
+  class Factory @Inject internal constructor(
+    webConfig: WebConfig
+  ) : FeatureBinding.Factory {
+
+    // This dispatcher is sized to the jetty thread pool size to make sure that
+    // no requests that are currently scheduled on a jetty thread are ever blocked
+    // from reading a streaming request
+    private val grpcMessageSourceChannelDispatcher =
+      Dispatchers.IO.limitedParallelism(
+        parallelism = webConfig.jetty_max_thread_pool_size,
+        name = "GrpcMessageSourceChannel.bridgeFromSource"
+      )
+
     override fun create(
       action: Action,
       pathPattern: PathPattern,
@@ -96,30 +139,35 @@ internal class GrpcFeatureBinding(
       val responseAdapter = if (action.parameters.size == 2) {
         claimer.claimParameter(1)
         val responseType: Type = action.parameters[1].type.streamElementType()
-          ?: error("@Grpc function's second parameter should be a MessageSource: $action")
+          ?: error("@Grpc function's second parameter should be a MessageSink(blocking) or SendChannel(suspending): $action")
         @Suppress("UNCHECKED_CAST") // Assume it's a proto type.
         ProtoAdapter.get(responseType as Class<Any>)
       } else {
         claimer.claimReturnValue()
         @Suppress("UNCHECKED_CAST") // Assume it's a proto type.
         ProtoAdapter.get(wireAnnotation.responseAdapter) as ProtoAdapter<Any>
       }
+      val isSuspending = action.function.isSuspend
 
       return if (streamingRequestType != null) {
         @Suppress("UNCHECKED_CAST") // Assume it's a proto type.
         GrpcFeatureBinding(
           requestAdapter = ProtoAdapter.get(streamingRequestType as Class<Any>),
           responseAdapter = responseAdapter,
           streamingRequest = true,
-          streamingResponse = streamingResponse
+          streamingResponse = streamingResponse,
+          isSuspend = isSuspending,
+          grpcMessageSourceChannelContext = grpcMessageSourceChannelDispatcher,
         )
       } else {
         @Suppress("UNCHECKED_CAST") // Assume it's a proto type.
         GrpcFeatureBinding(
           requestAdapter = ProtoAdapter.get(wireAnnotation.requestAdapter) as ProtoAdapter<Any>,
           responseAdapter = responseAdapter,
           streamingRequest = false,
-          streamingResponse = streamingResponse
+          streamingResponse = streamingResponse,
+          isSuspend = isSuspending,
+          grpcMessageSourceChannelContext = grpcMessageSourceChannelDispatcher,
         )
       }
     }

misk/src/main/kotlin/misk/grpc/GrpcMessageSinkChannel.kt

@@ -0,0 +1,33 @@
+package misk.grpc
+
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.SendChannel
+import kotlinx.coroutines.channels.consumeEach
+import kotlinx.coroutines.runInterruptible
+
+/**
+ * Bridges a [GrpcMessageSink] to a [Channel].
+ *
+ * This is the primary mechanism for suspending gRPC calls to handle
+ * Server response streaming. The [GrpcMessageSinkChannel] can be passed
+ * to the gRPC Action function as a SendChannel to write responses to.
+ *
+ */
+internal class GrpcMessageSinkChannel<T : Any>(
+  private val channel: Channel<T>,
+  private val sink: GrpcMessageSink<T>,
+) : SendChannel<T> by channel {
+
+  /**
+   * Bridges the channel to the sink.
+   *
+   * This will read from the [channel] and write the messages to the [sink]
+   * until the channel is closed for sending.
+   */
+  suspend fun bridgeToSink() =
+    channel.consumeEach {
+      runInterruptible {
+        sink.write(it)
+      }
+    }
+}

misk/src/main/kotlin/misk/grpc/GrpcMessageSourceChannel.kt

@@ -0,0 +1,41 @@
+package misk.grpc
+
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.ReceiveChannel
+import kotlinx.coroutines.runInterruptible
+import kotlin.coroutines.CoroutineContext
+
+/**
+ * Bridges a [GrpcMessageSource] to a [Channel].
+ *
+ * This is the primary mechanism for suspending gRPC calls to handle
+ * Client request streaming. The [GrpcMessageSourceChannel] can be passed
+ * to the gRPC Action function as a ReceiveChannel to read requests from.
+ *
+ */
+internal class GrpcMessageSourceChannel<T : Any>(
+  private val channel: Channel<T>,
+  private val source: GrpcMessageSource<T>,
+  private val coroutineContext: CoroutineContext,
+) : ReceiveChannel<T> by channel {
+
+  /**
+   * Bridges the source to the channel.
+   *
+   * This will read from the [source] and send the messages to the [channel]
+   * until the [source] is exhausted.
+   */
+  suspend fun bridgeFromSource() {
+    try {
+      while (true) {
+        val request = runInterruptible(coroutineContext) {
+          source.read()
+        } ?: break
+        channel.send(request)
+      }
+    } finally {
+      channel.close()
+    }
+  }
+
+}