grpc: Fix race condition bug

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

Author: Jozott00

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

@@ -46,4 +46,6 @@ internal class CallbackFuture<T : Any> {
             }
         }
     }
+
+    val isCompleted: Boolean get() = state.value is State.Done
 }

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

@@ -119,6 +119,8 @@ internal class CompletionQueue {
      * See [BatchResult] for possible outcomes.
      */
     fun runBatch(call: CPointer<grpc_call>, ops: CPointer<grpc_op>, nOps: ULong): BatchResult {
+        if (_shutdownDone.isCompleted) return BatchResult.CQShutdown
+
         val completion = CallbackFuture<Boolean>()
         val tag = newCbTag(completion, OPS_COMPLETE_CB)
 
@@ -194,8 +196,8 @@ private fun opsCompleteCb(functor: CPointer<grpc_completion_queue_functor>?, ok:
 private fun shutdownCb(functor: CPointer<grpc_completion_queue_functor>?, ok: Int) {
     val tag = functor!!.reinterpret<kgrpc_cb_tag>()
     val cq = tag.pointed.user_data!!.asStableRef<CompletionQueue>().get()
-    cq._shutdownDone.complete(Unit)
     cq._state.value = CompletionQueue.State.CLOSED
+    cq._shutdownDone.complete(Unit)
     grpc_completion_queue_destroy(cq.raw)
 }
 

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

@@ -67,6 +67,7 @@ internal class NativeServerCall<Request, Response>(
     private val callbackMutex = ReentrantLock()
     private var initialized = false
     private var cancelled = false
+    private var closed = false
     private val finalized = atomic(false)
 
     // tracks whether the initial metadata has been sent.
@@ -143,6 +144,7 @@ internal class NativeServerCall<Request, Response>(
     }
 
     fun cancel(status: grpc_status_code, message: String) {
+        cancelled = true
         grpc_call_cancel_with_status(raw, status, message, null)
     }
 
@@ -168,6 +170,9 @@ internal class NativeServerCall<Request, Response>(
         cleanup: () -> Unit = {},
         onSuccess: () -> Unit = {},
     ) {
+        // if we are already closed, we cannot run any more batches.
+        if (closed || cancelled) return cleanup()
+
         when (val result = cq.runBatch(raw, ops, nOps)) {
             is BatchResult.Submitted -> {
                 result.future.onComplete {
@@ -286,6 +291,8 @@ internal class NativeServerCall<Request, Response>(
 
     override fun close(status: Status, trailers: GrpcMetadata) {
         check(initialized) { internalError("Call not initialized") }
+        closed = true
+
         val arena = Arena()
 
         val details = status.getDescription()?.toGrpcSlice()
@@ -327,7 +334,7 @@ internal class NativeServerCall<Request, Response>(
     }
 
 
-    private fun <T> tryRun(block: () -> T): T {
+    private inline fun <T> tryRun(crossinline block: () -> T): T {
         try {
             return block()
         } catch (e: Throwable) {