zio · jdegoes · Jul 24, 2020 · Jul 24, 2020 · Jul 24, 2020 · Jul 24, 2020
diff --git a/core-tests/shared/src/test/scala/zio/ScheduleSpec.scala b/core-tests/shared/src/test/scala/zio/ScheduleSpec.scala
@@ -350,6 +350,31 @@ object ScheduleSpec extends ZIOBaseSpec {
         _   <- ref.getAndUpdate(_ + 1).repeat(schedule(ref))
         res <- ref.get
       } yield assert(res)(equalTo(8))
+    },
+    testM("Reset after some inactivity") {
+
+      def io(ref: Ref[Int], latch: Promise[Nothing, Unit]): ZIO[Clock, String, Unit] =
+        ref
+          .updateAndGet(_ + 1)
+          .flatMap(retries =>
+            // the 5th retry will fail after 10 seconds to let the schedule reset
+            if (retries == 5) latch.succeed(()) *> io(ref, latch).delay(10.seconds)
+            // the 10th retry will succeed, which is only possible if the schedule was reset
+            else if (retries == 10) UIO.unit
+            else ZIO.fail("Boom")
+          )
+
+      assertM {
+        for {
+          retriesCounter <- Ref.make(-1)
+          latch          <- Promise.make[Nothing, Unit]
+          fiber          <- io(retriesCounter, latch).retry(Schedule.recurs(5).resetAfter(5.seconds)).fork
+          _              <- latch.await
+          _              <- TestClock.adjust(10.seconds)
+          _              <- fiber.join
+          retries        <- retriesCounter.get
+        } yield retries
+      }(equalTo(10))
     }
   )
 

diff --git a/core/shared/src/main/scala/zio/Schedule.scala b/core/shared/src/main/scala/zio/Schedule.scala
@@ -465,7 +465,7 @@ trait Schedule[-R, -A, +B] extends Serializable { self =>
       def sleep(duration: Duration)   = ZIO.unit
     }
 
-    provideSome[R1](env => ev1.update[R1, Clock.Service](env, proxy(_)))
+    provideSome[R1](env => ev1.update[R1, Clock.Service](env, proxy))
   }
 
   /**
@@ -522,6 +522,25 @@ trait Schedule[-R, -A, +B] extends Serializable { self =>
   final def repetitions: Schedule[R, A, Int] =
     fold(0)((n: Int, _: B) => n + 1)
 
+  /**
+   * Return a new schedule that automatically resets the schedule to its initial state
+   * after some time of inactivity defined by `duration`.
+   */
+  final def resetAfter(duration: Duration): Schedule[R with Clock, A, B] =
+    new Schedule[R with Clock, A, B] {
+      type State = (self.State, Long)
+      val initial: URIO[R with Clock, State] = self.initial zip clock.nanoTime
+      val extract: (A, State) => B           = { case (a, (state, _)) => self.extract(a, state) }
+      val update: (A, State) => ZIO[R with Clock, Unit, State] = {
+        case (a, (state, lastUpdate)) =>
+          for {
+            cdt        <- clock.nanoTime
+            stateToUse <- if (lastUpdate + duration.toNanos < cdt) initial.map(_._1) else UIO(state)
+            result     <- self.update(a, stateToUse) zip clock.nanoTime
+          } yield result
+      }
+    }
+
   /**
    * Puts this schedule into the second element of a either, and passes along
    * another value unchanged as the first element of the either.
@@ -563,7 +582,7 @@ trait Schedule[-R, -A, +B] extends Serializable { self =>
    * Sends every output value to the specified sink.
    */
   final def tapOutput[R1 <: R](f: B => ZIO[R1, Nothing, Unit]): Schedule[R1, A, B] =
-    updated(update => (a, s) => update(a, s).flatMap(s1 => f(self.extract(a, s1)).as(s1)))
+    updated(update => (a, s) => update(a, s).tap(s1 => f(self.extract(a, s1))))
 
   /**
    * Returns a new schedule with the update function transformed by the