Last pull from internal before launch, hopefully. Stop implementing CoroutineScope in AbstractCoroutineServerImpl.

Author: lowasser

compiler/src/main/java/io/grpc/kotlin/generator/GrpcCoroutineServerGenerator.kt

@@ -19,15 +19,30 @@ package io.grpc.kotlin.generator
 import com.google.common.annotations.VisibleForTesting
 import com.google.protobuf.Descriptors.MethodDescriptor
 import com.google.protobuf.Descriptors.ServiceDescriptor
-import com.squareup.kotlinpoet.*
+import com.squareup.kotlinpoet.ClassName
+import com.squareup.kotlinpoet.CodeBlock
+import com.squareup.kotlinpoet.FunSpec
+import com.squareup.kotlinpoet.KModifier
+import com.squareup.kotlinpoet.MemberName
 import com.squareup.kotlinpoet.MemberName.Companion.member
+import com.squareup.kotlinpoet.ParameterSpec
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
+import com.squareup.kotlinpoet.TypeSpec
+import com.squareup.kotlinpoet.asClassName
 import io.grpc.ServerServiceDefinition
 import io.grpc.Status
 import io.grpc.StatusException
 import io.grpc.kotlin.AbstractCoroutineServerImpl
 import io.grpc.kotlin.ServerCalls
-import io.grpc.kotlin.generator.protoc.*
+import io.grpc.kotlin.generator.protoc.Declarations
+import io.grpc.kotlin.generator.protoc.GeneratorConfig
+import io.grpc.kotlin.generator.protoc.MemberSimpleName
+import io.grpc.kotlin.generator.protoc.builder
+import io.grpc.kotlin.generator.protoc.classBuilder
+import io.grpc.kotlin.generator.protoc.declarations
+import io.grpc.kotlin.generator.protoc.methodName
+import io.grpc.kotlin.generator.protoc.of
+import io.grpc.kotlin.generator.protoc.serviceName
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.flow.Flow
 import kotlin.coroutines.CoroutineContext
@@ -78,7 +93,7 @@ class GrpcCoroutineServerGenerator(config: GeneratorConfig): ServiceCodeGenerato
       .addModifiers(KModifier.ABSTRACT)
       .addKdoc(
         """
-        Skeletal implementation of the %L service based on Kotlin coroutines.
+        Skeletal implementation of the %L service based on Kotlin coroutines. 
         """.trimIndent(),
         service.fullName
       )
@@ -167,7 +182,7 @@ class GrpcCoroutineServerGenerator(config: GeneratorConfig): ServiceCodeGenerato
       CodeBlock.of(
         """
           %M(
-            scope = this,
+            context = this.context,
             descriptor = %L,
             implementation = ::%N
           )
@@ -204,7 +219,7 @@ class GrpcCoroutineServerGenerator(config: GeneratorConfig): ServiceCodeGenerato
           If creating or collecting the returned flow fails with a [%statusException:T], the RPC
           will fail with the corresponding [%status:T].  If it fails with a
           [%cancellationException:T], the RPC will fail with status `Status.CANCELLED`.  If creating
-          or collecting the returned flow fails for any other reason, the RPC will fail with
+          or collecting the returned flow fails for any other reason, the RPC will fail with 
           `Status.UNKNOWN` with the exception as a cause.
         """.trimIndent()
       )
@@ -241,4 +256,4 @@ class GrpcCoroutineServerGenerator(config: GeneratorConfig): ServiceCodeGenerato
       .addNamed(kDocSections.joinToString("\n\n"), kDocBindings)
       .build()
   }
-}
+}

compiler/src/test/java/io/grpc/kotlin/generator/CoroutineServerImplGeneratorTest.kt

