diff --git a/core/shared/src/main/scala/cats/effect/IO.scala b/core/shared/src/main/scala/cats/effect/IO.scala index 112a125882..11b395b6a3 100644 --- a/core/shared/src/main/scala/cats/effect/IO.scala +++ b/core/shared/src/main/scala/cats/effect/IO.scala @@ -1599,6 +1599,13 @@ object IO extends IOCompanionPlatform with IOLowPriorityImplicits with TuplePara def parTraverseN_[T[_]: Foldable, A, B](n: Int)(ta: T[A])(f: A => IO[B]): IO[Unit] = _asyncForIO.parTraverseN_(n)(ta)(f) + /** + * Like `Parallel.parFlatTraverse`, but limits the degree of parallelism. + */ + def parFlatTraverseN[T[_]: Traverse: cats.FlatMap, A, B](n: Int)(ta: T[A])( + f: A => IO[T[B]]): IO[T[B]] = + _asyncForIO.parFlatTraverseN(n)(ta)(f) + /** * Like `Parallel.parSequence` */ @@ -1623,6 +1630,12 @@ object IO extends IOCompanionPlatform with IOLowPriorityImplicits with TuplePara def parSequenceN_[T[_]: Foldable, A](n: Int)(tma: T[IO[A]]): IO[Unit] = _asyncForIO.parSequenceN_(n)(tma) + /** + * Like `Parallel.parFlatSequence`, but limits the degree of parallelism. + */ + def parFlatSequenceN[T[_]: Traverse: cats.FlatMap, A](n: Int)(tmta: T[IO[T[A]]]): IO[T[A]] = + _asyncForIO.parFlatSequenceN(n)(tmta) + /** * Like `Parallel.parReplicateA`, but limits the degree of parallelism. */ diff --git a/kernel/shared/src/main/scala/cats/effect/kernel/GenConcurrent.scala b/kernel/shared/src/main/scala/cats/effect/kernel/GenConcurrent.scala index df6c4222d6..da74552d65 100644 --- a/kernel/shared/src/main/scala/cats/effect/kernel/GenConcurrent.scala +++ b/kernel/shared/src/main/scala/cats/effect/kernel/GenConcurrent.scala @@ -16,7 +16,7 @@ package cats.effect.kernel -import cats.{Foldable, Monoid, Semigroup, Traverse} +import cats.{FlatMap, Foldable, Monoid, Semigroup, Traverse} import cats.data.{EitherT, IorT, Kleisli, OptionT, WriterT} import cats.effect.kernel.instances.spawn._ import cats.effect.kernel.syntax.all._ @@ -143,6 +143,18 @@ trait GenConcurrent[F[_], E] extends GenSpawn[F, E] { * limit. */ def parTraverseN[T[_]: Traverse, A, B](n: Int)(ta: T[A])(f: A => F[B]): F[T[B]] = { + implicit val F: GenConcurrent[F, E] = this + parTraverseNImpl[T, A, B, B](n)(ta)(f)(_.sequence[F, B]) + } + + /** + * Shared core implementation for both [[parTraverseN]] and [[parFlatTraverseN]], taking a + * function that decides on how to sequence the result. + * @tparam B + * the intermediate result type of the function `f`, needs to be sequenceable to `C` + */ + private def parTraverseNImpl[T[_]: Traverse, A, B, C](n: Int)(ta: T[A])(f: A => F[B])( + seq: T[F[B]] => F[T[C]]): F[T[C]] = { require(n >= 1, s"Concurrency limit should be at least 1, was: $n") implicit val F: GenConcurrent[F, E] = this @@ -242,7 +254,7 @@ trait GenConcurrent[F[_], E] extends GenSpawn[F, E] { } } - results.flatMap(_.sequence).onCancel(cancelAllAndJoin) + results.flatMap(seq).onCancel(cancelAllAndJoin) } } } @@ -336,6 +348,23 @@ trait GenConcurrent[F[_], E] extends GenSpawn[F, E] { } } + /** + * Like `Parallel.parFlatSequence`, but limits the degree of parallelism. See [[parSequenceN]] + * for fairness considerations, the same semantics apply here. + */ + def parFlatSequenceN[T[_]: Traverse: FlatMap, A](n: Int)(tma: T[F[T[A]]]): F[T[A]] = + parFlatTraverseN(n)(tma)(identity) + + /** + * Like `Parallel.parFlatTraverse`, but limits the degree of parallelism. See [[parTraverseN]] + * for fairness considerations, the same semantics apply here. + */ + def parFlatTraverseN[T[_]: Traverse: FlatMap, A, B](n: Int)(ta: T[A])( + f: A => F[T[B]]): F[T[B]] = { + implicit val F: GenConcurrent[F, E] = this + parTraverseNImpl[T, A, T[B], B](n)(ta)(f)(_.flatSequence) + } + override def racePair[A, B](fa: F[A], fb: F[B]) : F[Either[(Outcome[F, E, A], Fiber[F, E, B]), (Fiber[F, E, A], Outcome[F, E, B])]] = { implicit val F: GenConcurrent[F, E] = this diff --git a/kernel/shared/src/main/scala/cats/effect/kernel/syntax/GenConcurrentSyntax.scala b/kernel/shared/src/main/scala/cats/effect/kernel/syntax/GenConcurrentSyntax.scala index 35b7604aac..164d83822b 100644 --- a/kernel/shared/src/main/scala/cats/effect/kernel/syntax/GenConcurrentSyntax.scala +++ b/kernel/shared/src/main/scala/cats/effect/kernel/syntax/GenConcurrentSyntax.scala @@ -16,7 +16,7 @@ package cats.effect.kernel.syntax -import cats.{Foldable, Traverse} +import cats.{FlatMap, Foldable, Traverse} import cats.effect.kernel.GenConcurrent trait GenConcurrentSyntax { @@ -34,6 +34,11 @@ trait GenConcurrentSyntax { ): ConcurrentParSequenceNOps[T, F, A] = new ConcurrentParSequenceNOps(wrapped) + implicit def concurrentParFlatSequenceOps[T[_], F[_], A]( + wrapped: T[F[T[A]]] + ): ConcurrentParFlatSequenceNOps[T, F, A] = + new ConcurrentParFlatSequenceNOps(wrapped) + } final class GenConcurrentOps_[F[_], A] private[syntax] (private val wrapped: F[A]) @@ -57,6 +62,11 @@ final class ConcurrentParTraverseNOps[T[_], A] private[syntax] ( f: A => F[B] )(implicit T: Foldable[T], F: GenConcurrent[F, ?]): F[Unit] = F.parTraverseN_(n)(wrapped)(f) + + def parFlatTraverseN[F[_], B](n: Int)( + f: A => F[T[B]] + )(implicit T: Traverse[T], FM: FlatMap[T], F: GenConcurrent[F, ?]): F[T[B]] = + F.parFlatTraverseN(n)(wrapped)(f) } final class ConcurrentParSequenceNOps[T[_], F[_], A] private[syntax] ( @@ -68,3 +78,11 @@ final class ConcurrentParSequenceNOps[T[_], F[_], A] private[syntax] ( def parSequenceN_(n: Int)(implicit T: Foldable[T], F: GenConcurrent[F, ?]): F[Unit] = F.parSequenceN_(n)(wrapped) } + +final class ConcurrentParFlatSequenceNOps[T[_], F[_], A] private[syntax] ( + private val wrapped: T[F[T[A]]] +) extends AnyVal { + def parFlatSequenceN( + n: Int)(implicit T: Traverse[T], FM: FlatMap[T], F: GenConcurrent[F, ?]): F[T[A]] = + F.parFlatSequenceN(n)(wrapped) +} diff --git a/kernel/shared/src/test/scala/cats/effect/kernel/SyntaxSuite.scala b/kernel/shared/src/test/scala/cats/effect/kernel/SyntaxSuite.scala index d71b8463a4..993d1344b7 100644 --- a/kernel/shared/src/test/scala/cats/effect/kernel/SyntaxSuite.scala +++ b/kernel/shared/src/test/scala/cats/effect/kernel/SyntaxSuite.scala @@ -61,6 +61,16 @@ class SyntaxSuite extends FunSuite { result: F[Unit] } + { + val result = List(target).parFlatTraverseN(3)(t => t.map(List(_))) + result: F[List[A]] + } + + { + val result = List(target.map(List(_))).parFlatSequenceN(3) + result: F[List[A]] + } + { val result = target.parReplicateAN(3)(5) result: F[List[A]] diff --git a/tests/shared/src/test/scala/cats/effect/IOSuite.scala b/tests/shared/src/test/scala/cats/effect/IOSuite.scala index 4a4ecd5dd2..d9a2b276f4 100644 --- a/tests/shared/src/test/scala/cats/effect/IOSuite.scala +++ b/tests/shared/src/test/scala/cats/effect/IOSuite.scala @@ -2268,6 +2268,72 @@ class IOSuite extends BaseScalaCheckSuite with DisciplineSuite with IOPlatformSu assertCompleteAs(test.attempt.void, ()) } + ticked("parFlatTraverseN - run tasks in parallel and flatten the results (ticked)") { + implicit ticker => + val p = List(1, 2, 3) + .parFlatTraverseN(2) { (n: Int) => IO.sleep(2.seconds) >> IO.pure(List.fill(n)(())) } + .timeoutTo(5.seconds, IO(fail("parFlatSequenceN took too long"))) + + assertCompleteAs(p, List.fill(1 + 2 + 3)(())) + } + + real("parFlatTraverseN - throw when n < 1") { + IO.defer { + List.empty[Int].parFlatTraverseN(0)(List(_).pure[IO]) + }.mustFailWith[IllegalArgumentException] + } + + real("parFlatTraverseN - propagate errors") { + List(1, 2, 3) + .parFlatTraverseN(2) { (n: Int) => + if (n == 2) IO.raiseError(new RuntimeException) else List(n).pure[IO] + } + .mustFailWith[RuntimeException] + } + + ticked("parFlatTraverseN - be cancelable") { implicit ticker => + val p = for { + f <- List(1, 2, 3).parFlatTraverseN(2)(_ => IO.never[List[Int]]).start + _ <- IO.sleep(100.millis) + _ <- f.cancel + } yield true + + assertCompleteAs(p, true) + } + + ticked("parFlatSequenceN - run tasks in parallel and flatten the results (ticked)") { + implicit ticker => + val p = List(1, 2, 3) + .map { (n: Int) => IO.sleep(2.seconds) >> IO.pure(List.fill(n)(())) } + .parFlatSequenceN(2) + .timeoutTo(5.seconds, IO(fail("parFlatSequenceN took too long"))) + + assertCompleteAs(p, List.fill(1 + 2 + 3)(())) + } + + real("parFlatSequenceN - throw when n < 1") { + IO.defer { + List.empty[IO[List[Int]]].parFlatSequenceN(0) + }.mustFailWith[IllegalArgumentException] + } + + real("parFlatSequenceN - propagate errors") { + List(1, 2, 3) + .map { (n: Int) => if (n == 2) IO.raiseError(new RuntimeException) else List(n).pure[IO] } + .parFlatSequenceN(2) + .mustFailWith[RuntimeException] + } + + ticked("parFlatSequenceN - be cancelable") { implicit ticker => + val p = for { + f <- List(1, 2, 3).map(_ => IO.never[List[IO[Int]]]).parFlatSequenceN(2).start + _ <- IO.sleep(100.millis) + _ <- f.cancel + } yield true + + assertCompleteAs(p, true) + } + real("parallel - run parallel actually in parallel") { val x = IO.sleep(2.seconds) >> IO.pure(1) val y = IO.sleep(2.seconds) >> IO.pure(2)