From c958f8afc35f44a7b228a1f8ae183d5021f922e4 Mon Sep 17 00:00:00 2001
From: Sergey Chelombitko <sergey.chelombitko@team.bumble.com>
Date: Fri, 20 Dec 2024 13:42:56 +0000
Subject: [PATCH] Handle cancellation in withRetry

---
 .../com/malinskiy/marathon/execution/Retry.kt | 24 ++++++++++++-------
 .../marathon/execution/device/DeviceActor.kt  | 22 ++++++++---------
 .../marathon/android/AndroidAppInstaller.kt   |  6 +++--
 3 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/core/src/main/kotlin/com/malinskiy/marathon/execution/Retry.kt b/core/src/main/kotlin/com/malinskiy/marathon/execution/Retry.kt
index 299b58db1..2d0c76ded 100644
--- a/core/src/main/kotlin/com/malinskiy/marathon/execution/Retry.kt
+++ b/core/src/main/kotlin/com/malinskiy/marathon/execution/Retry.kt
@@ -1,18 +1,24 @@
 package com.malinskiy.marathon.execution
 
-import kotlinx.coroutines.delay
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.ensureActive
+import kotlinx.coroutines.isActive
+import kotlinx.coroutines.time.delay
+import java.time.Duration
+
+suspend fun withRetry(maxAttempts: Int, retryDelay: Duration, block: suspend () -> Unit) {
+    check(maxAttempts >= 1) { "maxAttempts must be >= 1" }
 
-@Suppress("TooGenericExceptionCaught")
-suspend fun <T> withRetry(attempts: Int, delayTime: Long = 0, f: suspend () -> T): T {
     var attempt = 1
-    while (true) {
+    while (currentCoroutineContext().isActive) {
         try {
-            return f()
-        } catch (th: Throwable) {
-            if (attempt == attempts) {
-                throw th
+            return block()
+        } catch (@Suppress("TooGenericExceptionCaught") e: Throwable) {
+            currentCoroutineContext().ensureActive()
+            if (attempt == maxAttempts) {
+                throw e
             } else {
-                delay(delayTime)
+                delay(retryDelay)
             }
         }
         ++attempt
diff --git a/core/src/main/kotlin/com/malinskiy/marathon/execution/device/DeviceActor.kt b/core/src/main/kotlin/com/malinskiy/marathon/execution/device/DeviceActor.kt
index 711df6628..adeec87b3 100644
--- a/core/src/main/kotlin/com/malinskiy/marathon/execution/device/DeviceActor.kt
+++ b/core/src/main/kotlin/com/malinskiy/marathon/execution/device/DeviceActor.kt
@@ -21,8 +21,8 @@ import kotlinx.coroutines.CompletionHandler
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.async
 import kotlinx.coroutines.channels.SendChannel
-import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
+import java.time.Duration
 import java.time.Instant
 import kotlin.coroutines.CoroutineContext
 
@@ -157,18 +157,16 @@ class DeviceActor(
 
     private fun initialize() {
         logger.debug("[{}] Initializing", device.serialNumber)
-        job = scope.async {
+        job = scope.launch {
             try {
-                withRetry(30, 10000) {
-                    if (isActive) {
-                        try {
-                            device.prepare(configuration)
-                        } catch (e: CancellationException) {
-                            throw e
-                        } catch (@Suppress("TooGenericExceptionCaught") e: Exception) {
-                            logger.debug("[{}] Initialization failed. Retrying", device.serialNumber, e)
-                            throw e
-                        }
+                withRetry(maxAttempts = 30, retryDelay = Duration.ofSeconds(10)) {
+                    try {
+                        device.prepare(configuration)
+                    } catch (e: CancellationException) {
+                        throw e
+                    } catch (@Suppress("TooGenericExceptionCaught") e: Exception) {
+                        logger.debug("[{}] Initialization failed. Retrying", device.serialNumber, e)
+                        throw e
                     }
                 }
                 state.transition(DeviceEvent.Complete)
diff --git a/vendor/vendor-android/ddmlib/src/main/kotlin/com/malinskiy/marathon/android/AndroidAppInstaller.kt b/vendor/vendor-android/ddmlib/src/main/kotlin/com/malinskiy/marathon/android/AndroidAppInstaller.kt
index 2b67e4085..b42b81e1e 100644
--- a/vendor/vendor-android/ddmlib/src/main/kotlin/com/malinskiy/marathon/android/AndroidAppInstaller.kt
+++ b/vendor/vendor-android/ddmlib/src/main/kotlin/com/malinskiy/marathon/android/AndroidAppInstaller.kt
@@ -7,6 +7,7 @@ import com.malinskiy.marathon.execution.withRetry
 import com.malinskiy.marathon.io.FileHasher
 import com.malinskiy.marathon.log.MarathonLogging
 import java.io.File
+import java.time.Duration
 import java.time.Instant
 import kotlin.system.measureTimeMillis
 
@@ -39,7 +40,7 @@ class AndroidAppInstaller(
 
     @Suppress("TooGenericExceptionThrown")
     private suspend fun ensureInstalled(device: AndroidDevice, appPackage: String, appApk: File) {
-        withRetry(attempts = MAX_RETIRES, delayTime = 1000) {
+        withRetry(maxAttempts = MAX_INSTALLATION_ATTEMPTS, retryDelay = INSTALLATION_RETRY_DELAY) {
             try {
                 val checkStarted = Instant.now()
                 val fileHash = fileHasher.getHash(appApk)
@@ -138,7 +139,8 @@ class AndroidAppInstaller(
     }
 
     companion object {
-        private const val MAX_RETIRES = 3
+        private const val MAX_INSTALLATION_ATTEMPTS = 3
+        private val INSTALLATION_RETRY_DELAY = Duration.ofSeconds(1)
         private const val MARSHMALLOW_VERSION_CODE = 23
         private const val MD5_HASH_SIZE = 32
         private const val INSTALLED_TEST_APPS_SCRIPT = "pm list packages -3 | grep -E '\\.test\$' | tr -d '\\r' | cut -d ':' -f 2"