Merge pull request #332 from nplevin/0.4-inf-loop-fix

Stop exceptions from causing an infinite loop.
davenverse · Dec 13, 2021 · b5bd22d · b5bd22d
2 parents c568174 + 83fa483
commit b5bd22d
Showing 1 changed file with 19 additions and 6 deletions.
diff --git a/modules/core/src/main/scala/io/chrisdavenport/mules/DispatchOneCache.scala b/modules/core/src/main/scala/io/chrisdavenport/mules/DispatchOneCache.scala
@@ -62,15 +62,18 @@ final class DispatchOneCache[F[_], K, V] private[DispatchOneCache] (
         emptyFV
     }
 
-  private def insertAtomic(k: K, action: K => F[V]): F[Unit] = {
+  private def insertAtomic(k: K, action: K => F[V]): F[Option[Either[Throwable, V]]] = {
     mapRef(k).modify{
       case None =>
         (None, createEmptyIfUnset(k))
       case s@Some(cacheItem) =>
         (s, updateIfFailedThenCreate(k, cacheItem))
     }.flatMap{ maybeDeferred =>
-        maybeDeferred.bracketCase(_.traverse_{ deferred =>
-          action(k).attempt.flatMap(e => deferred.complete(e).attempt.void)
+        maybeDeferred.bracketCase(_.flatTraverse{ deferred =>
+          action(k).attempt.flatMap(e => deferred.complete(e).attempt map {
+            case Left(_) => Option.empty // Either.left[Throwable, V](err) //only happened if complete action fails
+            case Right(_) => Option(e)
+          })
         }){
           case (Some(deferred), ExitCase.Canceled) => deferred.complete(CancelationDuringDispatchOneCacheInsertProcessing.asLeft).attempt.void
           case (Some(deferred), ExitCase.Error(e)) => deferred.complete(e.asLeft).attempt.void
@@ -99,10 +102,20 @@ final class DispatchOneCache[F[_], K, V] private[DispatchOneCache] (
       }
       .flatMap{
         case Some(s) => s.item.get.flatMap{
-          case Left(_) => insertAtomic(k, action) >> lookupOrLoad(k, action)
-          case Right(v) => F.pure(v)
+          case Left(_) => insertAtomic(k, action)
+          case Right(v) => F.pure(Option(Either.right[Throwable,V](v)))
         }
-        case None => insertAtomic(k, action) >> lookupOrLoad(k, action)
+        case None => insertAtomic(k, action)
+      }
+      .flatMap {
+        case Some(res) => res match {
+          case Right(v) => v.pure[F]
+          case Left(err) => err match {
+            case CancelationDuringDispatchOneCacheInsertProcessing => lookupOrLoad(k,action)
+            case _ => F.raiseError(err)
+          }
+        }
+        case _ => lookupOrLoad(k,action) //cache miss case?
       }
   }