TimeUtils.timeCase/FOps.timeCase added

cats-effect/src/main/scala/com/avast/sst/catseffect/TimeUtils.scala

@@ -1,7 +1,7 @@
 package com.avast.sst.catseffect
 
 import cats.effect.syntax.bracket._
-import cats.effect.{Bracket, Clock}
+import cats.effect.{Bracket, Clock, ExitCase}
 import cats.syntax.flatMap._
 import cats.syntax.functor._
 
@@ -10,9 +10,10 @@ import scala.concurrent.duration.Duration
 
 object TimeUtils {
 
+  private final val unit = TimeUnit.NANOSECONDS
+
   /** Measures the time it takes the effect to finish and records it using the provided function. */
   def time[F[_], A](f: F[A])(record: Duration => F[Unit])(implicit F: Bracket[F, Throwable], C: Clock[F]): F[A] = {
-    val unit = TimeUnit.NANOSECONDS
     for {
       start <- C.monotonic(unit)
       result <- f.guarantee {
@@ -21,6 +22,22 @@ object TimeUtils {
     } yield result
   }
 
+  /** Measures the time it takes the effect to finish and records it using the provided function. It distinguishes between successful
+    * and failure state.
+    * Please note, that in case of the effect cancellation the `record` is not invoked at all.
+    */
+  def timeCase[F[_], A](f: F[A])(record: Either[Duration, Duration] => F[Unit])(implicit F: Bracket[F, Throwable], C: Clock[F]): F[A] = {
+    def calculateAndRecordAs(start: Long)(wrap: Duration => Either[Duration, Duration]): F[Unit] = {
+      C.monotonic(unit).map(computeTime(start)).flatMap(d => record(wrap(d)))
+    }
+
+    F.bracketCase(C.monotonic(unit))(_ => f) {
+      case (start, ExitCase.Completed) => calculateAndRecordAs(start)(Right(_))
+      case (start, ExitCase.Error(_))  => calculateAndRecordAs(start)(Left(_))
+      case _                           => F.unit
+    }
+  }
+
   private def computeTime(start: Long)(end: Long) = Duration.fromNanos(end - start)
 
 }

cats-effect/src/main/scala/com/avast/sst/catseffect/syntax/TimeSyntax.scala

@@ -20,6 +20,14 @@ object TimeSyntax {
     /** Measures the time it takes the effect to finish and records it using the provided function. */
     def time(record: Duration => F[Unit])(implicit F: Bracket[F, Throwable], C: Clock[F]): F[A] = TimeUtils.time(f)(record)
 
+    /** Measures the time it takes the effect to finish and records it using the provided function. It distinguishes between successful
+      * and failure state.
+      * Please note, that in case of the effect cancellation the `record` is not invoked at all.
+      */
+    def timeCase(record: Either[Duration, Duration] => F[Unit])(implicit F: Bracket[F, Throwable], C: Clock[F]): F[A] = {
+      TimeUtils.timeCase(f)(record)
+    }
+
   }
 
 }

cats-effect/src/test/scala/com/avast/sst/catseffect/syntax/FOpsTest.scala

@@ -32,4 +32,54 @@ class FOpsTest extends AsyncFunSuite {
     io.unsafeToFuture()
   }
 
+  test("timeCase success") {
+    val sleepTime = Duration.fromNanos(500000000)
+    implicit val mockClock: Clock[IO] = new Clock[IO] {
+      var values = List(0L, sleepTime.toNanos)
+      override def monotonic(unit: TimeUnit): IO[Long] = {
+        val time = values.head
+        values = values.tail
+        IO.pure(time)
+      }
+      override def realTime(unit: TimeUnit): IO[Long] = ???
+    }
+
+    val io = for {
+      ref <- Ref.of[IO, Option[Either[Duration, Duration]]](None)
+      _ <- IO.sleep(sleepTime).timeCase { eitherD =>
+        ref.set(Some(eitherD))
+      }
+      result <- ref.get
+    } yield assert(result === Some(Right(sleepTime)))
+
+    io.unsafeToFuture()
+  }
+
+  test("timeCase failure") {
+    val sleepTime = Duration.fromNanos(500000000)
+    implicit val mockClock: Clock[IO] = new Clock[IO] {
+      var values = List(0L, sleepTime.toNanos)
+      override def monotonic(unit: TimeUnit): IO[Long] = {
+        val time = values.head
+        values = values.tail
+        IO.pure(time)
+      }
+      override def realTime(unit: TimeUnit): IO[Long] = ???
+    }
+
+    val io = for {
+      ref <- Ref.of[IO, Option[Either[Duration, Duration]]](None)
+      _ <- IO.sleep(sleepTime)
+      _ <- IO
+        .raiseError(new RuntimeException("my exception"))
+        .timeCase { eitherD =>
+          ref.set(Some(eitherD))
+        }
+        .attempt
+      result <- ref.get
+    } yield assert(result === Some(Left(sleepTime)))
+
+    io.unsafeToFuture()
+  }
+
 }