diff --git a/core/shared/src/main/scala/cats/effect/IO.scala b/core/shared/src/main/scala/cats/effect/IO.scala index 7138a0b915f..0c51d9da36f 100644 --- a/core/shared/src/main/scala/cats/effect/IO.scala +++ b/core/shared/src/main/scala/cats/effect/IO.scala @@ -819,9 +819,12 @@ sealed abstract class IO[+A] private () extends IOPlatform[A] { */ def timeoutTo[A2 >: A](duration: Duration, fallback: IO[A2]): IO[A2] = { handleDuration[IO[A2]](duration, this) { finiteDuration => - race(IO.sleep(finiteDuration)).flatMap { - case Right(_) => fallback - case Left(value) => IO.pure(value) + IO.uncancelable { poll => + poll(racePair(IO.sleep(finiteDuration))) flatMap { + case Left((oc, f)) => f.cancel *> oc.embed(poll(IO.canceled) *> IO.never) + case Right((f, _)) => + f.cancel *> f.join.flatMap { oc => oc.fold(fallback, IO.raiseError, identity) } + } } } } diff --git a/kernel/shared/src/main/scala/cats/effect/kernel/GenTemporal.scala b/kernel/shared/src/main/scala/cats/effect/kernel/GenTemporal.scala index 56d6d0d2702..5ed9f8d4518 100644 --- a/kernel/shared/src/main/scala/cats/effect/kernel/GenTemporal.scala +++ b/kernel/shared/src/main/scala/cats/effect/kernel/GenTemporal.scala @@ -91,9 +91,13 @@ trait GenTemporal[F[_], E] extends GenConcurrent[F, E] with Clock[F] { handleDuration(duration, fa)(timeoutTo(fa, _, fallback)) protected def timeoutTo[A](fa: F[A], duration: FiniteDuration, fallback: F[A]): F[A] = - flatMap(race(fa, sleep(duration))) { - case Left(a) => pure(a) - case Right(_) => fallback + uncancelable { poll => + implicit val F: GenTemporal[F, E] = this + + poll(racePair(fa, sleep(duration))) flatMap { + case Left((oc, f)) => f.cancel *> oc.embed(poll(F.canceled) *> F.never) + case Right((f, _)) => f.cancel *> f.join.flatMap { oc => oc.embed(fallback) } + } } /** @@ -115,9 +119,16 @@ trait GenTemporal[F[_], E] extends GenConcurrent[F, E] with Clock[F] { protected def timeout[A](fa: F[A], duration: FiniteDuration)( implicit ev: TimeoutException <:< E): F[A] = { - flatMap(race(fa, sleep(duration))) { - case Left(a) => pure(a) - case Right(_) => raiseError[A](ev(new TimeoutException(duration.toString()))) + uncancelable { poll => + implicit val F: GenTemporal[F, E] = this + + poll(racePair(fa, sleep(duration))) flatMap { + case Left((oc, f)) => f.cancel *> oc.embed(poll(F.canceled) *> F.never) + case Right((f, _)) => + f.cancel *> f.join.flatMap { oc => + oc.embed(raiseError[A](ev(new TimeoutException(duration.toString())))) + } + } } }