Stack safety of combineRequests

shared/src/main/scala/fetch.scala

@@ -34,13 +34,13 @@ object `package` {
 
   private[fetch] sealed trait FetchQuery[I, A] extends FetchRequest {
     def data: Data[I, A]
-    def identities: NonEmptyList[I]
+    def identities: Set[I]
   }
   private[fetch] final case class FetchOne[I, A](id: I, data: Data[I, A]) extends FetchQuery[I, A] {
-    override def identities: NonEmptyList[I] = NonEmptyList.one(id)
+    override def identities: Set[I] = Set(id)
   }
   private[fetch] final case class Batch[I, A](ids: NonEmptyList[I], data: Data[I, A]) extends FetchQuery[I, A] {
-    override def identities: NonEmptyList[I] = ids
+    override def identities: Set[I] = ids.toList.toSet
   }
 
   // Fetch result states
@@ -59,71 +59,151 @@ object `package` {
 
   // In-progress request
 
-  private[fetch] final case class BlockedRequest[F[_]](request: FetchRequest, result: FetchStatus => F[Unit])
+  private[fetch] final case class BlockedRequest[F[_]](request: FetchRequest, result: CombinationTailRec[F])
 
   /* Combines the identities of two `FetchQuery` to the same data source. */
   private def combineIdentities[I, A](x: FetchQuery[I, A], y: FetchQuery[I, A]): NonEmptyList[I] = {
-    y.identities.foldLeft(x.identities) {
-      case (acc, i) => if (acc.exists(_ == i)) acc else NonEmptyList(acc.head, acc.tail :+ i)
-    }
+    (x.identities ++ y.identities).toList.toNel.get
+  }
+
+  private[fetch] sealed trait CombinationTailRec[F[_]] extends Product with Serializable {
+    def flatMap(f: () => CombinationTailRec[F]): CombinationTailRec[F] = CombinationFlatMap(this, f)
   }
 
+  private[fetch] case class CombinationDone[F[_]](result: F[Unit]) extends CombinationTailRec[F]
+  private[fetch] case class CombinationLeaf[F[_]](f: FetchStatus => F[Unit]) extends CombinationTailRec[F]
+  private[fetch] case class CombinationBarrier[F[_]](sub: () => CombinationTailRec[F], status: FetchStatus) extends CombinationTailRec[F]
+  private[fetch] case class CombinationSuspend[F[_]](resume: FetchStatus => CombinationTailRec[F]) extends CombinationTailRec[F]
+  private[fetch] case class CombinationFlatMap[F[_]](sub: CombinationTailRec[F], f: () => CombinationTailRec[F]) extends CombinationTailRec[F]
+
+  /* Run combineRequests in the stack safe way */
+  def runCombinationResult[F[_]: Monad](
+    comb: CombinationTailRec[F],
+    status: FetchStatus
+  ): F[Unit] =
+    comb match {
+      case CombinationDone(r) =>
+        // Combination is finished, just return the result
+        r
+
+      case CombinationLeaf(f) =>
+        // Last step before completing a combination, apply the method and box the result
+        runCombinationResult(
+          CombinationDone(f(status)),
+          status
+        )
+
+      case CombinationBarrier(sub, newStatus) =>
+        // A barrier is boxing a CombinationTailRec to which a specific status must be applied
+        // Just make a recursive call with the status
+        runCombinationResult(sub(), newStatus)
+
+      case CombinationSuspend(resume) =>
+        // Base CombinationTailRec, derives a CombinationTailRec from the current status
+        runCombinationResult(resume(status), status)
+
+      case CombinationFlatMap(current, next) =>
+        current match {
+          case CombinationDone(r) =>
+            // Current combination is done, just flatMap the result with the next combination
+            r.flatMap(_ => runCombinationResult(next(), status))
+
+          case CombinationLeaf(f) =>
+            // Current combination is a lead. Compute and box the result, and still flatMap on the next combination
+            runCombinationResult(
+              CombinationDone(f(status)).flatMap(next),
+              status
+            )
+
+          case CombinationBarrier(sub, newStatus) =>
+            // Barrier, the current combination needs to be ran with the specific status, and the next one needs to work on the current status
+            runCombinationResult(
+              sub().flatMap(() => CombinationBarrier(next, status)),
+              newStatus
+            )
+
+          case CombinationSuspend(sub) =>
+            // Derive the combination from the current status and still flatMap on the next combination
+            runCombinationResult(
+              sub(status).flatMap(next),
+              status
+            )
+
+          case CombinationFlatMap(sub, fNext) =>
+            // FlatMap of a FlatMap
+            // Treat the current combination, then the fNext, and finally the next combination of the top level flatMap
+            runCombinationResult(
+              sub.flatMap(() => fNext().flatMap(next)),
+              status
+            )
+        }
+    }
+
   /* Combines two requests to the same data source. */
   private def combineRequests[F[_] : Monad](x: BlockedRequest[F], y: BlockedRequest[F]): BlockedRequest[F] = (x.request, y.request) match {
     case (a@FetchOne(aId, ds), b@FetchOne(anotherId, _)) =>
       if (aId == anotherId)  {
         val newRequest = FetchOne(aId, ds)
-        val newResult = (r: FetchStatus) => (x.result(r), y.result(r)).tupled.void
+        val newResult = x.result.flatMap(() => y.result)
         BlockedRequest(newRequest, newResult)
       } else {
         val combined = combineIdentities(a, b)
         val newRequest = Batch(combined, ds)
-        val newResult = (r: FetchStatus) => r match {
-          case FetchDone(m : Map[Any, Any]) => {
-            val xResult = m.get(aId).map(FetchDone(_)).getOrElse(FetchMissing())
-            val yResult = m.get(anotherId).map(FetchDone(_)).getOrElse(FetchMissing())
-              (x.result(xResult), y.result(yResult)).tupled.void
-          }
+        val newResult = CombinationSuspend(
+          (r: FetchStatus) => r match {
+            case FetchDone(m : Map[Any, Any]) => {
+              val xResult = m.get(aId).map(FetchDone(_)).getOrElse(FetchMissing())
+              val yResult = m.get(anotherId).map(FetchDone(_)).getOrElse(FetchMissing())
+              CombinationBarrier(() => x.result, xResult)
+                .flatMap(() => CombinationBarrier(() => y.result, yResult))
+            }
 
-          case FetchMissing() =>
-            (x.result(r), y.result(r)).tupled.void
-        }
+            case FetchMissing() =>
+              x.result.flatMap(() => y.result)
+          }
+        )
         BlockedRequest(newRequest, newResult)
       }
 
     case (a@FetchOne(oneId, ds), b@Batch(anotherIds, _)) =>
       val combined = combineIdentities(a, b)
       val newRequest = Batch(combined, ds)
-      val newResult = (r: FetchStatus) => r match {
-        case FetchDone(m : Map[Any, Any]) => {
-          val oneResult = m.get(oneId).map(FetchDone(_)).getOrElse(FetchMissing())
+      val newResult = CombinationSuspend(
+        (r: FetchStatus) => r match {
+          case FetchDone(m: Map[Any, Any]) => {
+            val oneResult = m.get(oneId).map(FetchDone(_)).getOrElse(FetchMissing())
+            CombinationBarrier(() => x.result, oneResult)
+              .flatMap(() => y.result)
+          }
 
-          (x.result(oneResult), y.result(r)).tupled.void
+          case FetchMissing() =>
+            x.result.flatMap(() => y.result)
         }
+      )
 
-        case FetchMissing() =>
-          (x.result(r), y.result(r)).tupled.void
-      }
       BlockedRequest(newRequest, newResult)
 
     case (a@Batch(manyId, ds), b@FetchOne(oneId, _)) =>
       val combined = combineIdentities(a, b)
       val newRequest = Batch(combined, ds)
-      val newResult = (r: FetchStatus) => r match {
-        case FetchDone(m : Map[Any, Any]) => {
-          val oneResult = m.get(oneId).map(FetchDone(_)).getOrElse(FetchMissing())
-            (x.result(r), y.result(oneResult)).tupled.void
-        }
+      val newResult = CombinationSuspend(
+        (r: FetchStatus) => r match {
+          case FetchDone(m: Map[Any, Any]) => {
+            val oneResult = m.get(oneId).map(FetchDone(_)).getOrElse(FetchMissing())
+            CombinationBarrier(() => y.result, oneResult)
+              .flatMap(() => x.result)
+          }
 
-        case FetchMissing() =>
-          (x.result(r), y.result(r)).tupled.void
-      }
+          case FetchMissing() =>
+            x.result.flatMap(() => y.result)
+        }
+      )
       BlockedRequest(newRequest, newResult)
 
     case (a@Batch(manyId, ds), b@Batch(otherId, _)) =>
       val combined = combineIdentities(a, b)
       val newRequest = Batch(combined, ds)
-      val newResult = (r: FetchStatus) => (x.result(r), y.result(r)).tupled.void
+      val newResult = x.result.flatMap(() => y.result)
       BlockedRequest(newRequest, newResult)
   }
 
@@ -275,7 +355,7 @@ object `package` {
           deferred <- Deferred[F, FetchStatus]
           request = FetchOne(id, ds.data)
           result = deferred.complete _
-          blocked = BlockedRequest(request, result)
+          blocked = BlockedRequest(request, CombinationLeaf(result))
           anyDs = ds.asInstanceOf[DataSource[F, Any, Any]]
           blockedRequest = RequestMap(Map(ds.data.identity -> (anyDs, blocked)))
         } yield Blocked(blockedRequest, Unfetch[F, A](
@@ -297,7 +377,7 @@ object `package` {
           deferred <- Deferred[F, FetchStatus]
           request = FetchOne(id, ds.data)
           result = deferred.complete _
-          blocked = BlockedRequest(request, result)
+          blocked = BlockedRequest(request, CombinationLeaf(result))
           anyDs = ds.asInstanceOf[DataSource[F, Any, Any]]
           blockedRequest = RequestMap(Map(ds.data.identity -> (anyDs, blocked)))
         } yield Blocked(blockedRequest, Unfetch[F, Option[A]](
@@ -514,7 +594,7 @@ object `package` {
   private def runFetchOne[F[_]](
     q: FetchOne[Any, Any],
     ds: DataSource[F, Any, Any],
-    putResult: FetchStatus => F[Unit],
+    putResult: CombinationTailRec[F],
     cache: Ref[F, DataCache[F]],
     log: Option[Ref[F, Log]]
   )(
@@ -527,7 +607,7 @@ object `package` {
       maybeCached <- c.lookup(q.id, q.data)
       result <- maybeCached match {
         // Cached
-        case Some(v) => putResult(FetchDone(v)).as(Nil)
+        case Some(v) => runCombinationResult(putResult, FetchDone(v)).as(Nil)
 
         // Not cached, must fetch
         case None => for {
@@ -539,12 +619,12 @@ object `package` {
             case Some(a) => for {
               newC <- c.insert(q.id, a, q.data)
               _ <- cache.set(newC)
-              result <- putResult(FetchDone[Any](a))
+              result <- runCombinationResult(putResult, FetchDone[Any](a))
             } yield List(Request(q, startTime, endTime))
 
             // Missing
             case None =>
-              putResult(FetchMissing()).as(List(Request(q, startTime, endTime)))
+              runCombinationResult(putResult, FetchMissing()).as(List(Request(q, startTime, endTime)))
           }
         } yield result
       }
@@ -558,7 +638,7 @@ object `package` {
   private def runBatch[F[_]](
     q: Batch[Any, Any],
     ds: DataSource[F, Any, Any],
-    putResult: FetchStatus => F[Unit],
+    putResult: CombinationTailRec[F],
     cache: Ref[F, DataCache[F]],
     log: Option[Ref[F, Log]]
   )(
@@ -579,7 +659,7 @@ object `package` {
       cachedResults = cached.toMap
       result <- uncachedIds.toNel match {
         // All cached
-        case None => putResult(FetchDone[Map[Any, Any]](cachedResults)).as(Nil)
+        case None => runCombinationResult(putResult, FetchDone[Map[Any, Any]](cachedResults)).as(Nil)
 
         // Some uncached
         case Some(uncached) => for {
@@ -603,7 +683,7 @@ object `package` {
           updatedCache <- c.bulkInsert(batchedRequest.results.toList, q.data)
           _ <- cache.set(updatedCache)
 
-          result <- putResult(FetchDone[Map[Any, Any]](resultMap))
+          result <- runCombinationResult(putResult, FetchDone[Map[Any, Any]](resultMap))
 
         } yield batchedRequest.batches.map(Request(_, startTime, endTime))
       }

shared/src/test/scala/FetchBatchingTests.scala

@@ -65,12 +65,44 @@ class FetchBatchingTests extends FetchSpec {
     }
   }
 
+  case class BatchedDataBigId(
+    str1: String,
+    str2: String,
+    str3: String
+  )
+
+  object BigIdData extends Data[BatchedDataBigId, String] {
+    def name = "Big id batching"
+
+    implicit def source[F[_] : ConcurrentEffect]: DataSource[F, BatchedDataBigId, String] = new DataSource[F, BatchedDataBigId, String] {
+      override def data = BigIdData
+
+      override def CF = ConcurrentEffect[F]
+
+      override def fetch(request: BatchedDataBigId): F[Option[String]] = {
+        batch(NonEmptyList.one(request)).map(_.get(request))
+      }
+
+      override def batch(ids: NonEmptyList[BatchedDataBigId]): F[Map[BatchedDataBigId, String]] =
+        CF.pure(
+          ids.map { id =>
+            id -> id.toString
+          }.toList.toMap
+        )
+
+      override val batchExecution = InParallel
+    }
+  }
+
   def fetchBatchedDataSeq[F[_] : ConcurrentEffect](id: Int): Fetch[F, Int] =
     Fetch(BatchedDataSeq(id), SeqBatch.source)
 
   def fetchBatchedDataPar[F[_] : ConcurrentEffect](id: Int): Fetch[F, Int] =
     Fetch(BatchedDataPar(id), ParBatch.source)
 
+  def fetchBatchedDataBigId[F[_] : ConcurrentEffect](id: BatchedDataBigId): Fetch[F, String] =
+    Fetch(id, BigIdData.source)
+
   "A large fetch to a datasource with a maximum batch size is split and executed in sequence" in {
     def fetch[F[_] : ConcurrentEffect]: Fetch[F, List[Int]] =
       List.range(1, 6).traverse(fetchBatchedDataSeq[F])
@@ -170,20 +202,22 @@ class FetchBatchingTests extends FetchSpec {
   }
 
   "Very deep fetches don't overflow stack or heap" in {
-    val depth = 200
-    val list  = (1 to depth).toList
-    def fetch[F[_] : ConcurrentEffect]: Fetch[F, List[Int]] =
-      list.map(x => (0 until x).toList.traverse(fetchBatchedDataSeq[F]))
-        .foldLeft(
-          Fetch.pure[F, List[Int]](List.empty[Int])
-        )(_ >> _)
-
-    val io = Fetch.runLog[IO](fetch)
+    val depth = 5000
+    val ids = for {
+      id <- 0 to depth
+    } yield BatchedDataBigId(
+      str1 = "longString" + id,
+      str2 = "longString" + (id + 1),
+      str3 = "longString" + (id + 2)
+    )
+
+    val io = Fetch.runLog[IO](
+      ids.toList.traverse(fetchBatchedDataBigId[IO])
+    )
 
     io.map({
       case (log, result) => {
-        result shouldEqual (0 until depth).toList
-        log.rounds.size shouldEqual depth
+        result shouldEqual ids.map(_.toString)
       }
     }).unsafeToFuture
   }