added retryRaise and retryEither functions (#3373)

arrow-kt · Feb 16, 2024 · af3814b · af3814b
1 parent d42429f
commit af3814b
Show file tree

Hide file tree

Showing 3 changed files with 114 additions and 2 deletions.
diff --git a/arrow-libs/resilience/arrow-resilience/api/arrow-resilience.api b/arrow-libs/resilience/arrow-resilience/api/arrow-resilience.api
@@ -252,8 +252,11 @@ public final class arrow/resilience/Schedule$Decision$Done : arrow/resilience/Sc
 
 public final class arrow/resilience/ScheduleKt {
 	public static final fun retry-4AuOtiA (Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun retry-YL6hcnA (Larrow/core/raise/Raise;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun retryEither-4AuOtiA (Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function0;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun retryOrElse-aZo8_V4 (Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function3;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun retryOrElseEither-aZo8_V4 (Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function3;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun retryRaise-4AuOtiA (Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
 public final class arrow/resilience/common/Platform : java/lang/Enum {

diff --git a/arrow-libs/resilience/arrow-resilience/src/commonMain/kotlin/arrow/resilience/Schedule.kt b/arrow-libs/resilience/arrow-resilience/src/commonMain/kotlin/arrow/resilience/Schedule.kt
@@ -1,3 +1,5 @@
+@file:OptIn(ExperimentalTypeInference::class)
+
 package arrow.resilience
 
 import arrow.core.Either
@@ -7,6 +9,9 @@ import arrow.core.identity
 import arrow.core.left
 import arrow.core.merge
 import arrow.core.nonFatalOrThrow
+import arrow.core.raise.Raise
+import arrow.core.raise.either
+import arrow.core.raise.fold
 import arrow.core.right
 import arrow.core.some
 import arrow.resilience.Schedule.Companion.identity
@@ -17,6 +22,7 @@ import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.retry
+import kotlin.experimental.ExperimentalTypeInference
 import kotlin.jvm.JvmInline
 import kotlin.math.pow
 import kotlin.random.Random
@@ -444,3 +450,54 @@ public suspend fun <Input, Output, A> Schedule<Throwable, Output>.retryOrElseEit
     }
   }
 }
+
+/**
+ * Retries [action] using any [Error] that occurred as the input to the [Schedule].
+ * It will return the last [Error] if the [Schedule] is exhausted, and ignores the output of the [Schedule].
+ */
+public suspend inline fun <Error, Result, Output> Schedule<Error, Output>.retryRaise(
+  @BuilderInference action: Raise<Error>.() -> Result,
+): Either<Error, Result> = either {
+  retry(this@retryRaise, action)
+}
+
+/**
+ * Retries [action] using any [Error] that occurred as the input to the [Schedule].
+ * It will return the last [Error] if the [Schedule] is exhausted, and ignores the output of the [Schedule].
+ */
+public suspend inline fun <Error, Result, Output> Schedule<Error, Output>.retryEither(
+  @BuilderInference action: () -> Either<Error, Result>,
+): Either<Error, Result> = retryRaise {
+  action().bind()
+}
+
+/**
+ * Retries [action] using any [Error] that occurred as the input to the [Schedule].
+ * It will return the last [Error] if the [Schedule] is exhausted, and ignores the output of the [Schedule].
+ */
+public suspend inline fun <Error, Result, Output> Raise<Error>.retry(
+  schedule: Schedule<Error, Output>,
+  @BuilderInference action: Raise<Error>.() -> Result,
+): Result {
+  var step: ScheduleStep<Error, Output> = schedule.step
+
+  while (true) {
+    currentCoroutineContext().ensureActive()
+    fold(
+      action,
+      recover = { error ->
+        when (val decision = step(error)) {
+          is Continue -> {
+            if (decision.delay != ZERO) delay(decision.delay)
+            step = decision.step
+          }
+
+          is Done -> raise(error)
+        }
+      },
+      transform = { result ->
+        return result
+      },
+    )
+  }
+}
diff --git a/...w-libs/resilience/arrow-resilience/src/commonTest/kotlin/arrow/resilience/ScheduleTest.kt b/...w-libs/resilience/arrow-resilience/src/commonTest/kotlin/arrow/resilience/ScheduleTest.kt
@@ -3,9 +3,10 @@ package arrow.resilience
 import arrow.atomic.AtomicLong
 import arrow.atomic.updateAndGet
 import arrow.core.Either
+import arrow.core.left
+import arrow.core.right
 import arrow.resilience.Schedule.Decision.Continue
 import arrow.resilience.Schedule.Decision.Done
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestResult
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.withTimeout
@@ -27,7 +28,6 @@ internal data class SideEffect(var counter: Int = 0) {
   }
 }
 
-@OptIn(ExperimentalCoroutinesApi::class)
 @ExperimentalTime
 @Suppress("UNREACHABLE_CODE", "UNUSED_VARIABLE")
 class ScheduleTest {
@@ -280,6 +280,56 @@ class ScheduleTest {
 
     result.fold({ assertEquals(ex, it) }, { fail("The impossible happened") })
   }
+
+  @Test
+  fun retryRaiseIsStackSafe(): TestResult = runTest {
+    val count = AtomicLong(0)
+    val iterations = stackSafeIteration().toLong()
+
+    suspend fun increment() {
+      count.incrementAndGet()
+    }
+
+    val result = Schedule.recurs<CustomError>(iterations).retryRaise {
+      increment()
+      raise(CustomError)
+    }
+
+    assertTrue { result is Either.Left }
+    assertEquals(iterations + 1, count.get())
+  }
+
+  @Test
+  fun retryRaiseSucceedsIfErrorIsNotRaised(): TestResult = runTest {
+    val result = Schedule.recurs<CustomError>(0).retryRaise { 1 }
+
+    assertTrue { result is Either.Right && result.value == 1 }
+  }
+
+  @Test
+  fun retryEitherIsStackSafe(): TestResult = runTest {
+    val count = AtomicLong(0)
+    val iterations = stackSafeIteration().toLong()
+
+    suspend fun increment() {
+      count.incrementAndGet()
+    }
+
+    val result = Schedule.recurs<CustomError>(iterations).retryEither {
+      increment()
+      CustomError.left()
+    }
+
+    assertTrue { result is Either.Left }
+    assertEquals(iterations + 1, count.get())
+  }
+
+  @Test
+  fun retryEitherSucceedsIfErrorIsNotRaised(): TestResult = runTest {
+    val result = Schedule.recurs<CustomError>(0).retryEither { 1.right() }
+
+    assertTrue { result is Either.Right && result.value == 1 }
+  }
 }
 
 fun <A, B> Schedule.Decision<A, B>.delay(): Duration? = when (this) {
@@ -355,3 +405,5 @@ private suspend fun <B> checkRepeat(schedule: Schedule<Long, List<B>>, expected:
 
   assertContentEquals(expected, result)
 }
+
+private object CustomError