From 069e51f2ffa6aaae85afbd910c3a6694176fbbc6 Mon Sep 17 00:00:00 2001
From: InsanusMokrassar <ovsyannikov.alexey95@gmail.com>
Date: Tue, 22 Dec 2020 15:16:06 +0600
Subject: [PATCH] add defaultSafelyExceptionHandler and SafelyExceptionHandler

---
 CHANGELOG.md                                  | 10 +++++--
 .../micro_utils/coroutines/HandleSafely.kt    | 30 +++++++++++++++++--
 .../ktor/client/FlowsWebsockets.kt            |  2 +-
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 23da22e99ba..b185780e18f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,9 +3,13 @@
 ## 0.4.14
 
 * `Versions`:
-      * `Kotlin`: `1.4.20` -> `1.4.21`
-      * `Ktor`: `1.4.3` -> `1.5.0`
-      * `Klock`: `2.0.1` -> `2.0.2`
+    * `Kotlin`: `1.4.20` -> `1.4.21`
+    * `Ktor`: `1.4.3` -> `1.5.0`
+    * `Klock`: `2.0.1` -> `2.0.2`
+* `Coroutines`:
+    * Add global variable `defaultSafelyExceptionHandler`
+    * Add `SafelyExceptionHandlerKey` and `SafelyExceptionHandler` classes to be able to overwrite
+    `defaultSafelyExceptionHandler` using context of coroutine
 
 ## 0.4.13
 
diff --git a/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/HandleSafely.kt b/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/HandleSafely.kt
index 64d3dc86c1d..91daf1de16c 100644
--- a/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/HandleSafely.kt
+++ b/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/HandleSafely.kt
@@ -2,23 +2,47 @@ package dev.inmo.micro_utils.coroutines
 
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.supervisorScope
-
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.coroutineContext
 
 typealias ExceptionHandler<T> = suspend (Throwable) -> T
+
+/**
+ * This instance will be used in all calls of [safely] where exception handler has not been passed
+ */
+var defaultSafelyExceptionHandler: ExceptionHandler<Nothing> = { throw it }
+
+class SafelyExceptionHandlerKey<T> : CoroutineContext.Key<SafelyExceptionHandler<T>>
+inline fun <T> safelyExceptionHandlerKey() = SafelyExceptionHandlerKey<T>()
+class SafelyExceptionHandler<T>(
+    val handler: ExceptionHandler<T>
+) : CoroutineContext.Element {
+
+    override val key: CoroutineContext.Key<*> = safelyExceptionHandlerKey<T>()
+}
+
 /**
  * It will run [block] inside of [supervisorScope] to avoid problems with catching of exceptions
  *
  * @param [onException] Will be called when happen exception inside of [block]. By default will throw exception - this
  * exception will be available for catching
+ *
+ * @see defaultSafelyExceptionHandler
+ * @see safelyWithoutExceptions
  */
 suspend inline fun <T> safely(
-    noinline onException: ExceptionHandler<T> = { throw it },
+    noinline onException: ExceptionHandler<T> = defaultSafelyExceptionHandler,
     noinline block: suspend CoroutineScope.() -> T
 ): T {
     return try {
         supervisorScope(block)
     } catch (e: Throwable) {
-        onException(e)
+        val handler = if (onException == defaultSafelyExceptionHandler) {
+            coroutineContext[safelyExceptionHandlerKey<T>()] ?.handler ?: onException
+        } else {
+            onException
+        }
+        handler(e)
     }
 }
 
diff --git a/ktor/client/src/commonMain/kotlin/dev/inmo/micro_utils/ktor/client/FlowsWebsockets.kt b/ktor/client/src/commonMain/kotlin/dev/inmo/micro_utils/ktor/client/FlowsWebsockets.kt
index 0a8243f1598..567d87e5b01 100644
--- a/ktor/client/src/commonMain/kotlin/dev/inmo/micro_utils/ktor/client/FlowsWebsockets.kt
+++ b/ktor/client/src/commonMain/kotlin/dev/inmo/micro_utils/ktor/client/FlowsWebsockets.kt
@@ -25,7 +25,7 @@ inline fun <T> HttpClient.createStandardWebsocketFlow(
         val producerScope = this@channelFlow
         do {
             val reconnect = try {
-                safely ({ throw it }) {
+                safely {
                     ws(correctedUrl) {
                         for (received in incoming) {
                             when (received) {