From 91668585912a86953586bdf0bc2561d7e8e596b7 Mon Sep 17 00:00:00 2001
From: Artem Bukhonov <artem.bukhonov@jetbrains.com>
Date: Tue, 13 Jun 2023 22:38:02 +0200
Subject: [PATCH] More proper overloads with suspend handlers: pass
 CoroutineContext and CoroutineStart

---
 .../util/IRdEndpointCoroutineUtil.kt          | 26 +++++++++++++++++++
 .../rd/framework/util/ISourceCoroutineUtil.kt | 11 +++++---
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/IRdEndpointCoroutineUtil.kt b/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/IRdEndpointCoroutineUtil.kt
index c342d931d..a2285715a 100644
--- a/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/IRdEndpointCoroutineUtil.kt
+++ b/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/IRdEndpointCoroutineUtil.kt
@@ -5,8 +5,13 @@ import com.jetbrains.rd.framework.impl.RdEndpoint
 import com.jetbrains.rd.util.lifetime.Lifetime
 import com.jetbrains.rd.util.reactive.IScheduler
 import com.jetbrains.rd.util.threading.SynchronousScheduler
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.async
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 
+@Deprecated("Use the overload with CoroutineScope and CoroutineContext and pass all required context elements")
 fun <TReq, TRes> IRdEndpoint<TReq, TRes>.setSuspend(
     cancellationScheduler: IScheduler? = null,
     handlerScheduler: IScheduler? = null,
@@ -21,4 +26,25 @@ fun <TReq, TRes> IRdEndpoint<TReq, TRes>.setSuspend(
     set(cancellationScheduler, handlerScheduler) { lt, req ->
         lt.startAsync(coroutineDispatcher) { handler(lt, req) }.toRdTask()
     }
+}
+
+/**
+ * Sets suspend handler for the [IRdEndpoint].
+ *
+ * When a protocol call is occurred it starts a new coroutine on [scope] passing [coroutineContext] and [coroutineStart] to it.
+ * [cancellationScheduler] and [handlerScheduler] are passed to [IRdEndpoint.set]
+ */
+fun <TReq, TRes> IRdEndpoint<TReq, TRes>.setSuspend(
+    scope: CoroutineScope,
+    coroutineContext: CoroutineContext = EmptyCoroutineContext,
+    coroutineStart: CoroutineStart = CoroutineStart.DEFAULT,
+    cancellationScheduler: IScheduler? = null,
+    handlerScheduler: IScheduler? = null,
+    handler: suspend (Lifetime, TReq) -> TRes
+) {
+    set(cancellationScheduler, handlerScheduler) { lt, req ->
+        scope.async(coroutineContext, coroutineStart) {
+            handler(lt, req)
+        }.toRdTask()
+    }
 }
\ No newline at end of file
diff --git a/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/ISourceCoroutineUtil.kt b/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/ISourceCoroutineUtil.kt
index 9998dc727..ccd9ee4db 100644
--- a/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/ISourceCoroutineUtil.kt
+++ b/rd-kt/rd-framework/src/main/kotlin/com/jetbrains/rd/framework/util/ISourceCoroutineUtil.kt
@@ -4,9 +4,10 @@ import com.jetbrains.rd.util.lifetime.Lifetime
 import com.jetbrains.rd.util.reactive.IScheduler
 import com.jetbrains.rd.util.reactive.ISource
 import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Deferred
-import kotlinx.coroutines.launch
 import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 
 fun <T> ISource<T>.nextValueAsync(
     lifetime: Lifetime,
@@ -53,13 +54,15 @@ suspend fun <T : Any> ISource<T?>.nextNotNullValue(lifetime: Lifetime = Lifetime
         nextNotNullValueAsync(lifetime).await()
     }
 
+@Deprecated("Use overload with CoroutineContext")
 fun<T> ISource<T>.adviseSuspend(lifetime: Lifetime, scheduler: IScheduler, handler: suspend (T) -> Unit) {
-    adviseSuspend(lifetime, scheduler.asCoroutineDispatcher(allowInlining = true), handler)
+    // AIR allowInlining is a bad practice, ask @daniil for the details
+    adviseSuspend(lifetime, scheduler.asCoroutineDispatcher(allowInlining = true), CoroutineStart.UNDISPATCHED, handler)
 }
 
-fun<T> ISource<T>.adviseSuspend(lifetime: Lifetime, context: CoroutineContext, handler: suspend (T) -> Unit) {
+fun<T> ISource<T>.adviseSuspend(lifetime: Lifetime, context: CoroutineContext = EmptyCoroutineContext, coroutineStart: CoroutineStart = CoroutineStart.DEFAULT, handler: suspend (T) -> Unit) {
     advise(lifetime) {
-        lifetime.launch(context) {
+        lifetime.launch(context, coroutineStart) {
             handler(it)
         }
     }