From e0c644cb9d1d3e6e4f34814b6a87131bbf5ed9c4 Mon Sep 17 00:00:00 2001 From: Arman Bilge Date: Sun, 17 Sep 2023 22:04:06 +0000 Subject: [PATCH] Track cancelation in generic sleep --- .../unsafe/WorkStealingThreadPool.scala | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/core/jvm/src/main/scala/cats/effect/unsafe/WorkStealingThreadPool.scala b/core/jvm/src/main/scala/cats/effect/unsafe/WorkStealingThreadPool.scala index dec3dcb8f1..f008a8fa56 100644 --- a/core/jvm/src/main/scala/cats/effect/unsafe/WorkStealingThreadPool.scala +++ b/core/jvm/src/main/scala/cats/effect/unsafe/WorkStealingThreadPool.scala @@ -627,13 +627,18 @@ private[effect] final class WorkStealingThreadPool( cancel } - override def sleep(delay: FiniteDuration, task: Runnable): Runnable = - sleepInternal( - delay, - new AtomicBoolean with (Right[Nothing, Unit] => Unit) { // run at most once - def apply(ru: Right[Nothing, Unit]) = if (compareAndSet(false, true)) task.run() - } - ) + override def sleep(delay: FiniteDuration, task: Runnable): Runnable = { + val cb = new AtomicBoolean with (Right[Nothing, Unit] => Unit) { // run at most once + def apply(ru: Right[Nothing, Unit]) = if (compareAndSet(false, true)) task.run() + } + + val cancel = sleepInternal(delay, cb) + + () => { + cb.set(true) + cancel.run() + } + } /** * Shut down the thread pool and clean up the pool state. Calling this method after the pool