grpc-native: Refactor to use callbacks instead of coroutines

Signed-off-by: Johannes Zottele <[email protected]>

Author: Jozott00

grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/ManagedChannel.native.kt

@@ -58,16 +58,15 @@ internal class NativeManagedChannel(
     target: String,
     // we must store them, otherwise the credentials are getting released
     credentials: GrpcCredentials,
-    dispatcher: CoroutineDispatcher = Dispatchers.Default,
 ) : ManagedChannel, ManagedChannelPlatform() {
 
-    // A reference to make sure the grpc_init() was called. (it is released after shutdown)
+    // a reference to make sure the grpc_init() was called. (it is released after shutdown)
     @Suppress("unused")
     private val rt = GrpcRuntime.acquire()
 
-    private val channelJob = SupervisorJob()
-    private val callJobSupervisor = SupervisorJob(channelJob)
-    private val channelScope = CoroutineScope(channelJob + dispatcher)
+    // job bundling all the call jobs created by this channel.
+    // this allows easy cancellation of ongoing calls.
+    private val callJobSupervisor = SupervisorJob()
 
     // the channel's completion queue, handling all request operations
     private val cq = CompletionQueue()
@@ -111,23 +110,15 @@ internal class NativeManagedChannel(
             return
         }
         if (force) {
+            // TODO: replace jobs by custom pendingCallClass.
             callJobSupervisor.cancelChildren(CancellationException("Channel is shutting down"))
         }
-        // prevent any start() calls on already call jobs
-        callJobSupervisor.children.forEach {
-            (it as CompletableJob).complete()
-        }
-        callJobSupervisor.complete()
-        channelScope.launch {
-            withContext(NonCancellable) {
-                // wait for all child jobs to complete.
-                callJobSupervisor.join()
-                // wait for the completion queue to shut down.
-                cq.shutdown(force).await()
-                if (isTerminatedInternal.complete(Unit)) {
-                    // release the grpc runtime, so it might call grpc_shutdown()
-                    rt.close()
-                }
+
+        // wait for the completion queue to shut down.
+        cq.shutdown(force).onComplete {
+            if (isTerminatedInternal.complete(Unit)) {
+                // release the grpc runtime, so it might call grpc_shutdown()
+                rt.close()
             }
         }
     }

grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/CallbackFuture.kt

@@ -6,24 +6,44 @@ package kotlinx.rpc.grpc.internal
 
 import kotlinx.atomicfu.atomic
 
+/**
+ * Thread safe future for callbacks.
+ */
 internal class CallbackFuture<T : Any> {
-    private val value = atomic<T?>(null)
-    private val callback = atomic<((T) -> Unit)?>(null)
+    private sealed interface State<out T> {
+        data class Pending<T>(val cbs: List<(T) -> Unit> = emptyList()) : State<T>
+        data class Done<T>(val value: T) : State<T>
+    }
+
+    private val state = atomic<State<T>>(State.Pending())
 
     fun complete(result: T) {
-        if (value.compareAndSet(null, result)) {
-            callback.getAndSet(null)?.invoke(result)
-        } else {
-            error("Already completed")
+        var toInvoke: List<(T) -> Unit>
+        while (true) {
+            when (val s = state.value) {
+                is State.Pending -> if (state.compareAndSet(s, State.Done(result))) {
+                    toInvoke = s.cbs
+                    break
+                }
+
+                is State.Done -> error("Already completed")
+            }
         }
+        for (cb in toInvoke) cb(result)
     }
 
-    fun onComplete(cb: (T) -> Unit) {
-        val r = value.value
-        if (r != null) cb(r)
-        else if (!callback.compareAndSet(null, cb)) {
-            // Already someone registered → run immediately
-            value.value?.let(cb)
+    fun onComplete(callback: (T) -> Unit) {
+        while (true) {
+            when (val s = state.value) {
+                is State.Done -> {
+                    callback(s.value); return
+                }
+
+                is State.Pending -> {
+                    val next = State.Pending(s.cbs + callback) // copy-on-write append
+                    if (state.compareAndSet(s, next)) return
+                }
+            }
         }
     }
 }

grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/CompletionQueue.kt

@@ -6,22 +6,19 @@
 
 package kotlinx.rpc.grpc.internal
 
-import cnames.structs.grpc_call
 import kotlinx.atomicfu.atomic
 import kotlinx.atomicfu.locks.SynchronizedObject
 import kotlinx.atomicfu.locks.synchronized
 import kotlinx.cinterop.*
-import kotlinx.coroutines.CompletableDeferred
 import libgrpcpp_c.*
 import platform.posix.memset
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.native.ref.createCleaner
 
 internal sealed interface BatchResult {
-    object Success : BatchResult
-    object ResultError : BatchResult
     object CQShutdown : BatchResult
     data class CallError(val error: grpc_call_error) : BatchResult
+    data class Called(val future: CallbackFuture<Boolean>) : BatchResult
 }
 
 /**
@@ -45,7 +42,7 @@ internal class CompletionQueue {
     // internal as it must be accessible from the SHUTDOWN_CB,
     // but it shouldn't be used from outside this file.
     @Suppress("PropertyName")
-    internal val _shutdownDone = CompletableDeferred<Unit>()
+    internal val _shutdownDone = CallbackFuture<Unit>()
 
     // used for spinning lock. false means not used (available)
     private val batchStartGuard = SynchronizedObject()
@@ -72,40 +69,41 @@ internal class CompletionQueue {
         require(kgrpc_iomgr_run_in_background()) { "The gRPC iomgr is not running background threads, required for callback based APIs." }
     }
 
-    // TODO: Remove this method
-    fun runBatch(call: NativeClientCall<*, *>, ops: CPointer<grpc_op>, nOps: ULong) =
-        runBatch(call.raw, ops, nOps)
-
-    fun runBatch(call: CPointer<grpc_call>, ops: CPointer<grpc_op>, nOps: ULong): CompletableDeferred<BatchResult> {
-        val completion = CompletableDeferred<BatchResult>()
+    fun runBatch(call: NativeClientCall<*, *>, ops: CPointer<grpc_op>, nOps: ULong): BatchResult {
+        val completion = CallbackFuture<Boolean>()
         val tag = newCbTag(completion, OPS_COMPLETE_CB)
 
         var err = grpc_call_error.GRPC_CALL_ERROR
+
         synchronized(batchStartGuard) {
-            // synchronizes access to grpc_call_start_batch
+            if (_state.value == State.SHUTTING_DOWN && ops.pointed.op == GRPC_OP_RECV_STATUS_ON_CLIENT) {
+                // if the queue is in the process of a SHUTDOWN,
+                // new call status receive batches will be rejected.
+                deleteCbTag(tag)
+                return BatchResult.CQShutdown
+            }
+
             if (forceShutdown || _state.value == State.CLOSED) {
                 // if the queue is either closed or in the process of a FORCE shutdown,
                 // new batches will instantly fail.
                 deleteCbTag(tag)
-                completion.complete(BatchResult.CQShutdown)
-                return completion
+                return BatchResult.CQShutdown
             }
 
-            err = grpc_call_start_batch(call, ops, nOps, tag, null)
+            err = grpc_call_start_batch(call.raw, ops, nOps, tag, null)
         }
 
         if (err != grpc_call_error.GRPC_CALL_OK) {
             // if the call was not successful, the callback will not be invoked.
             deleteCbTag(tag)
-            completion.complete(BatchResult.CallError(err))
-            return completion
+            return BatchResult.CallError(err)
         }
 
-        return completion
+        return BatchResult.Called(completion)
     }
 
     // must not be canceled as it cleans resources and sets the state to CLOSED
-    fun shutdown(force: Boolean = false): CompletableDeferred<Unit> {
+    fun shutdown(force: Boolean = false): CallbackFuture<Unit> {
         if (force) {
             forceShutdown = true
         }
@@ -130,10 +128,9 @@ internal class CompletionQueue {
 @CName("kq_ops_complete_cb")
 private fun opsCompleteCb(functor: CPointer<grpc_completion_queue_functor>?, ok: Int) {
     val tag = functor!!.reinterpret<grpc_cb_tag>()
-    val cont = tag.pointed.user_data!!.asStableRef<CompletableDeferred<BatchResult>>().get()
+    val cont = tag.pointed.user_data!!.asStableRef<CallbackFuture<Boolean>>().get()
     deleteCbTag(tag)
-    if (ok != 0) cont.complete(BatchResult.Success)
-    else cont.complete(BatchResult.ResultError)
+    cont.complete(ok != 0)
 }
 
 @CName("kq_shutdown_cb")