Address #66 by providing an interceptor-based API to inject CoroutineContexts. (#143)

* Revert "Add a ServerInterceptor API for injecting custom CoroutineContext elements."

This reverts commit 989f1f1.

* Add CoroutineContextServerInterceptor, an API to add elements to the CoroutineContext of the server.

Author: lowasser

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

@@ -0,0 +1,56 @@
+package io.grpc.kotlin
+
+import io.grpc.Metadata
+import io.grpc.ServerCall
+import io.grpc.ServerCallHandler
+import io.grpc.ServerInterceptor
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
+import io.grpc.Context as GrpcContext
+
+/**
+ * A [ServerInterceptor] subtype that can install elements in the [CoroutineContext] where server
+ * logic is executed.  These elements are applied "after" the
+ * [AbstractCoroutineServerImpl.context]; that is, the interceptor overrides the server's context.
+ */
+abstract class CoroutineContextServerInterceptor : ServerInterceptor {
+  companion object {
+    // This is deliberately kept visibility-restricted; it's intentional that the only way to affect
+    // the CoroutineContext is to extend CoroutineContextServerInterceptor.
+    internal val COROUTINE_CONTEXT_KEY : GrpcContext.Key<CoroutineContext> =
+      GrpcContext.keyWithDefault("grpc-kotlin-coroutine-context", EmptyCoroutineContext)
+
+    private fun GrpcContext.extendCoroutineContext(coroutineContext: CoroutineContext): GrpcContext {
+      val oldCoroutineContext: CoroutineContext = COROUTINE_CONTEXT_KEY[this]
+      val newCoroutineContext = oldCoroutineContext + coroutineContext
+      return withValue(COROUTINE_CONTEXT_KEY, newCoroutineContext)
+    }
+  }
+
+  /**
+   * Override this function to return a [CoroutineContext] in which to execute [call] and [headers].
+   * The returned [CoroutineContext] will override any corresponding context elements in the
+   * server object.
+   *
+   * This function will be called each time a [call] is executed.
+   */
+  abstract fun coroutineContext(call: ServerCall<*, *>, headers: Metadata): CoroutineContext
+
+  private inline fun <R> withGrpcContext(context: GrpcContext, action: () -> R): R {
+    val oldContext: GrpcContext = context.attach()
+    return try {
+      action()
+    } finally {
+      context.detach(oldContext)
+    }
+  }
+
+  final override fun <ReqT, RespT> interceptCall(
+    call: ServerCall<ReqT, RespT>,
+    headers: Metadata,
+    next: ServerCallHandler<ReqT, RespT>
+  ): ServerCall.Listener<ReqT> =
+    withGrpcContext(GrpcContext.current().extendCoroutineContext(coroutineContext(call, headers))) {
+      next.startCall(call, headers)
+    }
+}

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

@@ -196,7 +196,9 @@ object ServerCalls {
   ): ServerCallHandler<RequestT, ResponseT> =
     ServerCallHandler {
       call, _ -> serverCallListener(
-        context + GrpcContextElement.current(),
+        context
+          + CoroutineContextServerInterceptor.COROUTINE_CONTEXT_KEY.get()
+          + GrpcContextElement.current(),
         call,
         implementation
       )
@@ -235,9 +237,7 @@ object ServerCalls {
     }
 
     val rpcScope = CoroutineScope(context)
-    val rpcJob = rpcScope.async(
-      CoroutineName("${call.methodDescriptor.fullMethodName} implementation")
-    ) {
+    val rpcJob = rpcScope.async {
       runCatching {
         implementation(requests).collect {
           readiness.suspendUntilReady()

stub/src/test/java/io/grpc/kotlin/CoroutineContextServerInterceptorTest.kt

@@ -0,0 +1,100 @@
+package io.grpc.kotlin
+
+import com.google.common.truth.Truth.assertThat
+import io.grpc.ServerCall
+import io.grpc.ServerInterceptors
+import io.grpc.examples.helloworld.GreeterGrpcKt.GreeterCoroutineImplBase
+import io.grpc.examples.helloworld.GreeterGrpcKt.GreeterCoroutineStub
+import io.grpc.examples.helloworld.HelloReply
+import io.grpc.examples.helloworld.HelloRequest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlin.coroutines.coroutineContext
+import io.grpc.Metadata as GrpcMetadata
+
+/** Tests for [CoroutineContextServerInterceptor]. */
+@RunWith(JUnit4::class) /* inserted by Copybara: */ @com.google.testing.testsize.MediumTest
+class CoroutineContextServerInterceptorTest : AbstractCallsTest() {
+  class ArbitraryContextElement(val message: String = "") : CoroutineContext.Element {
+    companion object Key : CoroutineContext.Key<ArbitraryContextElement>
+    override val key: CoroutineContext.Key<*>
+      get() = Key
+  }
+
+  class HelloReplyWithContextMessage(
+    message: String? = null
+  ) : GreeterCoroutineImplBase(
+    message?.let { ArbitraryContextElement(it) } ?: EmptyCoroutineContext
+  ) {
+    override suspend fun sayHello(request: HelloRequest): HelloReply =
+      helloReply(coroutineContext[ArbitraryContextElement]!!.message)
+  }
+
+  @Test
+  fun injectContext() {
+    val interceptor = object : CoroutineContextServerInterceptor() {
+      override fun coroutineContext(
+        call: ServerCall<*, *>,
+        headers: GrpcMetadata
+      ): CoroutineContext = ArbitraryContextElement("success")
+    }
+
+    val channel = makeChannel(HelloReplyWithContextMessage(), interceptor)
+    val client = GreeterCoroutineStub(channel)
+
+    runBlocking {
+      assertThat(client.sayHello(helloRequest("")).message).isEqualTo("success")
+    }
+  }
+
+  @Test
+  fun conflictingInterceptorsInnermostWins() {
+    val interceptor1 = object : CoroutineContextServerInterceptor() {
+      override fun coroutineContext(
+        call: ServerCall<*, *>,
+        headers: GrpcMetadata
+      ): CoroutineContext = ArbitraryContextElement("first")
+    }
+    val interceptor2 = object : CoroutineContextServerInterceptor() {
+      override fun coroutineContext(
+        call: ServerCall<*, *>,
+        headers: GrpcMetadata
+      ): CoroutineContext = ArbitraryContextElement("second")
+    }
+
+    val channel = makeChannel(
+      ServerInterceptors.intercept(
+        ServerInterceptors.intercept(
+          HelloReplyWithContextMessage(),
+          interceptor2
+        ),
+        interceptor1
+      )
+    )
+    val client = GreeterCoroutineStub(channel)
+
+    runBlocking {
+      assertThat(client.sayHello(helloRequest("")).message).isEqualTo("second")
+    }
+  }
+
+  @Test
+  fun interceptorContextTakesPriority() {
+    val interceptor = object : CoroutineContextServerInterceptor() {
+      override fun coroutineContext(
+        call: ServerCall<*, *>,
+        headers: GrpcMetadata
+      ): CoroutineContext = ArbitraryContextElement("interceptor")
+    }
+
+    val channel = makeChannel(HelloReplyWithContextMessage("server"), interceptor)
+    val client = GreeterCoroutineStub(channel)
+
+    runBlocking {
+      assertThat(client.sayHello(helloRequest("")).message).isEqualTo("interceptor")
+    }
+  }
+}