import java.util.{Timer, TimerTask}
import org.joda.time.Duration
import scala.concurrent.{Await, ExecutionContext, Future, Promise}
import scala.concurrent.duration.Duration.Inf
import scala.math.min
import scala.util.Random

def withBackoff[A](maxRetry: Int,
                   baseDelay: Duration,
                   maxBackoffTime: Duration)
                  (f: => Future[A])
                  (implicit executionContext: ExecutionContext,
                            random: Random,
                            timer: Timer): Future[A] = {
  val maxMillis = maxBackoffTime.getMillis.toInt
  val promise = Promise[A]

  def task(retry: Int, delay: Int)
          (implicit executionContext: ExecutionContext): TimerTask =
    new TimerTask {
      def run() {
        val future = f

        future.onSuccess { case result => promise.success(result) }

        future.onFailure {
          case _ if retry > 0 =>
            val jitter = random.nextInt(min(delay, maxMillis)) + 1

            println(s"Retrying in ${jitter}ms")
            timer.schedule(task(retry - 1, delay << 1), jitter)

          case exception =>
            promise.failure(exception)
        }
      }
    }

  task(maxRetry, baseDelay.getMillis.toInt).run()

  promise.future
}


def flaky(probability: Int)(implicit random: Random) = {
  val number = random.nextInt(probability)

  if (number == 0)
    Future.successful(number)
  else
    Future.failed(new RuntimeException(s"It was ${number}"))
}


import scala.concurrent.ExecutionContext.Implicits.global

implicit val random = Random
implicit val timer: Timer = new Timer(true)

Await.result(withBackoff(10, new Duration(500), new Duration(20000))(flaky(5)), Inf)