@@ -64,7 +64,7 @@ class CoroutineServerImplGeneratorTest {
     """.trimIndent())
     assertThat(stub.serverMethodDef.toString()).isEqualTo("""
       io.grpc.kotlin.ServerCalls.unaryServerMethodDefinition(
-        scope = this,
+        context = this.context,
         descriptor = io.grpc.examples.helloworld.GreeterGrpc.getSayHelloMethod(),
         implementation = ::sayHello
       )
@@ -91,7 +91,7 @@ class CoroutineServerImplGeneratorTest {
     """.trimIndent())
     assertThat(stub.serverMethodDef.toString()).isEqualTo("""
       io.grpc.kotlin.ServerCalls.clientStreamingServerMethodDefinition(
-        scope = this,
+        context = this.context,
         descriptor = io.grpc.examples.helloworld.GreeterGrpc.getClientStreamSayHelloMethod(),
         implementation = ::clientStreamSayHello
       )
@@ -117,7 +117,7 @@ class CoroutineServerImplGeneratorTest {
     """.trimIndent())
     assertThat(stub.serverMethodDef.toString()).isEqualTo("""
       io.grpc.kotlin.ServerCalls.serverStreamingServerMethodDefinition(
-        scope = this,
+        context = this.context,
         descriptor = io.grpc.examples.helloworld.GreeterGrpc.getServerStreamSayHelloMethod(),
         implementation = ::serverStreamSayHello
       )
@@ -145,7 +145,7 @@ class CoroutineServerImplGeneratorTest {
     """.trimIndent())
     assertThat(stub.serverMethodDef.toString()).isEqualTo("""
       io.grpc.kotlin.ServerCalls.bidiStreamingServerMethodDefinition(
-        scope = this,
+        context = this.context,
         descriptor = io.grpc.examples.helloworld.GreeterGrpc.getBidiStreamSayHelloMethod(),
         implementation = ::bidiStreamSayHello
       )

compiler/src/test/java/io/grpc/kotlin/generator/GreeterCoroutineImplBase.expected

@@ -20,9 +20,8 @@ import kotlinx.coroutines.flow.Flow
 /**
  * Skeletal implementation of the helloworld.Greeter service based on Kotlin coroutines.
  */
-abstract class GreeterCoroutineImplBase(
-  coroutineContext: CoroutineContext = EmptyCoroutineContext
-) : AbstractCoroutineServerImpl(coroutineContext) {
+abstract class GreeterCoroutineImplBase(coroutineContext: CoroutineContext = EmptyCoroutineContext)
+    : AbstractCoroutineServerImpl(coroutineContext) {
   /**
    * Returns the response to an RPC for helloworld.Greeter.SayHello.
    *
@@ -34,8 +33,10 @@ abstract class GreeterCoroutineImplBase(
    *
    * @param request The request from the client.
    */
-  open suspend fun sayHello(request: HelloRequest): HelloReply = throw
-      StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.SayHello is unimplemented"))
+  open suspend fun sayHello(request: HelloRequest): HelloReply {
+    throw
+        StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.SayHello is unimplemented"))
+  }
 
   /**
    * Returns the response to an RPC for helloworld.Greeter.ClientStreamSayHello.
@@ -50,8 +51,10 @@ abstract class GreeterCoroutineImplBase(
    *        collected only once and throws [java.lang.IllegalStateException] on attempts to collect
    *        it more than once.
    */
-  open suspend fun clientStreamSayHello(requests: Flow<HelloRequest>): HelloReply = throw
-      StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.ClientStreamSayHello is unimplemented"))
+  open suspend fun clientStreamSayHello(requests: Flow<HelloRequest>): HelloReply {
+    throw
+        StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.ClientStreamSayHello is unimplemented"))
+  }
 
   /**
    * Returns a [Flow] of responses to an RPC for helloworld.Greeter.ServerStreamSayHello.
@@ -65,8 +68,10 @@ abstract class GreeterCoroutineImplBase(
    *
    * @param request The request from the client.
    */
-  open fun serverStreamSayHello(request: MultiHelloRequest): Flow<HelloReply> = throw
-      StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.ServerStreamSayHello is unimplemented"))
+  open fun serverStreamSayHello(request: MultiHelloRequest): Flow<HelloReply> {
+    throw
+        StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.ServerStreamSayHello is unimplemented"))
+  }
 
   /**
    * Returns a [Flow] of responses to an RPC for helloworld.Greeter.BidiStreamSayHello.
@@ -82,27 +87,29 @@ abstract class GreeterCoroutineImplBase(
    *        collected only once and throws [java.lang.IllegalStateException] on attempts to collect
    *        it more than once.
    */
-  open fun bidiStreamSayHello(requests: Flow<HelloRequest>): Flow<HelloReply> = throw
-      StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.BidiStreamSayHello is unimplemented"))
+  open fun bidiStreamSayHello(requests: Flow<HelloRequest>): Flow<HelloReply> {
+    throw
+        StatusException(UNIMPLEMENTED.withDescription("Method helloworld.Greeter.BidiStreamSayHello is unimplemented"))
+  }
 
   final override fun bindService(): ServerServiceDefinition = builder(getServiceDescriptor())
     .addMethod(unaryServerMethodDefinition(
-    scope = this,
+    context = this.context,
     descriptor = GreeterGrpc.getSayHelloMethod(),
     implementation = ::sayHello
   ))
     .addMethod(clientStreamingServerMethodDefinition(
-    scope = this,
+    context = this.context,
     descriptor = GreeterGrpc.getClientStreamSayHelloMethod(),
     implementation = ::clientStreamSayHello
   ))
     .addMethod(serverStreamingServerMethodDefinition(
-    scope = this,
+    context = this.context,
     descriptor = GreeterGrpc.getServerStreamSayHelloMethod(),
     implementation = ::serverStreamSayHello
   ))
     .addMethod(bidiStreamingServerMethodDefinition(
-    scope = this,
+    context = this.context,
     descriptor = GreeterGrpc.getBidiStreamSayHelloMethod(),
     implementation = ::bidiStreamSayHello
   )).build()

stub/src/main/java/io/grpc/kotlin/AbstractCoroutineServerImpl.kt

@@ -17,24 +17,23 @@
 package io.grpc.kotlin
 
 import io.grpc.BindableService
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.SupervisorJob
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 
 /**
  * Skeleton implementation of a coroutine-based gRPC server implementation.  Intended to be
  * subclassed by generated code.
  */
-abstract class AbstractCoroutineServerImpl private constructor(
-  private val delegateScope: CoroutineScope
-) : CoroutineScope, BindableService {
-
-  constructor(coroutineContext: CoroutineContext = EmptyCoroutineContext) :
-    this(CoroutineScope(coroutineContext + SupervisorJob(coroutineContext[Job])))
-  // We want a SupervisorJob so one failed RPC does not bring down the entire server.
-
-  final override val coroutineContext: CoroutineContext
-    get() = delegateScope.coroutineContext
-}
+abstract class AbstractCoroutineServerImpl(
+  /** The context in which to run server coroutines. */
+  val context: CoroutineContext = EmptyCoroutineContext
+) : BindableService {
+  /*
+   * Each RPC is executed in its own coroutine scope built from [context].  We could have a parent
+   * scope, but it doesn't really add anything: we don't want users to be able to launch tasks
+   * in that scope easily, since almost all coroutines should be scoped to the RPC and cancelled
+   * if the RPC is cancelled.  Users who don't want that behavior should manage their own scope for
+   * it.  Additionally, gRPC server objects don't have their own notion of shutdown: shutting down
+   * a server means cancelling the RPCs, not calling a teardown on the server object.
+   */
+}

stub/src/main/java/io/grpc/kotlin/ClientCalls.kt

@@ -109,15 +109,13 @@ object ClientCalls {
     headers: suspend () -> GrpcMetadata = { GrpcMetadata() }
   ): (RequestT) -> Flow<ResponseT> = {
     flow {
-      emitAll(
-        serverStreamingRpc(
-          channel,
-          method,
-          it,
-          callOptions,
-          headers()
-        )
-      )
+      serverStreamingRpc(
+        channel,
+        method,
+        it,
+        callOptions,
+        headers()
+      ).collect { emit(it) }
     }
   }
 
@@ -207,15 +205,13 @@ object ClientCalls {
   ): (Flow<RequestT>) -> Flow<ResponseT> =
     {
       flow {
-        emitAll(
-          bidiStreamingRpc(
-            channel,
-            method,
-            it,
-            callOptions,
-            headers()
-          )
-        )
+        bidiStreamingRpc(
+          channel,
+          method,
+          it,
+          callOptions,
+          headers()
+        ).collect { emit(it) }
       }
     }
 

stub/src/main/java/io/grpc/kotlin/GrpcContextElement.kt

@@ -16,11 +16,9 @@
 
 package io.grpc.kotlin
 
-import kotlinx.coroutines.CoroutineScope
-import io.grpc.Context as GrpcContext
-import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.ThreadContextElement
-import kotlinx.coroutines.withContext
+import kotlin.coroutines.CoroutineContext
+import io.grpc.Context as GrpcContext
 
 /**
  * A [CoroutineContext] that propagates an associated [io.grpc.Context] to coroutines run using

stub/src/main/java/io/grpc/kotlin/Helpers.kt

@@ -23,15 +23,16 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.FlowCollector
 import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.single
 import kotlinx.coroutines.runBlocking
 
 /**
  * Extracts the value of a [Deferred] known to be completed, or throws its exception if it was
  * not completed successfully.  (Non-experimental variant of `getDone`.)
  */
-internal val <T> Deferred<T>.doneValue: T
+val <T> Deferred<T>.doneValue: T
   get() {
     check(isCompleted) { "doneValue should only be called on completed Deferred values" }
     return runBlocking(Dispatchers.Unconfined) {
@@ -42,52 +43,44 @@ internal val <T> Deferred<T>.doneValue: T
 /**
  * Cancels a [Job] with a cause and suspends until the job completes/is finished cancelling.
  */
-internal suspend fun Job.cancelAndJoin(message: String, cause: Exception? = null) {
+suspend fun Job.cancelAndJoin(message: String, cause: Exception? = null) {
   cancel(message, cause)
   join()
 }
 
-internal suspend fun <T> FlowCollector<T>.emitAll(flow: Flow<T>) {
-  flow.collect { emit(it) }
-}
-
 /**
- * Extracts the one and only element of this flow, throwing an appropriate [StatusException] if
- * there is not exactly one element.  (Otherwise this is fully equivalent to `Flow.single()`.)
+ * Returns this flow, save that if there is not exactly one element, it throws a [StatusException].
+ *
+ * The purpose of this function is to enable the one element to get processed before we have
+ * confirmation that the input flow is done.
  */
-internal suspend fun <T> Flow<T>.singleOrStatus(
+fun <T> Flow<T>.singleOrStatusFlow(
   expected: String,
   descriptor: Any
-): T {
-  // We could call Flow.single() and catch exceptions, but if the underlying flow throws
-  // IllegalStateException or NoSuchElementException for its own reasons, we'd swallow those
-  // and give misleading errors.  Instead, we reimplement single() ourselves.
-  var result: T? = null
+): Flow<T> = flow {
   var found = false
   collect {
     if (!found) {
       found = true
-      result = it
+      emit(it)
     } else {
       throw StatusException(
         Status.INTERNAL.withDescription("Expected one $expected for $descriptor but received two")
       )
     }
   }
-  @Suppress("UNCHECKED_CAST")
   if (!found) {
     throw StatusException(
       Status.INTERNAL.withDescription("Expected one $expected for $descriptor but received none")
     )
-  } else {
-    return result as T
   }
 }
 
-/** Runs [block] and returns any exception it throws, or `null` if it does not throw. */
-internal inline fun thrownOrNull(block: () -> Unit): Throwable? = try {
-  block()
-  null
-} catch (thrown: Throwable) {
-  thrown
-}
+/**
+ * Returns the one and only element of this flow, and throws a [StatusException] if there is not
+ * exactly one element.
+ */
+suspend fun <T> Flow<T>.singleOrStatus(
+  expected: String,
+  descriptor: Any
+): T = singleOrStatusFlow(expected, descriptor).single()